feat: add --embed-weights host for a pinned host embedding table

docs/cli.md

@@ -71,6 +71,7 @@ ft serve --model ... --gpu GPU-9e8d7c6b  # the same card by UUID (a unique prefi
 |---|---|---|
 | `--memory-ratio` | 0.9 | Fraction of free VRAM the engine may use (weights + MoE cache + KV) |
 | `--num-pages` / `--num-tokens` | auto | KV capacity override in pages / tokens (mutually exclusive; auto sizes from VRAM left after weights and MoE cache) |
+| `--embed-weights` | gpu | Where the input token embedding table lives. `host` keeps it in pinned host RAM and each lookup reads its rows over PCIe, so its VRAM goes to the expert cache and KV. Refused when the table is tied to the LM head |
 | `--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` |
 | `--attention-backend`, `--attn` | auto | `trtllm`/`fi`/`fa`/`triton`/`dsv4_sparse`/`dsa`; `prefill,decode` pair allowed; auto picks per model + GPU |

python/freetoken/engine/config.py

@@ -34,6 +34,8 @@ 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"
+    # --embed-weights: "gpu" (default) keeps the input embedding tables in VRAM; "host" keeps the untied ones in pinned host RAM, read zero-copy by the lookup.
+    embed_weights: str = "gpu"
     # 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"

python/freetoken/engine/engine.py

@@ -15,7 +15,7 @@
 from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info
 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 import OffloadMoELayer, VocabParallelEmbedding, iter_moe_layers, iter_ops, set_rope_device
 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
@@ -285,10 +285,44 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att
             )
 
 
+def _host_embedding_keys(model, config: EngineConfig) -> frozenset[str]:
+    """State-dict keys of the tables --embed-weights host keeps in pinned host RAM."""
+    if config.embed_weights != "host":
+        return frozenset()
+    if config.model_config.tie_word_embeddings:
+        raise ValueError(
+            "--embed-weights host: this model's embedding table is tied to its LM head, which reads the whole table "
+            "every step; serve it with --embed-weights gpu"
+        )
+    tables = {id(op.weight) for op in iter_ops(model) if type(op) is VocabParallelEmbedding}
+    keys = frozenset(key for key, t in model.state_dict().items() if id(t) in tables)
+    if not keys:
+        raise ValueError(
+            f"--embed-weights host: {type(model).__name__} has no plain embedding table to move (a DF11 table stays "
+            "on the GPU); serve it with --embed-weights gpu"
+        )
+    return keys
+
+
+def _pinned_copy(key: str, t: torch.Tensor) -> torch.Tensor:
+    from freetoken.kernel.pinned import alloc_pinned_tensor
+
+    try:
+        host = alloc_pinned_tensor(*t.shape, dtype=t.dtype)
+    except RuntimeError as exc:
+        raise PinFailed(
+            f"cudaHostAlloc refused {t.nbytes / 2**20:.0f} MiB for {key} (--embed-weights host); "
+            "free host RAM or serve with --embed-weights gpu"
+        ) from exc
+    host.copy_(t)
+    return host
+
+
 def _make_dummy_weight_state_dict(
     model_state: Dict[str, torch.Tensor],
     *,
     device: torch.device,
+    host_keys: frozenset[str] = frozenset(),
 ) -> Dict[str, torch.Tensor]:
     state_dict: Dict[str, torch.Tensor] = {}
     fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2)
@@ -311,6 +345,8 @@ def _make_dummy_weight_state_dict(
             state_dict[key] = torch.full(param.shape, 127, dtype=param.dtype, device=device)
         else:
             state_dict[key] = torch.zeros(param.shape, dtype=param.dtype, device=device)
+        if key in host_keys:
+            state_dict[key] = _pinned_copy(key, state_dict[key])
     return state_dict
 
 
@@ -344,6 +380,7 @@ def _materialize_loaded_weight_state_dict(
     weights: Iterable[Tuple[str, torch.Tensor]],
     *,
     device: torch.device,
+    host_keys: frozenset[str] = frozenset(),
 ) -> Dict[str, torch.Tensor]:
     state_dict: Dict[str, torch.Tensor] = {}
     for key, weight in weights:
@@ -353,6 +390,8 @@ def _materialize_loaded_weight_state_dict(
             if key.endswith(".input_scale"):
                 continue
             state_dict[key] = weight.to(device=device)
+        elif key in host_keys:
+            state_dict[key] = _pinned_copy(key, weight.to(dtype=expected.dtype))
         else:
             state_dict[key] = weight.to(device=device, dtype=expected.dtype)
     return state_dict
@@ -418,13 +457,16 @@ def __init__(self, config: EngineConfig):
         self._post_weights_free = post_weights_free
         self.moe_offload_cache = None
         self.cpu_moe_executor = None
-        # Host-side auxiliary stores (qwen4_exp's pinned PLE table): after the weights so a
+        # Host-side auxiliary stores (--embed-weights host tables, 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
+        host_embeds = [op for op in iter_ops(self.model) if type(op) is VocabParallelEmbedding and not op.weight.is_cuda]
+        self._host_tables_bytes = sum(op.weight.nbytes for op in host_embeds)
+        if host_embeds:
+            logger.info_rank0(f"Embedding tables in pinned host RAM (--embed-weights host): {mem_GB(self._host_tables_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"):
@@ -574,16 +616,17 @@ def _load_weights(self, config: EngineConfig) -> None:
         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`")
+        host_keys = _host_embedding_keys(self.model, config)
         with _weight_load_context():
-            self.model.load_state_dict(self._load_weight_state_dict(config))
+            self.model.load_state_dict(self._load_weight_state_dict(config, host_keys))
         finalize_quant(self.model)
         if resident:
             self._load_resident_experts(config, resident)
 
-    def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tensor]:
+    def _load_weight_state_dict(self, config: EngineConfig, host_keys: frozenset[str]) -> Dict[str, torch.Tensor]:
         model_state = self.model.state_dict()
         if config.use_dummy_weight:
-            return _make_dummy_weight_state_dict(model_state, device=self.device)
+            return _make_dummy_weight_state_dict(model_state, device=self.device, host_keys=host_keys)
         # _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;
         # MoE models exclude routed experts (they load as expert banks, not dense weights).
@@ -596,6 +639,7 @@ def _load_weight_state_dict(self, config: EngineConfig) -> Dict[str, torch.Tenso
                 include_vision=bool(config.active_encoders),
             ),
             device=self.device,
+            host_keys=host_keys,
         )
 
     def _load_resident_experts(self, config: EngineConfig, layers) -> None:

python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu

@@ -10,6 +10,8 @@
 #include "mmq.cuh"
 #include "mmvq.cuh"
 #include "quantize.cuh"
+// after the ggml headers, so their HIP macros win over hip_compat.h
+#include <freetoken/pinned.h>
 
 #include <tvm/ffi/container/tensor.h>
 #include <tvm/ffi/extra/c_env_api.h>
@@ -436,7 +438,8 @@ void op_dequantize(TensorView w, int64_t type_id, int64_t k, TensorView out) {
 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());
+    // w may be a pinned host table read zero-copy, so the launch device comes from out
+    const DeviceGuard guard(out.device());
     const ggml_type type = (ggml_type) type_id;
     check_weight(w, type, k);
     const int64_t n = ids.size(0);
@@ -446,8 +449,8 @@ void op_get_rows(TensorView w, int64_t type_id, int64_t k, TensorView ids, Tenso
         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()));
+    get_rows_cuda(host::device_alias(w.data_ptr(), w.device()), 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(out.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); }

python/freetoken/kernel/csrc/include/freetoken/pinned.h

@@ -0,0 +1,40 @@
+#pragma once
+
+#include <freetoken/hip_compat.h>
+#include <freetoken/utils.h>
+
+#include <dlpack/dlpack.h>
+
+namespace host {
+
+// Pinned host memory is GPU-dereferenceable at its host VA only where UVA identity
+// holds (Linux). On Windows/WDDM, cudaHostRegister'd memory maps to a different device
+// address, so host-resident tensors are translated where their pointer enters kernel params.
+// Cached once per process: FreeToken pins one CUDA device per process (set at engine launch).
+inline bool host_ptr_identity() {
+    static const bool identity = [] {
+        int device = 0;
+        if (cudaGetDevice(&device) != cudaSuccess) {
+            return false;  // fail closed: translate (and surface errors), don't assume identity
+        }
+        int uva = 0, reg = 0;
+        cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device);
+        cudaDeviceGetAttribute(&reg, cudaDevAttrCanUseHostPointerForRegisteredMem, device);
+        return uva == 1 && reg == 1;
+    }();
+    return identity;
+}
+
+inline void* device_alias(void* ptr, DLDevice dev) {
+    if (dev.device_type == kDLCUDA || dev.device_type == kDLROCM || host_ptr_identity()) {
+        return ptr;
+    }
+    void* mapped = nullptr;
+    const auto err = cudaHostGetDevicePointer(&mapped, ptr, 0);
+    RuntimeCheck(err == cudaSuccess,
+        "host tensor must be pinned+mapped (cudaHostGetDevicePointer: ",
+        cudaGetErrorString(err), ")");
+    return mapped;
+}
+
+}  // namespace host

python/freetoken/kernel/csrc/jit/fast_index_copy.cuh

@@ -1,3 +1,4 @@
+#include <freetoken/pinned.h>
 #include <freetoken/tensor.h>
 #include <freetoken/utils.cuh>
 #include <freetoken/utils.h>
@@ -162,37 +163,6 @@ __always_inline __device__ void store_vec(void* __restrict__ dst, const Tp& vec)
 
 
 
-// Pinned host memory is GPU-dereferenceable at its host VA only where UVA identity
-// holds (Linux). On Windows/WDDM, cudaHostRegister'd memory maps to a different device
-// address, so host-resident tensors are translated here -- the one point their pointer
-// enters kernel params. Cached once per process: FreeToken pins one CUDA device per
-// process (set at engine launch).
-inline bool host_ptr_identity() {
-    static const bool identity = [] {
-        int device = 0;
-        if (cudaGetDevice(&device) != cudaSuccess) {
-            return false;  // fail closed: translate (and surface errors), don't assume identity
-        }
-        int uva = 0, reg = 0;
-        cudaDeviceGetAttribute(&uva, cudaDevAttrUnifiedAddressing, device);
-        cudaDeviceGetAttribute(&reg, cudaDevAttrCanUseHostPointerForRegisteredMem, device);
-        return uva == 1 && reg == 1;
-    }();
-    return identity;
-}
-
-inline void* device_alias(void* ptr, DLDevice dev) {
-    if (dev.device_type == kDLCUDA || dev.device_type == kDLROCM || host_ptr_identity()) {
-        return ptr;
-    }
-    void* mapped = nullptr;
-    const auto err = cudaHostGetDevicePointer(&mapped, ptr, 0);
-    host::RuntimeCheck(err == cudaSuccess,
-        "fast_index_copy: host tensor must be pinned+mapped (cudaHostGetDevicePointer: ",
-        cudaGetErrorString(err), ")");
-    return mapped;
-}
-
 struct IndexKernelParams {
     void* __restrict__ dst;
     const void* __restrict__ indices_dst;

python/freetoken/kernel/csrc/jit/index.cu

@@ -1,3 +1,4 @@
+#include <freetoken/pinned.h>
 #include <freetoken/tensor.h>
 #include <freetoken/utils.cuh>
 #include <freetoken/utils.h>
@@ -114,9 +115,10 @@ struct IndexKernel {
     auto weights_dtype_ = SymbolicDType{};
     auto indices_dtype_ = SymbolicDType{};
 
+    // the table may live in pinned host memory, read zero-copy
     TensorMatcher({-1, D}) //
         .with_dtype(weights_dtype_)
-        .with_device<kDLCUDA, kDLROCM>(device_)
+        .with_device<kDLCUDA, kDLROCM, kDLCUDAHost, kDLROCMHost, kDLCPU>()
         .verify(weights);
     TensorMatcher({L, D}) //
         .with_dtype(weights_dtype_)
@@ -128,6 +130,10 @@ struct IndexKernel {
         .verify(indices);
 
     const auto device = device_.unwrap();
+    const auto weights_device = weights.device();
+    RuntimeCheck(weights_device.device_type != device.device_type ||
+                     weights_device.device_id == device.device_id,
+                 "IndexKernel: weights are on another GPU than the output.");
     const auto use_int32 = indices_dtype_.unwrap().bits == 32;
     const auto num_indices = L.unwrap();
     const auto entry_size = dtype_bytes(weights_dtype_.unwrap()) * D.unwrap();
@@ -144,7 +150,8 @@ struct IndexKernel {
     const auto num_blocks = div_ceil(num_warps, kWarpPerBlock);
     const auto params = IndexKernelParams{
         .output = static_cast<char *>(output.data_ptr()),
-        .weight = static_cast<const char *>(weights.data_ptr()),
+        .weight = static_cast<const char *>(
+            device_alias(weights.data_ptr(), weights_device)),
         .indice = indices.data_ptr(),
         .num_warps = num_warps,
     };

python/freetoken/kernel/ggml.py

@@ -149,7 +149,7 @@ def dequantize(w: torch.Tensor, ggml_type: int, k: int, dtype: torch.dtype = tor
 
 
 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)
+    out = torch.empty(ids.shape[0], k, dtype=dtype, device=ids.device)
     _module().get_rows(w, ggml_type, k, ids, out)
     return out
 

python/freetoken/kernel/index.py

@@ -46,7 +46,7 @@ def indexing(
     vocab_range: Tuple[int, int] | None = None,  # (start, length)
 ) -> torch.Tensor:
     if output is None:
-        output = weights.new_empty(indices.shape[0], weights.shape[1])
+        output = weights.new_empty(indices.shape[0], weights.shape[1], device=indices.device)
 
     element_size = weights.shape[1] * weights.element_size()
     module = _jit_index_module(element_size, num_splits=num_splits_for(element_size))

python/freetoken/layers/__init__.py

@@ -7,7 +7,7 @@
     swiglu_clamp_and_mul,
     swigluoai_and_mul,
 )
-from .base import BaseOP, OPList, StateLessOP
+from .base import BaseOP, OPList, StateLessOP, iter_ops
 from .embedding import ParallelLMHead, VocabParallelEmbedding
 from .linear import (
     LinearColParallelMerged,
@@ -39,6 +39,7 @@
     "BaseOP",
     "StateLessOP",
     "OPList",
+    "iter_ops",
     "VocabParallelEmbedding",
     "ParallelLMHead",
     "LinearColParallelMerged",

python/freetoken/layers/base.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import Any, Dict, Generic, List, TypeAlias, TypeVar
+from typing import Any, Dict, Generic, Iterator, List, TypeAlias, TypeVar
 
 import torch
 
@@ -97,3 +97,23 @@ def load_state_dict(
 
         if not _internal and state_dict:
             raise RuntimeError(f"Unexpected keys in state_dict: {list(state_dict.keys())}")
+
+
+def iter_ops(root: BaseOP) -> Iterator[BaseOP]:
+    """Each op under ``root`` once, ``root`` first; an op's attributes are read after it is yielded."""
+    seen: set[int] = set()
+
+    def walk(op: BaseOP) -> Iterator[BaseOP]:
+        if id(op) in seen:
+            return
+        seen.add(id(op))
+        yield op
+        for value in vars(op).values():
+            if isinstance(value, BaseOP):
+                yield from walk(value)
+            elif isinstance(value, (list, tuple)):
+                for item in value:
+                    if isinstance(item, BaseOP):
+                        yield from walk(item)
+
+    return walk(root)

python/freetoken/layers/quantization/method.py

@@ -63,27 +63,12 @@ def scheme(self):
 
 def finalize_quant(root: Any) -> int:
     """Call ``quant_method.finalize(layer)`` on every layer under ``root``; returns the count."""
-    from freetoken.layers.base import BaseOP
+    from freetoken.layers.base import iter_ops
 
-    seen: set[int] = set()
     count = 0
-
-    def walk(op: Any) -> None:
-        nonlocal count
-        if id(op) in seen:
-            return
-        seen.add(id(op))
+    for op in iter_ops(root):
         method = getattr(op, "quant_method", None)
         if method is not None:
             method.finalize(op)
             count += 1
-        for value in vars(op).values():
-            if isinstance(value, BaseOP):
-                walk(value)
-            elif isinstance(value, (list, tuple)):
-                for item in value:
-                    if isinstance(item, BaseOP):
-                        walk(item)
-
-    walk(root)
     return count