"""
Implementation du moteur LLM Mirza basee sur MTPLX et Qwen 3.8-9B Distill MTPLX (Q6G64).
Fournit une acceleration native MTP (Multi-Token Prediction) sur Apple Silicon M4 via le runtime MTPLX.
"""
from __future__ import annotations
import json
import time
import urllib.error
import urllib.request
from collections.abc import Generator
from ...config import LlmConfig
from ...core.logger import getMirzaLogger
from ..interfaces.baseLlm import BaseLlm
from .mtplxProcessSupervisor import MtplxProcessSupervisor
from .repetitionGuard import RepetitionGuard
logger = getMirzaLogger("mirza.llm.mtplx")
try:
import orjson
def _fast_json_dumps(obj: object) -> bytes:
return orjson.dumps(obj)
def _fast_json_loads(data: bytes | str) -> object:
return orjson.loads(data)
except ImportError:
def _fast_json_dumps(obj: object) -> bytes:
return json.dumps(obj).encode("utf-8")
def _fast_json_loads(data: bytes | str) -> object:
return json.loads(data)
[docs]
class MtplxLlm(BaseLlm):
"""
Moteur de generation haute performance utilisant le runtime MTPLX.
Parameters
----------
modelConfig : LlmConfig
Configuration du moteur LLM.
"""
#: Le daemon prend un effort a trois crans, pas un nombre de jetons : ses
#: paliers sont traduits en effort par reasoningEffortFor.
reasoningSteps = (0, 512, 1024, 2048)
defaultReasoningBudget = 1024
reasoningParameter = "reasoning_effort"
[docs]
def __init__(self, modelConfig: LlmConfig) -> None:
self.config = modelConfig
self.host = getattr(self.config, "mtplxHost", "127.0.0.1")
self.port = getattr(self.config, "mtplxPort", 8766)
self.baseUrl = f"http://{self.host}:{self.port}/v1"
self.modelRepo = getattr(self.config, "mtplxModelRepo", "pooyakn73/Qwen3.8-9B-Distill-MTPLX-Q6G64")
self.autoStart = getattr(self.config, "mtplxAutoStart", True)
contextWindow = getattr(self.config, "maxKvCacheTokens", 16384)
kvBits = getattr(self.config, "kvBits", 8)
self.supervisor = MtplxProcessSupervisor(
host=self.host,
port=self.port,
modelRepo=self.modelRepo,
contextWindow=contextWindow,
kvBits=kvBits,
mtpDepth=getattr(self.config, "optiqMtpDepth", 0),
)
self.isLoaded = False
self.estimatedKvTokens = 0
self.currentSpeedTokPerSec = 0.0
self.lastTokensPerSec = 0.0
self.lastTtftMs = 0.0
self.lastTokenCount = 0
self.reasoningBudget = getattr(self.config, "reasoningBudget", 0)
self.enableThinking = bool(getattr(self.config, "enableThinking", False) and (self.reasoningBudget > 0))
self.sessionId = getattr(self.config, "mtplxSessionId", "mirza-session")
[docs]
def loadModel(self, systemPrompt: str = "") -> None:
"""
Charge le modele en s'assurant que le daemon MTPLX est actif.
Parameters
----------
systemPrompt : str, optional
Instructions systeme d'initialisation. Defaut : "".
Raises
------
ConnectionError
Si le daemon est injoignable et le demarrage automatique desactive.
"""
if not self.isLoaded:
if not self.supervisor.isHealthy() and self.autoStart:
self.supervisor.start()
elif not self.supervisor.isHealthy():
raise ConnectionError(f"Le serveur MTPLX n'est pas accessible sur {self.baseUrl}.")
logger.info("Modele MTPLX pret sur %s (Modele: %s)", self.baseUrl, self.modelRepo)
self.isLoaded = True
if systemPrompt:
self.prefill(systemPrompt)
[docs]
def prefill(self, systemPrompt: str = "") -> None:
"""
Pré-amorce le graphe de calcul Metal et le cache de session avec le prompt système.
Compile les shaders à l'avance et initialise les tables KV pour éliminer
le délai TTFT à froid lors du premier échange utilisateur.
Parameters
----------
systemPrompt : str, optional
Le prompt système complet à pré-remplir.
"""
if not systemPrompt:
return
logger.info("Pré-amorçage (prefill) du prompt système dans MTPLX...")
t0 = time.time()
try:
messages = [
{"role": "system", "content": systemPrompt},
{"role": "user", "content": "bonjour"},
]
payload = {
"model": self.modelRepo,
"messages": messages,
"session_id": self.sessionId,
"max_tokens": 1,
"temperature": 0.0,
"stream": False,
}
req = urllib.request.Request(
f"{self.baseUrl}/chat/completions",
data=_fast_json_dumps(payload),
headers={
"Content-Type": "application/json",
"User-Agent": "Mirza-Engine/1.0",
},
)
with urllib.request.urlopen(req, timeout=30) as resp:
resp.read()
dtMs = (time.time() - t0) * 1000
logger.info("Pré-amorçage MTPLX terminé en %.1f ms (shaders compilés, session amorcée).", dtMs)
except Exception as e:
logger.debug("Note pré-amorçage MTPLX : %s", e)
[docs]
def setThinkingMode(self, enabled: bool) -> None:
"""
Active ou desactive dynamiquement le mode reflexion (CoT).
Parameters
----------
enabled : bool
True pour le palier par defaut, False pour le couper.
"""
self.setReasoningBudget(self.defaultReasoningBudget if enabled else 0)
[docs]
def setReasoningBudget(self, budget: int) -> None:
"""
Definit le budget de tokens alloue a la reflexion.
Parameters
----------
budget : int
Budget en tokens. 0 desactive le mode reflexion.
"""
self.reasoningBudget = max(0, budget)
self.enableThinking = self.reasoningBudget > 0
logger.debug("Budget de reflexion MTPLX : %d tokens (actif: %s)", self.reasoningBudget, self.enableThinking)
[docs]
def clearCache(self) -> None:
"""
Reinitialise les etats de session sur le serveur MTPLX.
"""
# Un cache indisponible n'est pas une erreur : le calcul repart a froid.
try:
req = urllib.request.Request(
f"http://{self.host}:{self.port}/admin/cache/clear",
data=b"",
headers={"User-Agent": "Mirza-Client/1.0"},
method="POST",
)
with urllib.request.urlopen(req, timeout=1.0):
pass
logger.debug("Cache KV MTPLX reinitialise")
except Exception as e:
logger.debug("Reinitialisation cache MTPLX ignoree : %s", e)
[docs]
def streamResponse(self, formattedPrompt: str) -> Generator[str]:
"""
Genere les tokens en streaming temps reel via SSE depuis MTPLX.
Parameters
----------
formattedPrompt : str
Prompt produit par :meth:`formatPrompt`.
Yields
------
str
Fragments de texte. La reflexion est encadree par ``<think>`` et
``</think>`` pour que le filtre de flux la route vers le canal de
reflexion et ne fasse entendre que la reponse finale.
"""
if not self.isLoaded:
self.loadModel()
try:
messages = _fast_json_loads(formattedPrompt)
if not isinstance(messages, list):
messages = [{"role": "user", "content": str(formattedPrompt)}]
except Exception:
messages = [{"role": "user", "content": str(formattedPrompt)}]
maxToks = self.resolveMaxTokens(self.config.maxTokens, self.config.maxTokensThinking)
# Le cran vient des paliers de ce moteur : une comparaison ecrite en dur
# ne suivait pas le niveau choisi.
effort = self.reasoningEffortFor(self.reasoningBudget)
payload = {
"model": self.modelRepo,
"messages": messages,
"session_id": self.sessionId,
"max_tokens": maxToks,
"temperature": self.config.temperature,
"top_p": self.config.topP,
# Les penalties manquaient ici : la presence_penalty est le seul
# frein qui casse une boucle sans fin, et elle n'etait envoye par
# aucun des deux moteurs HTTP.
"presence_penalty": self.config.presencePenalty,
"repetition_penalty": self.config.repetitionPenalty,
"enable_thinking": self.enableThinking,
"stream": True,
"stream_options": {"include_usage": True},
}
if self.enableThinking:
payload["reasoning_effort"] = effort
req = urllib.request.Request(
f"{self.baseUrl}/chat/completions",
data=_fast_json_dumps(payload),
headers={
"Content-Type": "application/json",
"User-Agent": "Mirza-Engine/1.0",
},
)
startTime = time.time()
firstTokenTime = None
tokenCounter = 0
totalYieldedChars = 0
thinkingActive = False
reasoningGuard = RepetitionGuard()
degenerate = False
nativeSpeedTokPerSec = 0.0
nativeTtftMs = 0.0
try:
with urllib.request.urlopen(req, timeout=120) as resp:
for rawLine in resp:
line = rawLine.decode("utf-8").strip()
if not line or not line.startswith("data: "):
continue
if line == "data: [DONE]":
break
dataStr = line[6:].strip()
try:
chunkJson = _fast_json_loads(dataStr)
if "mtplx_progress" in chunkJson:
p = chunkJson["mtplx_progress"]
if "tok_s" in p and p["tok_s"] is not None:
self.currentSpeedTokPerSec = float(p["tok_s"])
if chunkJson.get("timings"):
t = chunkJson["timings"]
if "predicted_per_second" in t and t["predicted_per_second"] is not None:
nativeSpeedTokPerSec = float(t["predicted_per_second"])
if "prompt_ms" in t and t["prompt_ms"] is not None:
nativeTtftMs = float(t["prompt_ms"])
if chunkJson.get("mtplx_stats"):
s = chunkJson["mtplx_stats"]
if "decode_tok_s" in s and s["decode_tok_s"] is not None:
nativeSpeedTokPerSec = float(s["decode_tok_s"])
elif "tok_s" in s and s["tok_s"] is not None and nativeSpeedTokPerSec <= 0.0:
nativeSpeedTokPerSec = float(s["tok_s"])
if "ttft_s" in s and s["ttft_s"] is not None:
nativeTtftMs = float(s["ttft_s"]) * 1000.0
if chunkJson.get("usage"):
u = chunkJson["usage"]
if "completion_tokens" in u and u["completion_tokens"] is not None:
tokenCounter = max(tokenCounter, int(u["completion_tokens"]))
if "prompt_tokens" in u and u["prompt_tokens"] is not None:
self.estimatedKvTokens = int(u["prompt_tokens"]) + tokenCounter
choices = chunkJson.get("choices", [])
if not choices:
continue
delta = choices[0].get("delta", {})
reasoningChunk = delta.get("reasoning_content")
contentChunk = delta.get("content")
if reasoningChunk:
if firstTokenTime is None:
firstTokenTime = time.time()
tokenCounter += 1
totalYieldedChars += len(reasoningChunk)
reasoningGuard.feed(reasoningChunk)
if not thinkingActive:
thinkingActive = True
yield "<think>"
yield reasoningChunk
if reasoningGuard.triggered:
logger.warning(
"Reflexion degenerante coupee sur %r "
"(%d repetitions)",
reasoningGuard.triggeredLine,
reasoningGuard.maxRepeats,
)
degenerate = True
break
elif contentChunk:
if firstTokenTime is None:
firstTokenTime = time.time()
tokenCounter += 1
totalYieldedChars += len(contentChunk)
if thinkingActive:
thinkingActive = False
yield "</think>"
yield contentChunk
except Exception as exc:
logger.debug("Erreur parsing chunk SSE MTPLX: %s", exc)
continue
if degenerate:
break
if thinkingActive:
yield "</think>"
except Exception as e:
logger.error("Erreur de streaming MTPLX : %s", e)
yield f" [Erreur d'inférence MTPLX : {e}]"
finally:
totalDuration = time.time() - startTime
if nativeTtftMs > 0.0:
self.lastTtftMs = round(nativeTtftMs, 1)
elif firstTokenTime is not None:
self.lastTtftMs = round((firstTokenTime - startTime) * 1000.0, 1)
else:
self.lastTtftMs = round(totalDuration * 1000.0, 1)
if tokenCounter == 0 and totalYieldedChars > 0:
tokenCounter = max(1, totalYieldedChars // 4)
self.lastTokenCount = tokenCounter
if nativeSpeedTokPerSec > 0.0:
self.currentSpeedTokPerSec = round(nativeSpeedTokPerSec, 1)
else:
decodeDuration = (time.time() - firstTokenTime) if firstTokenTime else totalDuration
self.currentSpeedTokPerSec = round(tokenCounter / decodeDuration, 1) if decodeDuration > 0 else 0.0
self.lastTokensPerSec = self.currentSpeedTokPerSec
[docs]
def unload(self) -> None:
"""
Arrete le daemon MTPLX, qu'il ait ete lance ou adopte par cette instance.
"""
self.supervisor.stop()
self.isLoaded = False