refactor: reduce parameter passing

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,
     )
 
     # ---------------------------------------------------------------------------- #

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:

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,50 +58,23 @@ 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
         )
     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,
     )

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,