diff --git a/docs/cli.md b/docs/cli.md index 138683b01..feda4f7a3 100644 --- a/docs/cli.md +++ b/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 | diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 23afd3573..77aaa3eef 100644 --- a/python/freetoken/engine/config.py +++ b/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" diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 7982095ac..2de7f47a3 100644 --- a/python/freetoken/engine/engine.py +++ b/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: diff --git a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu b/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu index 170cbcac7..444b154c6 100644 --- a/python/freetoken/kernel/csrc/ggml/freetoken/ggml_ops.cu +++ b/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 #include #include @@ -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); } diff --git a/python/freetoken/kernel/csrc/include/freetoken/pinned.h b/python/freetoken/kernel/csrc/include/freetoken/pinned.h new file mode 100644 index 000000000..7c744de4a --- /dev/null +++ b/python/freetoken/kernel/csrc/include/freetoken/pinned.h @@ -0,0 +1,40 @@ +#pragma once + +#include +#include + +#include + +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(®, 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 diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index 279b15ae7..053f0b551 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -1,3 +1,4 @@ +#include #include #include #include @@ -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(®, 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; diff --git a/python/freetoken/kernel/csrc/jit/index.cu b/python/freetoken/kernel/csrc/jit/index.cu index 4832e0a46..f197ce087 100644 --- a/python/freetoken/kernel/csrc/jit/index.cu +++ b/python/freetoken/kernel/csrc/jit/index.cu @@ -1,3 +1,4 @@ +#include #include #include #include @@ -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(device_) + .with_device() .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(output.data_ptr()), - .weight = static_cast(weights.data_ptr()), + .weight = static_cast( + device_alias(weights.data_ptr(), weights_device)), .indice = indices.data_ptr(), .num_warps = num_warps, }; diff --git a/python/freetoken/kernel/ggml.py b/python/freetoken/kernel/ggml.py index 7a2765c5c..589bc4504 100644 --- a/python/freetoken/kernel/ggml.py +++ b/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 diff --git a/python/freetoken/kernel/index.py b/python/freetoken/kernel/index.py index 95e61ff47..cbc29e285 100644 --- a/python/freetoken/kernel/index.py +++ b/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)) diff --git a/python/freetoken/layers/__init__.py b/python/freetoken/layers/__init__.py index a07484620..590b58abb 100644 --- a/python/freetoken/layers/__init__.py +++ b/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", diff --git a/python/freetoken/layers/base.py b/python/freetoken/layers/base.py index b1e939d98..a05f6636d 100644 --- a/python/freetoken/layers/base.py +++ b/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) diff --git a/python/freetoken/layers/quantization/method.py b/python/freetoken/layers/quantization/method.py index 2a079539f..9e9063757 100644 --- a/python/freetoken/layers/quantization/method.py +++ b/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 diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index efe3b8c0b..b3eea9f8a 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -689,6 +689,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--embed-weights", + default=ServerArgs.embed_weights, + choices=["gpu", "host"], + help=( + "Where the input token embedding table lives. 'gpu' (default) keeps it in VRAM; '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. " + "Not for a table tied to the LM head." + ), + ) + parser.add_argument( "--swa-decoder-replay", default=ServerArgs.swa_decoder_replay,