From 335a0c128c4cc688acf8af814f5bee9368474c2a Mon Sep 17 00:00:00 2001 From: wooway777 Date: Wed, 8 Jul 2026 14:47:27 +0800 Subject: [PATCH] refactor: reduce parameter passing --- examples/bench.py | 36 ++--- examples/bench_videonsa.py | 23 +--- examples/test_infer.py | 87 +----------- python/infinilm/llm/llm.py | 71 ++++++++-- python/infinilm/llm/scheduler.py | 22 +++- python/infinilm/server/inference_server.py | 146 +++------------------ test/bench/test_benchmark.py | 50 +++---- 7 files changed, 155 insertions(+), 280 deletions(-) diff --git a/examples/bench.py b/examples/bench.py index bd424e836..a77ba4462 100644 --- a/examples/bench.py +++ b/examples/bench.py @@ -1,3 +1,4 @@ +import copy import json import os import sys @@ -197,6 +198,7 @@ def __init__( weight_load_mode="async", moe_ep_backend="disabled", moe_ep_size=1, + config=None, ) -> None: model_path = os.path.expanduser(model_path) self.draft_model_path = draft_model_path @@ -210,6 +212,7 @@ def __init__( self.use_mla = use_mla self.weight_load_mode = weight_load_mode self.skip_load = skip_load + self.config = config if draft_model_path is not None: self.processor = AutoInfinilmProcessor.from_pretrained(model_path) @@ -299,24 +302,17 @@ def run( # ---------------------------------------------------------------------------- # if self.draft_model_path is not None: prompt_text = self.tokenizer.decode(input_ids, skip_special_tokens=False) - llm = LLM( - model_path=self.model_path, - draft_model_path=self.draft_model_path, - num_draft_tokens=self.num_draft_tokens, - device=self.device_str, - tensor_parallel_size=self.tp, - cache_type="paged" if self.cache_config is not None else "static", - max_batch_size=batch_size, - max_tokens=output_len, - temperature=temperature, - top_p=top_p, - top_k=top_k, - enable_graph=self.enable_graph, - attn_backend=self.attn_backend, - use_mla=self.use_mla, - weight_load_mode=self.weight_load_mode, - skip_load=self.skip_load, - ) + if self.config is None: + raise ValueError( + "TestModel requires config when draft_model_path is set" + ) + llm_config = copy.copy(self.config) + llm_config.max_batch_size = batch_size + llm_config.max_new_tokens = output_len + llm_config.temperature = temperature + llm_config.top_p = top_p + llm_config.top_k = top_k + llm = LLM(llm_config) t1 = time.time() print("=================== start generate ====================") outputs = llm.generate( @@ -387,6 +383,9 @@ def run( skip_load = cfg.skip_load + cfg.moe_ep_backend = moe_ep_backend + cfg.ep = ep + batch_size = cfg.batch_size input_len = cfg.input_len output_len = cfg.output_len @@ -442,6 +441,7 @@ def run( weight_load_mode=cfg.weight_load_mode, moe_ep_backend=moe_ep_backend, moe_ep_size=ep, + config=cfg, ) # ---------------------------------------------------------------------------- # diff --git a/examples/bench_videonsa.py b/examples/bench_videonsa.py index faf0bf7f7..53c255c24 100644 --- a/examples/bench_videonsa.py +++ b/examples/bench_videonsa.py @@ -10,7 +10,6 @@ from infinilm.processors import AutoInfinilmProcessor from infinilm.processors.videonsa_processor import decode_video_frames - VIDEO_AUTO_MIN_FRAMES = 4 VIDEO_AUTO_MAX_FRAMES = 8 VIDEO_AUTO_SAMPLE_FPS = 1.0 @@ -186,7 +185,9 @@ def main(): output_lens = as_int_list(cfg.output_len) max_batch_size = max(int(cfg.batch_size), int(cfg.max_batch_size)) max_cache_len = max(max(input_lens) + max(output_lens) + 4096, cfg.max_cache_len) - cache_type = "paged" if cfg.enable_paged_attn else "static" + cfg.max_batch_size = max_batch_size + cfg.max_new_tokens = max(output_lens) + cfg.max_cache_len = max_cache_len video_meta = apply_video_auto_args(cfg) apply_multimodal_env(cfg) @@ -203,23 +204,7 @@ def main(): ) processor = AutoInfinilmProcessor.from_pretrained(cfg.model) tokenizer = processor.get_tokenizer() - model = LLM( - model_path=cfg.model, - device=cfg.get_device_str(cfg.device), - tensor_parallel_size=cfg.tp, - cache_type=cache_type, - max_batch_size=max_batch_size, - max_tokens=max(output_lens), - num_blocks=cfg.num_blocks, - block_size=cfg.block_size, - max_cache_len=max_cache_len, - temperature=cfg.temperature, - top_p=cfg.top_p, - top_k=cfg.top_k, - attn_backend=cfg.attn, - enable_graph=cfg.enable_graph, - weight_load_mode=cfg.weight_load_mode, - ) + model = LLM(cfg) for input_len in input_lens: for output_len in output_lens: diff --git a/examples/test_infer.py b/examples/test_infer.py index b02500422..941ae87c4 100644 --- a/examples/test_infer.py +++ b/examples/test_infer.py @@ -1,4 +1,3 @@ -import os import time from infinilm.base_config import BaseConfig @@ -9,60 +8,15 @@ def test( prompts: list[str], - model_path, - draft_model_path=None, - num_draft_tokens=4, - max_new_tokens=100, - device="cpu", - tp=1, - moe_ep_backend="disabled", - ep=1, - enable_paged_attn=False, - enable_graph=False, - num_blocks=512, - block_size=256, - top_k=1, - top_p=1.0, - temperature=1.0, - attn_backend="default", - use_mla=False, + config, image_path=None, video_path=None, video_num_frames=None, - skip_load=False, - weight_load_mode="async", - skip_legacy_moe=False, ): - model_path = os.path.expanduser(model_path) # ---------------------------------------------------------------------------- # # Create Model # ---------------------------------------------------------------------------- # - if enable_paged_attn and attn_backend == "default": - attn_backend = "paged-attn" - - model = LLM( - model_path=model_path, - draft_model_path=draft_model_path, - num_draft_tokens=num_draft_tokens, - device=device, - tensor_parallel_size=tp, - moe_ep_backend=moe_ep_backend, - moe_ep_size=ep, - cache_type="paged" if enable_paged_attn else "static", - max_batch_size=len(prompts), - max_tokens=max_new_tokens, - num_blocks=num_blocks, - block_size=block_size, - temperature=temperature, - top_k=top_k, - top_p=top_p, - enable_graph=enable_graph, - attn_backend=attn_backend, - use_mla=use_mla, - skip_load=skip_load, - weight_load_mode=weight_load_mode, - skip_legacy_moe=skip_legacy_moe, - ) + model = LLM(config) conversations = [ [{"role": "user", "content": [{"type": "text", "text": prompt}]}] @@ -104,20 +58,8 @@ def test( if __name__ == "__main__": cfg = BaseConfig() - device_str = cfg.get_device_str(cfg.device) - prompts = [cfg.prompt for _ in range(cfg.batch_size)] - model_path = cfg.model - - max_new_tokens = cfg.max_new_tokens - - tp = cfg.tp - - enable_paged_attn = cfg.enable_paged_attn - - enable_graph = cfg.enable_graph - if cfg.skip_legacy_moe: moe_ep_backend, ep = configure_moe_ep_backend( cfg.tp, cfg.dp, cfg.ep, cfg.moe_ep_backend, cfg.model @@ -125,29 +67,14 @@ def test( else: moe_ep_backend, ep = "disabled", 1 + cfg.moe_ep_backend = moe_ep_backend + cfg.ep = ep + cfg.max_batch_size = len(prompts) + test( prompts, - model_path, - draft_model_path=cfg.draft_model, - num_draft_tokens=cfg.num_draft_tokens, - max_new_tokens=max_new_tokens, - device=device_str, - tp=tp, - moe_ep_backend=moe_ep_backend, - ep=ep, - enable_paged_attn=enable_paged_attn, - enable_graph=enable_graph, - num_blocks=cfg.num_blocks, - block_size=cfg.block_size, - top_k=cfg.top_k, - top_p=cfg.top_p, - temperature=cfg.temperature, - attn_backend=cfg.attn, - use_mla=cfg.use_mla, + cfg, image_path=cfg.image, video_path=cfg.video, video_num_frames=cfg.video_num_frames, - skip_load=cfg.skip_load, - weight_load_mode=cfg.weight_load_mode, - skip_legacy_moe=cfg.skip_legacy_moe, ) diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59e9e94ba..1e8974cde 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -34,6 +34,61 @@ logger = logging.getLogger(__name__) +def _is_config_like(value) -> bool: + return isinstance(value, EngineConfig) or hasattr(value, "model") + + +def _normalize_engine_config(value, *, kv_transfer_config=None) -> EngineConfig: + if isinstance(value, EngineConfig): + if kv_transfer_config is None: + return value + return EngineConfig( + **{**value.__dict__, "kv_transfer_config": kv_transfer_config} + ) + if hasattr(value, "model"): + return EngineConfig( + model_path=os.path.expanduser(value.model), + draft_model_path=value.draft_model, + num_draft_tokens=value.num_draft_tokens, + device=value.get_device_str(value.device), + dtype=value.dtype, + tensor_parallel_size=value.tp, + moe_ep_backend=value.moe_ep_backend, + moe_ep_size=value.ep or 1, + cache_type="paged" if value.enable_paged_attn else "static", + max_batch_size=value.max_batch_size, + max_tokens=value.max_new_tokens, + num_blocks=value.num_blocks, + block_size=value.block_size, + max_cache_len=value.max_cache_len, + temperature=value.temperature, + top_p=value.top_p, + top_k=value.top_k, + enable_graph=value.enable_graph, + attn_backend=value.attn, + kv_transfer_config=kv_transfer_config, + use_mla=value.use_mla, + weight_load_mode=value.weight_load_mode, + skip_load=value.skip_load, + skip_legacy_moe=value.skip_legacy_moe, + ) + raise TypeError(f"Unsupported engine config type: {type(value)!r}") + + +def _build_engine_config( + model_path, *, kv_transfer_config=None, **kwargs +) -> EngineConfig: + if _is_config_like(model_path): + return _normalize_engine_config( + model_path, kv_transfer_config=kv_transfer_config + ) + return EngineConfig( + model_path=model_path, + kv_transfer_config=kv_transfer_config, + **kwargs, + ) + + class LLMEngine: """Low-level LLM engine that handles inference execution.""" @@ -88,9 +143,7 @@ def __init__(self, config: EngineConfig): assert 1024 <= max_num_batched_tokens <= max_position_embeddings self.scheduler = Scheduler( - max_batch_size=config.max_batch_size, - num_blocks=config.num_blocks, - block_size=config.block_size, + config=config, max_num_batched_tokens=max_num_batched_tokens, connector=connector, has_mamba_cache=has_mamba_cache, @@ -310,7 +363,7 @@ class LLM: def __init__( self, - model_path: str, + model_path: Union[str, EngineConfig, object], draft_model_path: Optional[str] = None, num_draft_tokens: int = 4, device: str = "cuda", @@ -355,8 +408,8 @@ def __init__( use_mla: Whether to use DeepSeek V2 MLA attention when supported. weight_load_mode: Weight loading mode across tensor-parallel workers. """ - config = EngineConfig( - model_path=model_path, + config = _build_engine_config( + model_path, draft_model_path=draft_model_path, num_draft_tokens=num_draft_tokens, device=device, @@ -517,7 +570,7 @@ class AsyncLLMEngine: def __init__( self, - model_path: str, + model_path: Union[str, EngineConfig, object], draft_model_path: Optional[str] = None, num_draft_tokens: int = 4, device: str = "cuda", @@ -565,8 +618,8 @@ def __init__( use_mla: Whether to use DeepSeek V2 MLA attention when supported. weight_load_mode: Weight loading mode across tensor-parallel workers. """ - config = EngineConfig( - model_path=model_path, + config = _build_engine_config( + model_path, draft_model_path=draft_model_path, num_draft_tokens=num_draft_tokens, device=device, diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index a1153e916..af842c03b 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -8,6 +8,7 @@ import janus +from infinilm.config.engine_config import EngineConfig from infinilm.llm.cache_manager import BlockManager, MambaCacheManager from infinilm.llm.request import InferenceRequest, RequestStatus @@ -71,9 +72,7 @@ class Scheduler: def __init__( self, - max_batch_size: int = 16, - num_blocks: int = 512, - block_size: int = 256, + config: EngineConfig, max_num_batched_tokens: int = 1024, connector=None, has_mamba_cache: bool = False, @@ -81,7 +80,7 @@ def __init__( ): self.waiting_queue = janus.Queue() self.running_queue = janus.Queue() - self.max_batch_size = max_batch_size + self.config = config self.finished_receiving_kv_req_ids: set[str] = set() self.failed_receiving_kv_req_ids: set[str] = set() @@ -89,18 +88,27 @@ def __init__( self.pending_kv_decode_blocks: int = 0 self.remote_kv_requests: dict[str, InferenceRequest] = {} - self.cache_manager = BlockManager(num_blocks=num_blocks, block_size=block_size) + self.cache_manager = BlockManager( + num_blocks=config.num_blocks, block_size=config.block_size + ) self.has_mamba_cache = has_mamba_cache self.mamba_cache_manager = ( - MambaCacheManager(num_mamba_cache_blocks or max(2, num_blocks // 4)) + MambaCacheManager(num_mamba_cache_blocks or max(2, config.num_blocks // 4)) if has_mamba_cache else None ) self.speculative_cache_ops = SpeculativeCacheOps(self.cache_manager) - self.block_size = block_size self.max_num_batched_tokens = max_num_batched_tokens self.connector = connector + @property + def max_batch_size(self) -> int: + return self.config.max_batch_size + + @property + def block_size(self) -> int: + return self.config.block_size + def add_request(self, request: InferenceRequest): if request is not None: request.status = RequestStatus.WAITING diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 242593e25..b707fc63d 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -94,83 +94,20 @@ class InferenceServer: def __init__( self, - model_path: str, - device: str = "cuda", - dtype: str = "float16", - tensor_parallel_size: int = 1, - moe_ep_backend: str = "disabled", - moe_ep_size: int = 1, - cache_type: str = "paged", - max_tokens: int = 4096, - max_batch_size: int = 16, - num_blocks: int = 512, - block_size: int = 256, - max_cache_len: int = 4096, - temperature: float = 1.0, - top_p: float = 0.8, - top_k: int = 1, - host: str = "0.0.0.0", - port: int = 8000, - enable_graph: bool = False, - attn_backend: str = "default", - use_mla: bool = False, - weight_load_mode: str = "async", - ignore_eos: bool = False, + config: BaseConfig, kv_transfer_config: Optional[KVTransferConfig] = None, ): - """Initialize inference server. - - Args: - model_path: Path to the model directory. - device: Device type ('cpu', 'cuda', 'mlu', 'moore'). - dtype: Data type ('float16', 'bfloat16', 'float32'). - tensor_parallel_size: Number of devices for tensor parallelism. - moe_ep_backend: MoE expert-parallel backend. - moe_ep_size: MoE expert-parallel size. - cache_type: Cache type ('paged' or 'static'). - max_tokens: Default maximum tokens to generate. - max_batch_size: Maximum batch size for inference (only for paged cache). - num_blocks: Number of KV cache blocks (only for paged cache). - block_size: Size of each KV cache block (only for paged cache). - max_cache_len: Maximum sequence length (only for static cache). - temperature: Default sampling temperature. - top_p: Default top-p sampling parameter. - top_k: Default top-k sampling parameter. - host: Server host address. - port: Server port number. - enable_graph: Whether to enable graph compiling. - attn_backend: Attention backend to use ('default', 'flash-attn'). - use_mla: Whether to use DeepSeek V2 MLA attention when supported. - weight_load_mode: Weight loading mode across tensor-parallel workers. - ignore_eos: Whether to ignore EOS tokens during generation. - kv_transfer_config: Optional configuration for the KV transfer mechanism. - """ - self.model_path = model_path - # vLLM-like served model id: directory name of model_path - self.model_id = os.path.basename(os.path.normpath(model_path)) or "model" - self.device = device - self.dtype = dtype - self.tensor_parallel_size = tensor_parallel_size - self.moe_ep_backend = moe_ep_backend - self.moe_ep_size = moe_ep_size - self.cache_type = cache_type - self.max_tokens = max_tokens - self.max_batch_size = max_batch_size - self.num_blocks = num_blocks - self.block_size = block_size - self.max_cache_len = max_cache_len - self.temperature = temperature - self.top_p = top_p - self.top_k = top_k - self.host = host - self.port = port - self.enable_graph = enable_graph - self.attn_backend = attn_backend - self.use_mla = use_mla - self.weight_load_mode = weight_load_mode - self.ignore_eos = ignore_eos + """Initialize inference server.""" + self.config = config self.kv_transfer_config = kv_transfer_config + self.model_path = config.model + self.model_id = os.path.basename(os.path.normpath(config.model)) or "model" + self.host = config.host + self.port = config.port + self.ignore_eos = config.ignore_eos + self.enable_graph = config.enable_graph + self.engine: AsyncLLMEngine = None def start(self): @@ -186,26 +123,7 @@ def _create_app(self): @asynccontextmanager async def lifespan(app: FastAPI): self.engine = AsyncLLMEngine( - model_path=self.model_path, - device=self.device, - dtype=self.dtype, - tensor_parallel_size=self.tensor_parallel_size, - moe_ep_backend=self.moe_ep_backend, - moe_ep_size=self.moe_ep_size, - cache_type=self.cache_type, - max_batch_size=self.max_batch_size, - max_tokens=self.max_tokens, - num_blocks=self.num_blocks, - block_size=self.block_size, - max_cache_len=self.max_cache_len, - temperature=self.temperature, - top_p=self.top_p, - top_k=self.top_k, - enable_graph=self.enable_graph, - attn_backend=self.attn_backend, - use_mla=self.use_mla, - weight_load_mode=self.weight_load_mode, - kv_transfer_config=self.kv_transfer_config, + self.config, kv_transfer_config=self.kv_transfer_config ) self.engine.start() logger.info(f"Engine initialized with model at {self.model_path}") @@ -328,6 +246,7 @@ def _build_sampling_params(self, data: dict) -> SamplingParams: # Support both: # - top-level OpenAI-ish fields: temperature/top_p/top_k/max_tokens/stop # - nested dict: sampling_params: { ... } + config = self.config sp = data.get("sampling_params") or {} if not isinstance(sp, dict): sp = {} @@ -341,19 +260,19 @@ def pick(key: str, default): return default # Accept common alias - max_tokens = pick("max_tokens", self.max_tokens) + max_tokens = pick("max_tokens", config.max_new_tokens) if max_tokens is None: # Some clients use max_new_tokens - max_tokens = pick("max_new_tokens", self.max_tokens) + max_tokens = pick("max_new_tokens", config.max_new_tokens) stop = pick("stop", None) if isinstance(stop, str): stop = [stop] return SamplingParams( - temperature=float(pick("temperature", self.temperature)), - top_p=float(pick("top_p", self.top_p)), - top_k=int(pick("top_k", self.top_k)), + temperature=float(pick("temperature", config.temperature)), + top_p=float(pick("top_p", config.top_p)), + top_k=int(pick("top_k", config.top_k)), max_tokens=int(max_tokens) if max_tokens is not None else None, stop=stop, ignore_eos=self.ignore_eos, @@ -583,8 +502,6 @@ def parse_kv_transfer_config(kv_transfer_config_str: str) -> KVTransferConfig: def main(): cfg = BaseConfig() setup_logging(cfg.log_level) - device = cfg.get_device_str(cfg.device) - kv_transfer_config = None if cfg.kv_transfer_config: kv_transfer_config = parse_kv_transfer_config(cfg.kv_transfer_config) @@ -600,31 +517,10 @@ def main(): ep, ) - server = InferenceServer( - model_path=cfg.model, - device=device, - dtype=cfg.dtype, - tensor_parallel_size=cfg.tp, - moe_ep_backend=moe_ep_backend, - moe_ep_size=ep, - cache_type="paged" if cfg.enable_paged_attn else "static", - max_tokens=cfg.max_new_tokens, - max_batch_size=cfg.max_batch_size, - num_blocks=cfg.num_blocks, - block_size=cfg.block_size, - max_cache_len=cfg.max_cache_len, - temperature=cfg.temperature, - top_p=cfg.top_p, - top_k=cfg.top_k, - host=cfg.host, - port=cfg.port, - enable_graph=cfg.enable_graph, - attn_backend=cfg.attn, - use_mla=cfg.use_mla, - weight_load_mode=cfg.weight_load_mode, - ignore_eos=cfg.ignore_eos, - kv_transfer_config=kv_transfer_config, - ) + cfg.moe_ep_backend = moe_ep_backend + cfg.ep = ep + + server = InferenceServer(cfg, kv_transfer_config=kv_transfer_config) server.start() diff --git a/test/bench/test_benchmark.py b/test/bench/test_benchmark.py index 60e63dc66..0279a289b 100644 --- a/test/bench/test_benchmark.py +++ b/test/bench/test_benchmark.py @@ -1,13 +1,13 @@ -import sys -import os -import time -import re -import csv import argparse +import csv import json -import numpy as np -from datasets import load_dataset, Dataset +import os +import re +import time from abc import ABC, abstractmethod + +import numpy as np +from datasets import Dataset, load_dataset from infinilm.base_config import BaseConfig from infinilm.processors import AutoInfinilmProcessor @@ -56,12 +56,21 @@ def __init__( enable_graph=False, attn_backend="default", ): - import transformers import infinicore - from infinilm.modeling_utils import load_model_state_dict_by_file + from infinilm.cache import PagedKVCacheConfig, StaticKVCacheConfig from infinilm.distributed import DistConfig - from infinilm.cache import StaticKVCacheConfig, PagedKVCacheConfig from infinilm.infer_engine import InferEngine + from infinilm.modeling_utils import load_model_state_dict_by_file + + device_name = None + if hasattr(model_dir_path, "model"): + config = model_dir_path + model_dir_path = config.model + device_name = config.get_device_str(config.device) + ndev = config.tp + enable_paged_attn = config.enable_paged_attn + enable_graph = config.enable_graph + attn_backend = config.attn self.benchmark = benchmark @@ -81,7 +90,8 @@ def __init__( "ali": "cuda", } - device_name = device_map.get(device_type_str.lower(), "cpu") + if device_name is None: + device_name = device_map.get(device_type_str.lower(), "cpu") # CUDA_VISIBLE_DEVICES is automatically respected by CUDA runtime API # When CUDA_VISIBLE_DEVICES=5 is set, CUDA only sees device 5 as device 0 # So device index 0 will automatically map to the first visible device @@ -298,9 +308,10 @@ def render_input_content(self, *args, **kwargs): raise ValueError(f"Unknown benchmark: {self.benchmark}") def _generate_step(self, tokens, max_steps, topp_, topk_, temperature_): - import torch import time + import torch + input_ids = torch.tensor([tokens], device=self.device) if self.device.type == "cuda": @@ -763,7 +774,7 @@ def load_one(subj): continue if not all_samples: raise FileNotFoundError( - f"No MMLU cached data found for any subject. Please ensure datasets are cached." + "No MMLU cached data found for any subject. Please ensure datasets are cached." ) return all_samples, "all" @@ -1005,7 +1016,7 @@ def _load_mmlu_subject(subj): continue if not samples: raise FileNotFoundError( - f"No MMLU data found for any subject in the list" + "No MMLU data found for any subject in the list" ) return samples, "all" else: @@ -1118,14 +1129,9 @@ def main(): model = VLLMBenchmark(cfg.model, device_type_str, cfg.tp, cfg.bench) else: # cpp backend model = InfiniLMBenchmark( - cfg.model, - device_type_str, - cfg.tp, - cfg.backend, - cfg.bench, - cfg.enable_paged_attn, - cfg.enable_graph, - cfg.attn, + cfg, + backend=cfg.backend, + benchmark=cfg.bench, ) # Step 3: Evaluate each subject