danielkrause/colibri-strix
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Unify continuous batching + heterogeneous runtime: decode batching, physical-core planning, disjoint VRAM/RAM placement, topp-policy warning (CPU-validated, CUDA on 6x5090) (#68)

Browse Source 
* Fuse CUDA expert MLP execution

* Group CUDA expert transfers by device

* Instrument grouped CUDA expert execution

* Bound grouped CUDA decode scratch

* Execute expert groups across GPUs in parallel

* Release host backing for multi-GPU experts

* Define quality-preserving memory policies

* Overlap cold expert loading with resident compute

* Adapt expert placement with session LFRU

* Fuse q4 expert gate and up dispatch

* Plan CPU work on physical cores

* Batch grouped expert CUDA kernels

* Separate VRAM and RAM expert placement

* Add ragged multi-sequence decode forward

* feat(runtime): add continuous decode scheduler

* Route concurrent API requests through batch scheduler

* Harden multiplex request lifecycle and framing

* Cancel disconnected multiplex requests

* Bind API port before starting the engine

* fix automatic KV slot allocation

* add native int4 Tensor Core grouped GEMM

* add Tensor Core throughput benchmark

* optimize packed int4 low-row kernels

* add asynchronous CUDA staging streams

* document validated six-GPU dense acceleration

* tune six-GPU expert hot set

* raise validated expert hot-set target

* add CUDA MLA absorption core

* fuse grouped expert gate and up projections

* Warn for explicit lossy routing flags
This commit is contained in:
ZacharyZcR
2026-07-13 20:30:36 +08:00
committed by GitHub
parent 98759bfc40
commit cbd599024e
20 changed files with 1741 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/tests/test_openai_server.py
+241 -4
View File
@@ -1,19 +1,24 @@
import io
import json
import math
import socket
import threading
import unittest
from unittest.mock import patch
from urllib.error import HTTPError
from urllib.request import Request, urlopen
from openai_server import (APIError, APIServer, ClientCancelled, END, GenerationScheduler,
generation_options, read_engine_turn, render_chat, serve)
READY, Engine, generation_options, read_engine_turn, render_chat,
serve)
class FakeEngine:
def __init__(self):
self.calls = []
def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0):
def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0,
cancelled=None):
self.calls.append((prompt, maximum, temperature, top_p, cache_slot))
on_text("Hé")
on_text("llo")
@@ -26,10 +31,12 @@ class BlockingEngine(FakeEngine):
self.entered = threading.Event()
self.release = threading.Event()
def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0):
def generate(self, prompt, maximum, temperature, top_p, on_text, cache_slot=0,
cancelled=None):
self.entered.set()
self.release.wait(2)
return super().generate(prompt, maximum, temperature, top_p, on_text, cache_slot)
return super().generate(prompt, maximum, temperature, top_p, on_text, cache_slot,
cancelled)
class TemplateTest(unittest.TestCase):
@@ -65,6 +72,10 @@ class TemplateTest(unittest.TestCase):
(4, 0.0, 1.0))
with self.assertRaises(APIError):
generation_options({"max_tokens": 9}, 8)
with self.assertRaises(APIError):
generation_options({"temperature": math.nan}, 8)
with self.assertRaises(APIError):
generation_options({"top_p": math.inf}, 8)
self.assertEqual(generation_options({"temperature": None, "top_p": None}, 8),
(8, 0.7, 0.9))
@@ -82,8 +93,27 @@ class ProtocolTest(unittest.TestCase):
with self.assertRaisesRegex(ValueError, "kv_slots"):
serve("/missing", kv_slots=0)
def test_occupied_port_fails_before_engine_start(self):
listener = socket.socket()
listener.bind(("127.0.0.1", 0))
listener.listen()
try:
with patch("openai_server.subprocess.Popen") as popen:
with self.assertRaises(OSError):
serve("/missing", port=listener.getsockname()[1])
popen.assert_not_called()
finally:
listener.close()
class SchedulerTest(unittest.TestCase):
def test_admits_up_to_capacity_without_serializing(self):
scheduler = GenerationScheduler(max_queue=0, queue_timeout=1, capacity=2)
with scheduler.admit() as first:
with scheduler.admit() as second:
self.assertEqual({first[1], second[1]}, {0, 1})
self.assertEqual(scheduler.snapshot()["active"], 2)
def test_rejects_when_waiting_queue_is_full(self):
scheduler = GenerationScheduler(max_queue=0, queue_timeout=1)
with scheduler.admit():
@@ -160,6 +190,213 @@ class SchedulerTest(unittest.TestCase):
self.assertEqual(errors, ["scheduler_closed"])
class BlockingStream:
def __init__(self, initial=b""):
self.buffer = bytearray(initial)
self.closed = False
self.condition = threading.Condition()
def feed(self, data):
with self.condition:
self.buffer.extend(data)
self.condition.notify_all()
def read(self, size=1):
with self.condition:
while len(self.buffer) < size and not self.closed:
self.condition.wait()
if not self.buffer and self.closed:
return b""
size = min(size, len(self.buffer))
data = bytes(self.buffer[:size])
del self.buffer[:size]
return data
def readline(self):
with self.condition:
while b"\n" not in self.buffer and not self.closed:
self.condition.wait()
if not self.buffer and self.closed:
return b""
end = self.buffer.find(b"\n")
size = len(self.buffer) if end < 0 else end + 1
data = bytes(self.buffer[:size])
del self.buffer[:size]
return data
def close(self):
with self.condition:
self.closed = True
self.condition.notify_all()
class FakeProcess:
def __init__(self, on_write):
self.stdout = BlockingStream(READY + b"STAT 0 0 0 0\n")
self.stdin = self
self.on_write = on_write
self.writes = []
self.returncode = None
def write(self, data):
self.writes.append(data)
self.on_write(self, data)
return len(data)
def flush(self):
pass
def poll(self):
return self.returncode
def terminate(self):
self.returncode = 0
self.stdout.close()
def wait(self, timeout=None):
return self.returncode
def kill(self):
self.terminate()
class DispatcherTest(unittest.TestCase):
def test_dispatches_interleaved_requests_by_id(self):
submitted = []
def respond(process, frame):
fields = frame.split(b"\n", 1)[0].split()
self.assertEqual(fields[0], b"SUBMIT")
submitted.append(fields[1])
if len(submitted) == 2:
first, second = submitted
process.stdout.feed(b"DATA " + second + b" 3\nB-2\n")
process.stdout.feed(b"DATA " + first + b" 3\nA-1\n")
process.stdout.feed(b"DONE " + second + b" STAT 1 2.5 0 1.0 4 0\n")
process.stdout.feed(b"DATA " + first + b" 3\nA-2\n")
process.stdout.feed(b"DONE " + first + b" STAT 2 3.5 0 1.0 5 1\n")
process = FakeProcess(respond)
with patch("openai_server.subprocess.Popen", return_value=process):
engine = Engine("glm", "model", kv_slots=2)
results = {}
def generate(name, prompt, slot):
chunks = []
stats = engine.generate(prompt, 8, 0.7, 0.9, chunks.append, slot)
results[name] = ("".join(chunks), stats)
threads = [threading.Thread(target=generate, args=("a", "alpha", 0)),
threading.Thread(target=generate, args=("b", "beta", 1))]
for thread in threads:
thread.start()
for thread in threads:
thread.join(timeout=2)
self.assertFalse(thread.is_alive())
engine.close()
self.assertEqual(results["a"][0], "A-1A-2")
self.assertTrue(results["a"][1]["length_limited"])
self.assertEqual(results["b"][0], "B-2")
headers = [frame.split(b"\n", 1)[0].split() for frame in process.writes]
self.assertEqual({int(header[2]) for header in headers}, {0, 1})
self.assertEqual({header[3] for header in headers}, {b"4", b"5"})
def test_routes_engine_error_to_request(self):
def respond(process, frame):
request_id = frame.split()[1]
process.stdout.feed(b"ERROR " + request_id + b" slot is busy\n")
process = FakeProcess(respond)
with patch("openai_server.subprocess.Popen", return_value=process):
engine = Engine("glm", "model")
with self.assertRaisesRegex(RuntimeError, "slot is busy"):
engine.generate("hello", 4, 0.7, 0.9, lambda _: None)
engine.close()
def test_close_wakes_pending_generation_and_is_idempotent(self):
process = FakeProcess(lambda _process, _frame: None)
with patch("openai_server.subprocess.Popen", return_value=process):
engine = Engine("glm", "model")
errors = []
def generate():
try:
engine.generate("hello", 4, 0.7, 0.9, lambda _: None)
except RuntimeError as error:
errors.append(str(error))
thread = threading.Thread(target=generate)
thread.start()
for _ in range(100):
with engine.pending_lock:
if engine.pending:
break
threading.Event().wait(0.01)
engine.close()
engine.close()
thread.join(timeout=2)
self.assertFalse(thread.is_alive())
self.assertEqual(errors, ["colibri engine is shutting down"])
self.assertFalse(engine.dispatcher.is_alive())
with engine.pending_lock:
self.assertFalse(engine.pending)
with self.assertRaisesRegex(RuntimeError, "shutting down"):
engine.generate("again", 4, 0.7, 0.9, lambda _: None)
def test_protocol_corruption_fails_request_and_stops_dispatcher(self):
def respond(process, frame):
request_id = frame.split()[1]
process.stdout.feed(b"DATA " + request_id + b" -1\n")
process = FakeProcess(respond)
with patch("openai_server.subprocess.Popen", return_value=process):
engine = Engine("glm", "model")
with self.assertRaisesRegex(RuntimeError, "DATA size"):
engine.generate("hello", 4, 0.7, 0.9, lambda _: None)
with self.assertRaisesRegex(RuntimeError, "dispatcher stopped"):
engine.generate("again", 4, 0.7, 0.9, lambda _: None)
engine.close()
def test_decodes_utf8_split_across_data_frames(self):
def respond(process, frame):
request_id = frame.split()[1]
process.stdout.feed(b"DATA " + request_id + b" 1\n\xc3\n")
process.stdout.feed(b"DATA " + request_id + b" 1\n\xa9\n")
process.stdout.feed(b"DONE " + request_id + b" STAT 1 1 0 1 1 0\n")
process = FakeProcess(respond)
with patch("openai_server.subprocess.Popen", return_value=process):
engine = Engine("glm", "model")
chunks = []
engine.generate("hello", 4, 0.7, 0.9, chunks.append)
engine.close()
self.assertEqual(chunks, ["é"])
def test_cancels_generation_after_consumer_disconnects(self):
request_id = None
def respond(process, frame):
nonlocal request_id
fields = frame.split()
if fields[0] == b"SUBMIT":
request_id = fields[1]
process.stdout.feed(b"DATA " + request_id + b" 1\nx\n")
elif fields[0] == b"CANCEL":
self.assertEqual(fields[1], request_id)
process.stdout.feed(b"ERROR " + request_id + b" CANCELLED\n")
process = FakeProcess(respond)
with patch("openai_server.subprocess.Popen", return_value=process):
engine = Engine("glm", "model")
output = []
with self.assertRaises(ClientCancelled):
engine.generate("hello", 8, 0.7, 0.9, output.append, cancelled=lambda: True)
engine.close()
self.assertEqual(output, ["x"])
self.assertEqual(process.writes[-1].split(), [b"CANCEL", request_id])
class HTTPTest(unittest.TestCase):
@classmethod
def setUpClass(cls):