Method Predict

backend/python/mlx-distributed/backend.py:188–276

(self, request, context)

        return backend_pb2.Result(message="Model loaded successfully", success=True)

    async def Predict(self, request, context):
        prompt_cache = None
        cache_key = None

        try:
            import mlx.core as mx
            from mlx_lm import stream_generate
            from mlx_lm.sample_utils import make_logits_processors, make_sampler

            prompt_text = self._prepare_prompt(request)
            tokens = self._get_tokens_from_prompt(prompt_text)

            if self.coordinator:
                from coordinator import CMD_GENERATE
                self.coordinator.broadcast_command(CMD_GENERATE, len(tokens))
                self.coordinator.broadcast_tokens(tokens)

            max_tokens, sampler_params, logits_params, stop_words = self._build_generation_params(request)

            if self.coordinator:
                gen_params = self.coordinator.broadcast_generation_params(
                    max_tokens=max_tokens,
                    temperature=sampler_params.get('temp', 0.6),
                    top_p=sampler_params.get('top_p', 1.0),
                )
                max_tokens = gen_params["max_tokens"]

            sampler = make_sampler(**sampler_params)
            logits_processors = make_logits_processors(**logits_params) if logits_params else None

            # Use prompt cache in single-node mode
            gen_kwargs = {}
            if self.lru_cache is not None:
                from mlx_lm.models.cache import make_prompt_cache
                cache_key = list(tokens)
                prompt_cache, remaining_tokens = self.lru_cache.fetch_nearest_cache(
                    self.model_key, cache_key
                )
                if prompt_cache is None:
                    prompt_cache = make_prompt_cache(self.model, self.max_kv_size)
                    remaining_tokens = cache_key
                gen_kwargs['prompt_cache'] = prompt_cache
                tokens = remaining_tokens if remaining_tokens else cache_key

            generated = []
            last_response = None
            for response in stream_generate(
                self.model,
                self.tokenizer,
                prompt=tokens,
                max_tokens=max_tokens,
                sampler=sampler,
                logits_processors=logits_processors,
                **gen_kwargs,
            ):
                generated.append(response.text)
                last_response = response
                if cache_key is not None: