Fix O(n²) text concatenation in MetalModelRunner.generate() (#60)

This PR is:

- To avoid O(n²) string concatenation in non‑v1
vllm_metal/model_runner.py::MetalModelRunner.generate() by collecting
streamed segments and joining once
- To add unit tests for streamed segment accumulation (v1 and legacy
runners) and refactor the tests into clearer, separated files

Note: similar to one of the previous PR I submitted; just changing str
+= to .join()

---------

Signed-off-by: Yuan Lik Xun <lxyuan0420@gmail.com>

Author: LxYuan0420

tests/test_model_runner_generate.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import vllm_metal.model_runner as mr
+
+
+class TestMetalModelRunnerGenerate:
+    def _make_runner(self) -> mr.MetalModelRunner:
+        runner = mr.MetalModelRunner.__new__(mr.MetalModelRunner)
+        runner.model = object()
+        runner.tokenizer = object()
+        return runner
+
+    def test_accumulates_streamed_segments(self, monkeypatch) -> None:
+        captured: dict[str, object] = {}
+
+        def fake_make_sampler(*, temp: float):
+            captured["temp"] = temp
+            return object()
+
+        def fake_stream_generate(model, tokenizer, prompt, max_tokens=256, **kwargs):
+            captured["prompt"] = prompt
+            captured["max_tokens"] = max_tokens
+            captured["kwargs"] = kwargs
+            yield SimpleNamespace(text="hello")
+            yield SimpleNamespace(text=" ")
+            yield SimpleNamespace(text="world")
+
+        monkeypatch.setattr(mr, "make_sampler", fake_make_sampler)
+        monkeypatch.setattr(mr, "stream_generate", fake_stream_generate)
+
+        runner = self._make_runner()
+        out = runner.generate("p", max_tokens=3, temperature=0.7)
+
+        assert out == "hello world"
+        assert captured["prompt"] == "p"
+        assert captured["max_tokens"] == 3
+        assert captured["temp"] == 0.7
+        assert "sampler" in captured["kwargs"]

tests/test_v1_model_runner_generate.py

@@ -7,51 +7,50 @@
 import vllm_metal.v1.model_runner as mr
 
 
-def _make_runner(mr_module):
-    runner = mr_module.MetalModelRunner.__new__(mr_module.MetalModelRunner)
-    runner.model = object()
-    runner.tokenizer = object()
-    return runner
-
-
-def test_generate_accumulates_streamed_segments(monkeypatch) -> None:
-    captured: dict[str, object] = {}
-
-    def fake_stream_generate(model, tokenizer, prompt, max_tokens=256, **kwargs):
-        captured["prompt"] = prompt
-        captured["max_tokens"] = max_tokens
-        captured["kwargs"] = kwargs
-        yield SimpleNamespace(text="hello")
-        yield SimpleNamespace(text=" ")
-        yield SimpleNamespace(text="world")
-
-    monkeypatch.setattr(mr, "stream_generate", fake_stream_generate)
-
-    runner = _make_runner(mr)
-    out = runner.generate("p", max_tokens=3, temperature=0.0)
-
-    assert out == "hello world"
-    assert captured["prompt"] == "p"
-    assert captured["max_tokens"] == 3
-    # mlx_lm 0.29+ uses sampler parameter instead of temp
-    assert "sampler" in captured["kwargs"]
-    assert callable(captured["kwargs"]["sampler"])
-
-
-def test_generate_passes_sampler_for_temperature_sampling(monkeypatch) -> None:
-    captured: dict[str, object] = {}
-
-    def fake_stream_generate(model, tokenizer, prompt, max_tokens=256, **kwargs):
-        captured["kwargs"] = kwargs
-        assert "sampler" in kwargs
-        assert callable(kwargs["sampler"])
-        yield SimpleNamespace(text="a")
-        yield SimpleNamespace(text="b")
-
-    monkeypatch.setattr(mr, "stream_generate", fake_stream_generate)
-
-    runner = _make_runner(mr)
-    out = runner.generate("p", max_tokens=2, temperature=0.5)
-
-    assert out == "ab"
-    assert "sampler" in captured["kwargs"]
+class TestV1MetalModelRunnerGenerate:
+    def _make_runner(self) -> mr.MetalModelRunner:
+        runner = mr.MetalModelRunner.__new__(mr.MetalModelRunner)
+        runner.model = object()
+        runner.tokenizer = object()
+        return runner
+
+    def test_accumulates_streamed_segments(self, monkeypatch) -> None:
+        captured: dict[str, object] = {}
+
+        def fake_stream_generate(model, tokenizer, prompt, max_tokens=256, **kwargs):
+            captured["prompt"] = prompt
+            captured["max_tokens"] = max_tokens
+            captured["kwargs"] = kwargs
+            yield SimpleNamespace(text="hello")
+            yield SimpleNamespace(text=" ")
+            yield SimpleNamespace(text="world")
+
+        monkeypatch.setattr(mr, "stream_generate", fake_stream_generate)
+
+        runner = self._make_runner()
+        out = runner.generate("p", max_tokens=3, temperature=0.0)
+
+        assert out == "hello world"
+        assert captured["prompt"] == "p"
+        assert captured["max_tokens"] == 3
+        # mlx_lm 0.29+ uses sampler parameter instead of temp
+        assert "sampler" in captured["kwargs"]
+        assert callable(captured["kwargs"]["sampler"])
+
+    def test_passes_sampler_for_temperature_sampling(self, monkeypatch) -> None:
+        captured: dict[str, object] = {}
+
+        def fake_stream_generate(model, tokenizer, prompt, max_tokens=256, **kwargs):
+            captured["kwargs"] = kwargs
+            assert "sampler" in kwargs
+            assert callable(kwargs["sampler"])
+            yield SimpleNamespace(text="a")
+            yield SimpleNamespace(text="b")
+
+        monkeypatch.setattr(mr, "stream_generate", fake_stream_generate)
+
+        runner = self._make_runner()
+        out = runner.generate("p", max_tokens=2, temperature=0.5)
+
+        assert out == "ab"
+        assert "sampler" in captured["kwargs"]

vllm_metal/model_runner.py

@@ -229,7 +229,7 @@ def generate(
             raise RuntimeError(msg)
 
         # Generate tokens using stream_generate
-        generated_text = ""
+        segments: list[str] = []
 
         # Create sampler with temperature
         sampler = make_sampler(temp=temperature)
@@ -242,9 +242,9 @@ def generate(
             sampler=sampler,
         ):
             # Accumulate incremental text from each token
-            generated_text += response.text
+            segments.append(response.text)
 
-        return generated_text
+        return "".join(segments)
 
     def __del__(self) -> None:
         """Cleanup model resources."""