Fix --mem-fraction-static not accounting for EAGLE draft model KV cache (#23862)

2026-06-12 10:35:48 -07:00
parent 533b59d00c
commit 75998d0421
22 changed files with 688 additions and 295 deletions
@@ -321,6 +321,8 @@ def load_model(server_args, port_args, gpu_id, tp_rank):
        model_runner = MlxModelRunnerStub(**runner_kwargs)
    else:
        model_runner = ModelRunner(**runner_kwargs)
+        model_runner.alloc_memory_pool()
+        model_runner.init_backends()
    rank_print(f"max_total_num_tokens={model_runner.max_total_num_tokens}")
    tokenizer = get_tokenizer(
        server_args.tokenizer_path,