Benchmarking

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:

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)

src/exo/shared/types/api.py

@@ -225,7 +225,7 @@ class ChatCompletionRequest(BaseModel):
 
 
 class BenchChatCompletionRequest(ChatCompletionRequest):
-    pass
+    disaggregated: bool = False
 
 
 class AddCustomModelParams(BaseModel):

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

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
 

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:

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