Method PredictStream

backend/python/mlx/backend.py:291–392

Generates text based on the given prompt and sampling parameters, and streams the results using MLX. Uses thread-safe LRU prompt cache for efficient prefix reuse across requests. Args: request: The predict stream request. context: The gRPC context.

(self, request, context)

            return backend_pb2.Result(success=False, message=str(e))

    async def PredictStream(self, request, context):
        """
        Generates text based on the given prompt and sampling parameters, and streams the results using MLX.

        Uses thread-safe LRU prompt cache for efficient prefix reuse across requests.

        Args:
            request: The predict stream request.
            context: The gRPC context.

        Yields:
            backend_pb2.Reply: Streaming predict results.
        """
        prompt_cache = None
        cache_key = None

        try:
            # Prepare the prompt and tokenize for cache key
            prompt_text = self._prepare_prompt(request)
            cache_key = self._get_tokens_from_prompt(prompt_text)

            # Fetch nearest cache (exact, shorter prefix, or create new)
            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

            # Build generation parameters using request attributes and options
            max_tokens, sampler_params, logits_params, stop_words = self._build_generation_params(
                request, default_max_tokens=512
            )

            print(
                f"Streaming text with MLX - max_tokens: {max_tokens}, "
                f"cache_hit: {len(remaining_tokens) < len(cache_key)}",
                file=sys.stderr,
            )

            # Create sampler and optional logits processors (penalties)
            sampler = make_sampler(**sampler_params)
            logits_processors = make_logits_processors(**logits_params) if logits_params else None

            accumulated = []
            last_response = None
            for response in stream_generate(
                self.model,
                self.tokenizer,
                prompt=remaining_tokens if remaining_tokens else cache_key,
                max_tokens=max_tokens,
                sampler=sampler,
                logits_processors=logits_processors,
                prompt_cache=prompt_cache,
            ):
                cache_key.append(response.token)
                accumulated.append(response.text)
                last_response = response