mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-04-29 04:09:52 +00:00
(kt-kernel): add numa_nodes parameter for explicit NUMA node mapping (#1891)
Add numa_nodes parameter to BaseMoEWrapper and all subclasses, allowing users to explicitly specify which NUMA node IDs to use for subpool mapping instead of always defaulting to sequential [0, 1, ..., N-1]. This enables running multiple KTransformers instances on different NUMA nodes of the same machine, e.g. --kt-threadpool-count 1 --kt-numa-nodes 1 to bind to NUMA node 1. Previously this required external numactl workarounds since subpool_numa_map was hardcoded to start from 0.
This commit is contained in:
parent
bdf4bb76c5
commit
3903c9afcc
5 changed files with 34 additions and 6 deletions
|
|
@ -1,7 +1,7 @@
|
|||
import os
|
||||
import torch
|
||||
import ctypes
|
||||
from typing import Optional
|
||||
from typing import List, Optional
|
||||
|
||||
# Use relative imports for package structure
|
||||
from ..experts_base import BaseMoEWrapper
|
||||
|
|
@ -97,6 +97,7 @@ class GeneralMoEWrapper(BaseMoEWrapper):
|
|||
cpu_save=cpu_save,
|
||||
max_deferred_experts_per_token=max_deferred_experts_per_token,
|
||||
method=method,
|
||||
numa_nodes=numa_nodes,
|
||||
)
|
||||
|
||||
# moe-specific: Check if we should load merged safetensor weights
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue