From 9e6484a538d094eade74ac174cbc26b5ede22673 Mon Sep 17 00:00:00 2001 From: Oql <1692110604@qq.com> Date: Tue, 31 Mar 2026 17:50:22 +0800 Subject: [PATCH] [fix]: fix --numa-nodes handling (#1904) * [fix]: fix --numa-nodes handling --- kt-kernel/python/cli/commands/run.py | 38 +++++++++++++++++++++++----- kt-kernel/python/utils/amx.py | 14 ++++++---- kt-kernel/python/utils/llamafile.py | 1 + kt-kernel/python/utils/moe_kernel.py | 1 + third_party/sglang | 2 +- 5 files changed, 43 insertions(+), 13 deletions(-) diff --git a/kt-kernel/python/cli/commands/run.py b/kt-kernel/python/cli/commands/run.py index 08ad216b..abbd54a3 100644 --- a/kt-kernel/python/cli/commands/run.py +++ b/kt-kernel/python/cli/commands/run.py @@ -40,7 +40,14 @@ from kt_kernel.cli.utils.user_model_registry import UserModelRegistry @click.option("--port", "-p", type=int, default=None, help="Server port") @click.option("--gpu-experts", type=int, default=None, help="Number of GPU experts per layer") @click.option("--cpu-threads", type=int, default=None, help="Number of CPU inference threads") -@click.option("--numa-nodes", type=int, default=None, help="Number of NUMA nodes") +@click.option( + "--numa-nodes", + "numa_nodes", + type=int, + multiple=True, + default=(), + help="Number of KT threadpools, or explicit NUMA node IDs for each threadpool (e.g. --numa-nodes 2 or --numa-nodes 0 --numa-nodes 1)", +) @click.option( "--tensor-parallel-size", "--tp", "tensor_parallel_size", type=int, default=None, help="Tensor parallel size" ) @@ -82,7 +89,7 @@ def run( port: Optional[int], gpu_experts: Optional[int], cpu_threads: Optional[int], - numa_nodes: Optional[int], + numa_nodes: Optional[tuple[int, ...]], tensor_parallel_size: Optional[int], model_path: Optional[str], weights_path: Optional[str], @@ -166,7 +173,7 @@ def _run_impl( port: Optional[int], gpu_experts: Optional[int], cpu_threads: Optional[int], - numa_nodes: Optional[int], + numa_nodes: Optional[tuple[int, ...]], tensor_parallel_size: Optional[int], model_path: Optional[Path], weights_path: Optional[Path], @@ -226,7 +233,7 @@ def _run_impl( gpu_experts is None or tensor_parallel_size is None or cpu_threads is None - or numa_nodes is None + or not numa_nodes or max_total_tokens is None ): # Model specified but some parameters missing - use interactive @@ -254,7 +261,10 @@ def _run_impl( # Extract parameters gpu_experts = config["gpu_experts"] cpu_threads = config["cpu_threads"] - numa_nodes = config["numa_nodes"] + if config.get("numa_nodes") is not None: + numa_nodes = (int(config["numa_nodes"]),) + else: + numa_nodes = () tensor_parallel_size = config["tp_size"] # Get kt-method and other method-specific settings @@ -421,7 +431,15 @@ def _run_impl( # CPU/GPU configuration with smart defaults total_threads = cpu.threads # Use logical threads instead of physical cores final_cpu_threads = resolve(cpu_threads, "inference.cpu_threads", int(total_threads * 0.8)) - final_numa_nodes = resolve(numa_nodes, "inference.numa_nodes", cpu.numa_nodes) + final_numa_nodes = resolve(None, "inference.numa_nodes", cpu.numa_nodes) + final_kt_numa_nodes = None + + if numa_nodes: + if len(numa_nodes) == 1: + final_numa_nodes = numa_nodes[0] + else: + final_kt_numa_nodes = list(numa_nodes) + final_numa_nodes = len(final_kt_numa_nodes) final_gpu_experts = resolve(gpu_experts, "inference.gpu_experts", 1) # KT-kernel options @@ -436,7 +454,6 @@ def _run_impl( final_mem_fraction_static = resolve(mem_fraction_static, "inference.mem_fraction_static", 0.98) final_watchdog_timeout = resolve(watchdog_timeout, "inference.watchdog_timeout", 3000) final_served_model_name = resolve(served_model_name, "inference.served_model_name", "") - # Performance flags final_disable_shared_experts_fusion = resolve( disable_shared_experts_fusion, "inference.disable_shared_experts_fusion", True @@ -472,6 +489,7 @@ def _run_impl( watchdog_timeout=final_watchdog_timeout, served_model_name=final_served_model_name, disable_shared_experts_fusion=final_disable_shared_experts_fusion, + kt_numa_nodes=final_kt_numa_nodes, tool_call_parser=final_tool_call_parser, reasoning_parser=final_reasoning_parser, settings=settings, @@ -503,6 +521,8 @@ def _run_impl( console.print(f" GPU Experts: [cyan]{final_gpu_experts}[/cyan] per layer") console.print(f" CPU Threads (kt-cpuinfer): [cyan]{final_cpu_threads}[/cyan]") console.print(f" NUMA Nodes (kt-threadpool-count): [cyan]{final_numa_nodes}[/cyan]") + if final_kt_numa_nodes is not None: + console.print(f" NUMA Nodes (binding): [cyan]{', '.join(map(str, final_kt_numa_nodes))}[/cyan]") console.print(f" Tensor Parallel: [cyan]{final_tensor_parallel_size}[/cyan]") console.print(f" Method: [cyan]{final_kt_method}[/cyan]") console.print(f" Attention: [cyan]{final_attention_backend}[/cyan]") @@ -577,6 +597,7 @@ def _build_sglang_command( watchdog_timeout: int, served_model_name: str, disable_shared_experts_fusion: bool, + kt_numa_nodes: Optional[list[int]], tool_call_parser: Optional[str], reasoning_parser: Optional[str], settings, @@ -632,6 +653,8 @@ def _build_sglang_command( "--kt-enable-dynamic-expert-update", # Enable dynamic expert updates ] ) + if kt_numa_nodes is not None: + cmd.extend(["--kt-numa-nodes", *map(str, kt_numa_nodes)]) # Add SGLang options cmd.extend( @@ -681,6 +704,7 @@ def _build_sglang_command( "kt-num-gpu-experts", "kt-cpuinfer", "kt-threadpool-count", + "kt-numa-nodes", "kt-method", "kt-gpu-prefill-token-threshold", "attention-backend", diff --git a/kt-kernel/python/utils/amx.py b/kt-kernel/python/utils/amx.py index d8b96ce9..4fa399a6 100644 --- a/kt-kernel/python/utils/amx.py +++ b/kt-kernel/python/utils/amx.py @@ -5,7 +5,13 @@ from typing import List, Optional # Use relative imports for package structure from ..experts_base import BaseMoEWrapper -from .loader import SafeTensorLoader, CompressedSafeTensorLoader, FP8SafeTensorLoader, BF16SafeTensorLoader, GPTQSafeTensorLoader +from .loader import ( + SafeTensorLoader, + CompressedSafeTensorLoader, + FP8SafeTensorLoader, + BF16SafeTensorLoader, + GPTQSafeTensorLoader, +) from kt_kernel_ext.moe import MOEConfig import kt_kernel_ext.moe as _moe_mod @@ -351,6 +357,7 @@ class NativeMoEWrapper(BaseMoEWrapper): cpu_save: bool = False, max_deferred_experts_per_token: Optional[int] = None, method: str = "RAWINT4", + numa_nodes: Optional[List[int]] = None, ): if method == "RAWINT4" and not _HAS_RAWINT4_SUPPORT: raise RuntimeError( @@ -379,10 +386,7 @@ class NativeMoEWrapper(BaseMoEWrapper): "Please recompile kt_kernel_ext with AVX512+BF16 or AVX2 enabled." ) if method == "GPTQ_INT4" and not _HAS_AVX2_GPTQ_INT4_SUPPORT: - raise RuntimeError( - "GPTQ_INT4 backend not available.\n" - "Please recompile kt_kernel_ext with AVX2 enabled." - ) + raise RuntimeError("GPTQ_INT4 backend not available.\n" "Please recompile kt_kernel_ext with AVX2 enabled.") super().__init__( layer_idx=layer_idx, diff --git a/kt-kernel/python/utils/llamafile.py b/kt-kernel/python/utils/llamafile.py index 66ebbeaa..0f1c4403 100644 --- a/kt-kernel/python/utils/llamafile.py +++ b/kt-kernel/python/utils/llamafile.py @@ -41,6 +41,7 @@ class LlamafileMoEWrapper(BaseMoEWrapper): cpu_save: bool = False, max_deferred_experts_per_token: Optional[int] = None, method: str = "LLAMAFILE", + numa_nodes: Optional[List[int]] = None, ): """ Initialize Llamafile MoE Wrapper. diff --git a/kt-kernel/python/utils/moe_kernel.py b/kt-kernel/python/utils/moe_kernel.py index 87326f35..6818c6ab 100644 --- a/kt-kernel/python/utils/moe_kernel.py +++ b/kt-kernel/python/utils/moe_kernel.py @@ -49,6 +49,7 @@ class GeneralMoEWrapper(BaseMoEWrapper): cpu_save: bool = False, max_deferred_experts_per_token: Optional[int] = None, method: str = "MOE_INT8", + numa_nodes: Optional[List[int]] = None, ): """ Initialize general MoE Wrapper. diff --git a/third_party/sglang b/third_party/sglang index 6b5135e1..a53cb6d0 160000 --- a/third_party/sglang +++ b/third_party/sglang @@ -1 +1 @@ -Subproject commit 6b5135e1a1360f20e104eaa12218068e30b12058 +Subproject commit a53cb6d0788bf045974a723d4bbb84e99ac2b6e3