From 5d22805a77b5ee47b4bfb1ff3ab524bcf4d3c983 Mon Sep 17 00:00:00 2001 From: Ryuichi Leo Takashige Date: Wed, 1 Apr 2026 19:22:45 +0100 Subject: [PATCH] Benchmarking --- bench/harness.py | 4 ++++ src/exo/master/api.py | 4 +++- src/exo/shared/types/api.py | 2 +- src/exo/shared/types/text_generation.py | 1 + .../worker/engines/mlx/generator/batch_generate.py | 2 +- src/exo/worker/engines/vllm/vllm_generator.py | 12 +++++++++--- .../worker/runner/llm_inference/batch_generator.py | 10 ++++++++-- 7 files changed, 27 insertions(+), 8 deletions(-) diff --git a/bench/harness.py b/bench/harness.py index 573cf99d..bf5fece9 100644 --- a/bench/harness.py +++ b/bench/harness.py @@ -69,6 +69,10 @@ class ExoClient: def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]: return self.request_json("POST", "/bench/chat/completions", body=payload) + def post_bench_disaggregated(self, payload: dict[str, Any]) -> dict[str, Any]: + payload["disaggregated"] = True + return self.request_json("POST", "/bench/chat/completions", body=payload) + def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]: if len(instance) != 1: diff --git a/src/exo/master/api.py b/src/exo/master/api.py index 3810f743..3819c2ce 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -741,7 +741,9 @@ class API: ) task_params = task_params.model_copy(update={"model": resolved_model}) - task_params = task_params.model_copy(update={"stream": False, "bench": True}) + task_params = task_params.model_copy( + update={"stream": False, "bench": True, **({"disaggregated_bench": True} if payload.disaggregated else {})} + ) command = TextGeneration(task_params=task_params) await self._send(command) diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py index 272a5c81..f9aed55f 100644 --- a/src/exo/shared/types/api.py +++ b/src/exo/shared/types/api.py @@ -225,7 +225,7 @@ class ChatCompletionRequest(BaseModel): class BenchChatCompletionRequest(ChatCompletionRequest): - pass + disaggregated: bool = False class AddCustomModelParams(BaseModel): diff --git a/src/exo/shared/types/text_generation.py b/src/exo/shared/types/text_generation.py index 1e407bee..e8b85c7b 100644 --- a/src/exo/shared/types/text_generation.py +++ b/src/exo/shared/types/text_generation.py @@ -71,3 +71,4 @@ class TextGenerationTaskParams(BaseModel, frozen=True): repetition_penalty: float | None = None repetition_context_size: int | None = None prefill_endpoints: list[str] | None = None + disaggregated_bench: bool = False diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py index 45e6c973..811664b6 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -157,7 +157,7 @@ class ExoBatchGenerator: if ( uncached_count > 1000 and task_params.prefill_endpoints - and not is_bench + and (not is_bench or task_params.disaggregated_bench) ): from exo.disaggregated.prefill_client import remote_prefill diff --git a/src/exo/worker/engines/vllm/vllm_generator.py b/src/exo/worker/engines/vllm/vllm_generator.py index 2ba5c639..d75b2d6c 100644 --- a/src/exo/worker/engines/vllm/vllm_generator.py +++ b/src/exo/worker/engines/vllm/vllm_generator.py @@ -584,16 +584,22 @@ def load_vllm_engine( is_nvfp4 = "nvfp4" in model_path.lower() or "nvfp4" in str(model_id).lower() has_mamba = False + is_mxfp4 = False config_path = Path(model_path) / "config.json" if config_path.exists(): with open(config_path) as f: model_config = json.load(f) text_config = model_config.get("text_config", model_config) has_mamba = "mamba_ssm_dtype" in text_config or "linear_attention" in (text_config.get("layer_types") or []) - if is_nvfp4 and not has_mamba: - backends = ["FLASHINFER", "FLASH_ATTN", "TRITON_ATTN"] - else: + quant_config = model_config.get("quantization_config") or text_config.get("quantization_config") + if quant_config and quant_config.get("quant_method") == "mxfp4": + is_mxfp4 = True + if is_mxfp4: + os.environ.setdefault("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8", "1") + if has_mamba: backends = ["FLASH_ATTN", "TRITON_ATTN"] + else: + backends = ["FLASHINFER", "FLASH_ATTN", "TRITON_ATTN"] engine: LLMEngine | None = None for backend in backends: diff --git a/src/exo/worker/runner/llm_inference/batch_generator.py b/src/exo/worker/runner/llm_inference/batch_generator.py index a37e637d..3de07ddd 100644 --- a/src/exo/worker/runner/llm_inference/batch_generator.py +++ b/src/exo/worker/runner/llm_inference/batch_generator.py @@ -293,7 +293,10 @@ class SequentialGenerator(InferenceGenerator): ) def close(self) -> None: - del self.tokenizer, self.group + if hasattr(self, "tokenizer"): + del self.tokenizer + if hasattr(self, "group"): + del self.group @dataclass(eq=False) @@ -507,4 +510,7 @@ class BatchGenerator(InferenceGenerator): def close(self) -> None: self._gen.close() - del self.tokenizer, self.group + if hasattr(self, "tokenizer"): + del self.tokenizer + if hasattr(self, "group"): + del self.group