"""
Moteur de generation de texte MLX avec decodage speculatif MTP pour Qwen 3.5.
"""
from __future__ import annotations
import json
import time
from collections.abc import Generator
from typing import Any
from mlx_vlm import apc, load, stream_generate
from mlx_vlm.prompt_utils import apply_chat_template
from ...config import LlmConfig
from ...core.logger import getMirzaLogger
from ..interfaces.baseLlm import BaseLlm
logger = getMirzaLogger("mirza.llm.mlx")
#: Keys mlx_vlm's qwen3_5 config insists on inside rope_parameters.
REQUIRED_ROPE_KEYS = ("partial_rotary_factor", "type", "mrope_section", "rope_theta")
def _normaliseRopeParameters(config: dict) -> dict:
"""
Normalise les paramètres rope_parameters requis par mlx_vlm pour Qwen 3.5.
"""
textConfig = config.get("text_config")
if not isinstance(textConfig, dict):
return config
rope = textConfig.get("rope_parameters")
if not isinstance(rope, dict):
return config
missing = [key for key in REQUIRED_ROPE_KEYS if key not in rope]
if not missing:
return config
if "partial_rotary_factor" in missing and isinstance(rope.get("factor"), (int, float)):
try:
derived = 1.0 / float(rope["factor"])
except ZeroDivisionError:
return config
rope["partial_rotary_factor"] = derived
logger.info(
"rope_parameters : partial_rotary_factor derive de factor=%s -> %s",
rope["factor"],
derived,
)
return config
def _loadModelConfigDict(modelRepo: str) -> dict:
"""
Lit le ``config.json`` du modele, normalise pour mlx_vlm.
``mlx_vlm.load`` ne rend plus la config avec les poids, alors
qu'``apply_chat_template`` en a besoin (il lit ``model_type``). Lire le
fichier coute peu au regard des gigaoctets de poids.
Parameters
----------
modelRepo : str
L'identifiant de depot HuggingFace, ou un chemin local.
Returns
-------
dict
La config analysee, ou un dict vide si elle est illisible. Le modele
reste utilisable dans ce cas ; seul le gabarit se degrade.
"""
try:
from huggingface_hub import hf_hub_download
with open(hf_hub_download(modelRepo, "config.json")) as handle:
return _normaliseRopeParameters(json.load(handle))
except Exception as exc:
logger.warning("config.json de %s illisible (%s) ; gabarit de prompt degrade", modelRepo, exc)
return {}
[docs]
class MlxQwenLlm(BaseLlm):
"""
Implementation MLX locale de Qwen 3.5 (9B 5-bit).
Supporte le decodage speculatif (MTP).
Parameters
----------
modelConfig : LlmConfig
Configuration contenant les parametres du modele MLX et les
hyperparametres.
"""
[docs]
def __init__(self, modelConfig: LlmConfig) -> None:
self.config = modelConfig
self.modelRepo = getattr(modelConfig, "llmModelRepo", "")
self.model: Any = None
self.processor: Any = None
self.modelConfig: Any = None
self.apcManager: apc.APCManager | None = None
self.draftModel: Any = None
self.draftKind = getattr(self.config, "draftKind", "mtp")
self.promptCache: Any = None
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.is4bModel = bool(
"4b" in self.modelRepo.lower() or "siddhjagani" in self.modelRepo.lower()
)
self.isLoaded = False
[docs]
def setThinkingMode(self, enabled: bool) -> None:
"""
Active ou desactive dynamiquement le mode reflexion (CoT).
Parameters
----------
enabled : bool
True pour demander un raisonnement, False pour l'eviter.
"""
self.setReasoningBudget(512 if enabled else 0)
[docs]
def setReasoningBudget(self, budget: int) -> None:
"""
Definit a chaud le budget de tokens de 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 MLX : %d tokens (actif: %s)", self.reasoningBudget, self.enableThinking)
[docs]
def loadModel(self, systemPrompt: str = "") -> None:
"""
Charge le modele principal et le drafter MTP en memoire unifiee.
Parameters
----------
systemPrompt : str, optional
Prompt systeme utilise pour le prechauffage. Defaut : "".
"""
logger.info("Chargement du modele MLX : %s...", self.modelRepo)
# mlx_vlm rend (model, processor), et parfois un troisieme element de
# config. Un depaquetage fixe a 3 cassait tous les chargements, donc
# l'arite est toleree.
loaded = load(self.modelRepo, trust_remote_code=True)
self.model, self.processor = loaded[0], loaded[1]
# apply_chat_template lit config['model_type'] : il faut la config du
# modele, pas celle de l'application.
self.modelConfig = loaded[2] if len(loaded) > 2 else _loadModelConfigDict(self.modelRepo)
# Cache de prefixe automatique (APC)
enableApc = getattr(self.config, "enableApc", True)
if enableApc:
try:
block_size = getattr(self.config, "apcBlockSize", 64)
self.apcManager = apc.APCManager(block_size=block_size)
if self.is4bModel:
# Conserve le prefixe systeme ainsi que plusieurs tours en cache exact
self.apcManager._exact_cache_max = 8
logger.debug("Automatic Prefix Caching (APC) initialise (blocs de %d)", block_size)
except Exception as e:
logger.warning("Echec initialisation APCManager (%s)", e)
self.apcManager = None
if self.is4bModel:
# Pour le modele 4B, evite le prefill segmente (chunked) qui impose
# des synchronisations Metal / GPU barriers repetees sur les invites moyennes
self.config.prefillStepSize = max(getattr(self.config, "prefillStepSize", 512), 2048)
# Drafter pour le decodage speculatif MTP
enableSpeculative = getattr(self.config, "enableSpeculative", True)
draftRepo = getattr(self.config, "drafterModelRepo", None) or getattr(self.config, "draftModelRepo", None)
if enableSpeculative and draftRepo:
logger.info("Chargement du drafter MTP separe : %s...", draftRepo)
try:
draftLoaded = load(draftRepo, trust_remote_code=True)
self.draftModel = draftLoaded[0]
self.draftKind = getattr(self.config, "draftKind", self.draftKind)
logger.info("Drafter MTP separe actif.")
except (OSError, RuntimeError, ValueError) as e:
logger.warning("Drafter separe indisponible (%s). Sans speculation.", e)
self.draftModel = None
# Un profil unifie porte sa tete MTP dans les poids principaux, mais
# load_drafter y garde une SECONDE copie de ces couches : vram a 15,9 Go
# contre 9,4 Go de base, ce qui casse le budget d'un seul modele en RAM.
# La tete reste donc dormante plutot que de doubler la memoire en silence.
# Voir todo.md T-061.
else:
self.draftModel = None
self._warmUp(systemPrompt)
self.isLoaded = True
def _warmUp(self, systemPrompt: str = "") -> None:
"""
Prechauffe le graphe de calcul Metal et compile les shaders.
Parameters
----------
systemPrompt : str, optional
Prompt systeme utilise pour la passe a blanc. Defaut : "".
"""
# Pour le 4B, amorcage du prefixe systeme statique dans le cache exact APC
if self.is4bModel and systemPrompt:
sysPrompt = f"<|im_start|>system\n{systemPrompt}<|im_end|>\n"
try:
for _ in self.streamResponse(sysPrompt):
break
logger.debug("Prefixe systeme 4B amorce dans APC exact.")
except (OSError, RuntimeError, ValueError) as e:
logger.debug("Note warm-up prefixe systeme 4B : %s", e)
dummyHistory = [
{"role": "user", "content": "bonjour"},
{"role": "assistant", "content": "bonjour !"},
]
prompt = self.formatPrompt(dummyHistory, systemPrompt=systemPrompt)
try:
for _ in self.streamResponse(prompt):
pass
logger.debug("Shaders Metal et Prompt Cache APC pre-compiles (Warm-up complet).")
except (OSError, RuntimeError, ValueError) as e:
logger.debug("Note warm-up : %s", e)
[docs]
def streamResponse(self, formattedPrompt: str) -> Generator[str]:
"""
Genere le flux en streaming, avec decodage speculatif et reflexion.
Parameters
----------
formattedPrompt : str
Prompt produit par :meth:`formatPrompt`.
Yields
------
str
Fragments de texte au fur et a mesure.
"""
startTime = time.time()
tokenCounter = 0
firstTokenTime = None
lastResult = None
isThinking = self.enableThinking and self.reasoningBudget > 0
try:
if isThinking:
# PHASE 1 : reflexion
phase1Params: dict[str, Any] = {
"temperature": self.config.temperature,
"top_p": self.config.topP,
"repetition_penalty": self.config.repetitionPenalty,
"prefill_step_size": self.config.prefillStepSize,
"max_tokens": self.reasoningBudget,
"enable_thinking": True,
}
if self.apcManager is not None:
phase1Params["apc_manager"] = self.apcManager
if getattr(self.config, "kvBits", None) is not None:
phase1Params["kv_bits"] = self.config.kvBits
tok = self.processor.tokenizer
originalEosTokenIds = list(tok.stopping_criteria.eos_token_ids)
try:
thinkEndTokenId = tok.encode("</think>", add_special_tokens=False)[-1]
except (AttributeError, IndexError, ValueError):
thinkEndTokenId = 248069
tok.stopping_criteria.eos_token_ids = list(dict.fromkeys([*originalEosTokenIds, thinkEndTokenId]))
phase1Text = ""
try:
for responseChunk in stream_generate(
self.model,
self.processor,
formattedPrompt,
**phase1Params,
):
tokenCounter += 1
if firstTokenTime is None:
firstTokenTime = time.time()
self.lastTtftMs = (firstTokenTime - startTime) * 1000.0
phase1Text += responseChunk.text
yield responseChunk.text
finally:
tok.stopping_criteria.eos_token_ids = originalEosTokenIds
continuationPrompt = formattedPrompt + phase1Text
if not continuationPrompt.rstrip().endswith("</think>"):
closureTag = "\n</think>\n\n"
continuationPrompt = continuationPrompt.rstrip() + closureTag
else:
closureTag = "\n\n"
continuationPrompt += closureTag
yield closureTag
# PHASE 2 : bascule sur la parole, avec decodage speculatif
phase2Params: dict[str, Any] = {
"temperature": self.config.temperature,
"top_p": self.config.topP,
"repetition_penalty": self.config.repetitionPenalty,
"prefill_step_size": self.config.prefillStepSize,
"max_tokens": self.config.maxTokens,
}
if self.apcManager is not None:
phase2Params["apc_manager"] = self.apcManager
if getattr(self.config, "kvBits", None) is not None:
phase2Params["kv_bits"] = self.config.kvBits
if self.draftModel is not None:
phase2Params["draft_model"] = self.draftModel
phase2Params["draft_kind"] = self.draftKind
phase2Params["draft_block_size"] = getattr(self.config, "draftBlockSize", 2)
speechTokenCounter = 0
speechStartTime = time.time()
speechFirstTokenTime = None
for responseChunk in stream_generate(
self.model,
self.processor,
continuationPrompt,
**phase2Params,
):
speechTokenCounter += 1
tokenCounter += 1
lastResult = responseChunk
if speechFirstTokenTime is None:
speechFirstTokenTime = time.time()
yield responseChunk.text
speechDuration = time.time() - (speechFirstTokenTime or speechStartTime)
mlxTps = getattr(lastResult, "generation_tps", None)
if mlxTps is not None and mlxTps > 0:
self.lastTokensPerSec = round(float(mlxTps), 1)
elif speechDuration > 0 and speechTokenCounter > 0:
self.lastTokensPerSec = round(speechTokenCounter / speechDuration, 1)
else:
self.lastTokensPerSec = 0.0
else:
# Mode standard, reflexion coupee
generationParams: dict[str, Any] = {
"temperature": self.config.temperature,
"top_p": self.config.topP,
"repetition_penalty": self.config.repetitionPenalty,
"prefill_step_size": self.config.prefillStepSize,
"stop_words": ["<|im_end|>", "<end_of_turn>", "<eos>", "<reflection>"],
"max_tokens": self.config.maxTokens,
}
if self.apcManager is not None:
generationParams["apc_manager"] = self.apcManager
if getattr(self.config, "kvBits", None) is not None:
generationParams["kv_bits"] = self.config.kvBits
if self.draftModel is not None:
generationParams["draft_model"] = self.draftModel
generationParams["draft_kind"] = self.draftKind
generationParams["draft_block_size"] = getattr(self.config, "draftBlockSize", 2)
for responseChunk in stream_generate(
self.model,
self.processor,
formattedPrompt,
**generationParams,
):
tokenCounter += 1
lastResult = responseChunk
if firstTokenTime is None:
firstTokenTime = time.time()
self.lastTtftMs = (firstTokenTime - startTime) * 1000.0
yield responseChunk.text
duration = time.time() - (firstTokenTime or startTime)
mlxTps = getattr(lastResult, "generation_tps", None)
if mlxTps is not None and mlxTps > 0:
self.lastTokensPerSec = round(float(mlxTps), 1)
elif duration > 0 and tokenCounter > 0:
self.lastTokensPerSec = round(tokenCounter / duration, 1)
else:
self.lastTokensPerSec = 0.0
finally:
self.currentSpeedTokPerSec = self.lastTokensPerSec
self.lastTokenCount = getattr(lastResult, "generation_tokens", tokenCounter)
[docs]
def clearCache(self) -> None:
"""
Vide le pool de blocs de l'APC Manager.
"""
if self.apcManager is not None:
self.apcManager.clear()
logger.debug("Cache APC reinitialise avec succes.")
[docs]
def unload(self) -> None:
"""
Decharge le modele neural, le processeur et vide la memoire Metal allouee.
"""
if self.apcManager is not None:
self.apcManager.clear()
self.apcManager = None
self.model = None
self.processor = None
self.modelConfig = None
self.draftModel = None
self.promptCache = None
self.isLoaded = False