easydel.inference.vsurge.__init__

class easydel.inference.vsurge.__init__.oDriver(engine: oEngine)

Bases: AbstractDriver

oDriver is responsible for managing the inference process for the oEngine. It handles request submission, input preparation, inference execution, and processing of model outputs. It utilizes background threads for concurrent processing of different stages of the inference pipeline.

compile()

Compiles the underlying engines.

This method is intended to perform any necessary compilation steps for the inference engines. Currently, it’s a placeholder.

property driver_name

Returns the name of the driver, derived from the engine’s model name.

get_total_concurrent_requests() → int

Gets the total number of concurrent requests the driver can handle.

This is determined by the maximum number of concurrent decode operations supported by the engine.

Returns

The maximum number of concurrent requests.

Return type

int

property num_used_slots

place_request_on_prefill_queue(request: ActiveRequest)

Places a new request onto the prefill queue for processing.

This method is used internally to add requests that require prefilling and subsequent generation.

Parameters

request (ActiveRequest) – The active request to place on the queue.

property processor: Any

Returns the processor/tokenizer associated with the engines.

Assumes all engines (prefill and decode) use the same processor. Raises an error if no engines are configured.

start()

Starts the driver and its background processing threads.

Threads for input preparation, inference, and summary processing are started if the driver is not already live.

stop()

Stops the driver and all background threads gracefully.

Signals the background threads to exit by putting None into their respective queues and then waits for them to join.

submit_request(request: Any)

Submits a new request to the driver’s processing pipeline.

The request is placed on the prefill queue for initial processing.

Parameters

request (tp.Any) – The request object to submit. Must be of type ActiveRequest.

Raises

TypeError – If the submitted request is not an instance of ActiveRequest.

class easydel.inference.vsurge.__init__.oEngine(model: Any, processor: Any, storage: PagedAttentionCache, manager: HBMPageManager, max_concurrent_decodes: int | None = None, max_concurrent_prefill: int | None = None, prefill_lengths: int | None = None, max_prefill_length: int | None = None, max_length: int | None = None, batch_size: int | None = None, seed: int = 894)

Bases: AbstractInferenceEngine

Optimized inference engine for EasyDeL models using Paged Attention.

This engine manages the inference process, including KV cache management with Paged Attention, scheduling, and execution of model forward passes for both prefill and decode steps.

property batch_size: int | None

Returns the configured batch size for the engine, if specified.

bulk_insert(prefix: Any, decode_state: Any, slots: list[int]) → Any

Efficiently inserts multiple prefill results into the decode state.

This method is optimized for integrating the results of a batch prefill operation into the decode state for multiple sequence slots.

Parameters

• prefix – The generation state from the bulk prefill step.

• decode_state – The current decode state.

• slots – A list of slot indices to insert into.

Returns

The updated decode state.

property colocated_cpus: Optional[list[jaxlib.xla_extension.Device]]

Returns CPU devices colocated with the engine’s accelerator devices.

This information can be useful for optimizing data transfers between host (CPU) and accelerator (GPU/TPU) memory. Currently returns None as the implementation is pending.

Returns

A list of colocated JAX CPU devices, or None if not implemented or available.

decode(graphstate: State[Key, VariableState[Any]], graphothers: State[Key, VariableState[Any]], state: Any, rngs: PRNGKey) → tuple[Any, Any]

Performs a single decode step for active sequences.

This involves generating the next token for each sequence based on the current state and KV cache.

Parameters

• graphstate – The graph state of the model.

• graphothers – Other graph components of the model.

• state – The current generation state.

• rngs – The PRNG key for sampling.

Returns

A tuple containing the updated generation state and the generated result tokens.

property eos_token_ids: list[int]

Returns a list of end-of-sequence token IDs from the processor and model config.

forward(graphstate: State[Key, VariableState[Any]], graphothers: State[Key, VariableState[Any]], state: ActiveSequenceBatch, iteration_plan: NextIterationPlan) → ModelOutputBatch

Performs a forward pass of the model.

This method executes the compiled continuous_forward function, processing the input state and iteration plan to produce model outputs and update the KV cache storage.

Parameters

• graphstate – The graph state of the model.

• graphothers – Other graph components of the model.

• state – The current active sequence batch state.

• iteration_plan – The plan for the current inference iteration.

Returns

The output batch from the model.

free_resource(slot: int) → bool

Frees resources associated with a specific inference slot. (Not Implemented)

Parameters

slot – The index of the slot to free.

Returns

Always returns False as it’s not implemented.

get_prefix_destination_sharding() → Any

Returns the shardings necessary to transfer KV cache data between engines.

This method is intended for scenarios involving multiple engines or devices where KV cache data needs to be moved.

Returns

The sharding specification for prefix destinations.

get_state_shardings(is_decode: bool = False)

Returns the sharding specifications for the engine’s state.

Parameters

is_decode – A boolean indicating if the sharding is for the decode state.

Returns

A tuple representing the sharding specification.

init_decode_state(*args, **kwargs) → ActiveSequenceBatch

Initializes the decode state for active sequences.

Parameters

• *args – Variable length argument list.

• **kwargs – Arbitrary keyword arguments.

Returns

An initialized ActiveSequenceBatch instance.

insert(prefix: Any, decode_state: Any, slot: int) → Any

Inserts or updates a generation state for a specific slot.

This is typically used to integrate the results of a prefill step into the ongoing decode process for a particular sequence slot.

Parameters

• prefix – The generation state from the prefill step.

• decode_state – The current decode state.

• slot – The slot index to insert into.

Returns

The updated decode state.

property max_concurrent_decodes

Maximum number of sequences decoded concurrently.

property max_concurrent_prefill

property max_length: int

Maximum total sequence length (prompt + generation).

This defines the size of the KV cache allocated.

property max_prefill_length: int

Maximum allowed length for the initial prompt (prefill phase).

Prompts longer than this will be truncated or handled according to the padding/truncation logic.

property pad_token_id

Returns the pad token ID from the processor.

prefill(graphstate: State[Key, VariableState[Any]], graphothers: State[Key, VariableState[Any]], tokens: Array, valids: Array, true_length: int, temperature: Array, top_p: Array, top_k: Array, rngs: PRNGKey) → tuple[Any, Any]

Performs the prefill step for a batch of prompts.

This involves processing the initial prompt tokens and populating the KV cache.

Parameters

• graphstate – The graph state of the model.

• graphothers – Other graph components of the model.

• tokens – The input tokens for the prompts.

• valids – A boolean array indicating valid tokens.

• true_length – The true length of the sequences.

• temperature – The temperature for sampling.

• top_p – The top-p value for sampling.

• top_k – The top-k value for sampling.

• rngs – The PRNG key for sampling.

Returns

A tuple containing the generation state after prefill and the result tokens.

property prefill_lengths: list[int]

Returns the configured list of max prefill length buckets for the engine.

property prng_key: PRNGKey

Provides a new PRNG key split from the internal state for sampling.

Each call to this property consumes the current key and returns a new, unique key, ensuring that subsequent sampling operations use different randomness.

Returns

A new JAX PRNGKey.

property samples_per_slot: int

Number of samples generated per inference slot.

This determines how many independent generation results are produced for each logical slot managed by the engine. It’s often 1, but could be higher for techniques like parallel sampling.

class easydel.inference.vsurge.__init__.vDriver(prefill_engines: Optional[Union[list[easydel.inference.vsurge.engines.vengine.engine.vEngine], vEngine]] = None, decode_engines: Optional[Union[list[easydel.inference.vsurge.engines.vengine.engine.vEngine], vEngine]] = None, interleaved_mode: bool = False, detokenizing_blocks: int = 8)

Bases: AbstractDriver

Drives the engines.

compile()

Compiles engines.

property driver_name

get_total_concurrent_requests() → int

Gets the total number of concurrent requests the driver can handle.

place_request_on_prefill_queue(request: ActiveRequest)

Used to place new requests for prefilling and generation.

property processor: Any

Returns the processor/tokenizer associated with the engines.

Assumes all engines (prefill and decode) use the same processor. Raises an error if no engines are configured.

start()

stop()

Stops the driver and all background threads.

submit_request(request: Any)

Submits a new request to the driver’s processing queue.

class easydel.inference.vsurge.__init__.vEngine(model: Any, processor: Any, max_concurrent_decodes: int | None = None, max_concurrent_prefill: int | None = None, prefill_lengths: int | None = None, max_prefill_length: int | None = None, max_length: int | None = None, batch_size: int | None = None, seed: int = 894)

Bases: AbstractInferenceEngine

Core inference engine for EasyDeL models using NNX graphs.

This engine manages the model state (split into graph definition, state, and other parameters) and provides JIT-compiled functions for the prefill and decode steps of autoregressive generation. It handles KV caching and sampling.

property batch_size: int | None

Returns the configured batch size for the engine, if specified.

bulk_insert(prefix: GenerationState, decode_state: GenerationState, slots: list[int]) → GenerationState

Efficiently inserts multiple prefill results into the decode state.

This function takes a GenerationState (prefix) typically resulting from a batch prefill operation and inserts its relevant components (logits, cache, index, tokens, valids, position IDs, generated tokens) into the main decode_state at multiple specified slots. This is useful for initializing the decode state after processing a batch of prompts simultaneously. Both input states’ caches are donated.

Parameters

• prefix – The GenerationState containing the results from a prefill operation (or similar initialization). Its cache is marked for donation.

• decode_state – The target GenerationState (e.g., the main decode loop state) to be updated. Its cache is marked for donation.

• slots – A list of integer indices indicating the slots within the decode_state’s batch dimension where the corresponding data from the prefix state should be inserted.

Returns

An updated GenerationState (decode_state) with the prefill results inserted at the specified slots.

property colocated_cpus: Optional[list[jaxlib.xla_extension.Device]]

Returns CPU devices colocated with the engine’s accelerator devices.

This information can be useful for optimizing data transfers between host (CPU) and accelerator (GPU/TPU) memory. Currently returns None as the implementation is pending.

Returns

A list of colocated JAX CPU devices, or None if not implemented or available.

decode(graphstate: State[Key, VariableState[Any]], graphothers: State[Key, VariableState[Any]], state: GenerationState, rngs: PRNGKey) → tuple[easydel.inference.vsurge.engines.vengine.utilities.GenerationState, easydel.inference.vsurge.engines._utils.ResultTokens]

Performs a single decode step in the autoregressive generation loop.

Takes the previous GenerationState, generates the next token using the model and KV cache, and updates the state. This function is JIT-compiled and allows the input state’s cache to be modified in-place (donated).

Parameters

• graphstate – The NNX GraphState (parameters) of the model.

• graphothers – Other NNX state variables of the model.

• state – The current GenerationState from the previous step. state.cache is marked for donation.

• rngs – A JAX PRNG key for sampling the next token.

Returns

• next_generation_state: The updated GenerationState for the next iteration.

• result: A ResultTokens object containing the newly generated token.

Return type

A tuple containing

property eos_token_ids: list[int]

A list of End-of-Sequence token IDs.

free_resource(slot: int) → bool

Frees resources associated with a specific inference slot. (Not Implemented)

Parameters

slot – The index of the slot to free.

Returns

Always returns False as it’s not implemented.

get_prefix_destination_sharding() → Any

Returns the shardings necessary to transfer KV cache data between engines.

Currently returns None, indicating default or no specific sharding.

get_state_shardings(is_decode: bool = False)

init_decode_state(*args, **kwargs) → GenerationState

Initializes the GenerationState for a new sequence

insert(prefix: GenerationState, decode_state: GenerationState, slot: int) → GenerationState

Inserts or updates a generation state, potentially for managing batches. (JIT-compiled)

This function seems designed to merge or update parts of the generation state, possibly inserting a ‘prefix’ state (e.g., from a completed prefill) into a larger batch state (‘decode_state’) at a specific ‘slot’. The exact mechanism for insertion isn’t fully clear from the current implementation, as it primarily focuses on broadcasting the prefix cache and returning the prefix state. Both input states’ caches are donated.

Parameters

• prefix – The GenerationState to potentially insert (e.g., from prefill). Its cache is marked for donation.

• decode_state – The target GenerationState to update (e.g., the main decode loop state). Its cache is marked for donation.

• slot – The index within the batch where the insertion/update should occur.

Returns

An updated GenerationState. In the current implementation, it returns the prefix state with its cache potentially broadcasted. Needs clarification on the intended merging logic with decode_state and slot.

property max_concurrent_decodes: int

Maximum number of sequences that can be decoded concurrently.

This determines the batch size used during the decode phase.

property max_length: int

Maximum total sequence length (prompt + generation).

This defines the size of the KV cache allocated.

property max_prefill_length: int

Maximum allowed length for the initial prompt (prefill phase).

Prompts longer than this will be truncated or handled according to the padding/truncation logic.

property pad_token_id

The ID of the padding token.

prefill(graphstate: State[Key, VariableState[Any]], graphothers: State[Key, VariableState[Any]], tokens: Array, valids: Array, true_length: int, temperature: Array, top_p: Array, rngs: PRNGKey) → tuple[easydel.inference.vsurge.engines.vengine.utilities.GenerationState, easydel.inference.vsurge.engines._utils.ResultTokens]

Performs the prefill step for initializing the generation process.

Processes the initial prompt tokens, initializes the KV cache, and generates the first token of the sequence. This function is JIT-compiled.

Parameters

• graphstate – The NNX GraphState (parameters) of the model.

• graphothers – Other NNX state variables of the model.

• tokens – The input prompt token IDs (batch_size, sequence_length).

• valids – A boolean array indicating valid token positions in the input (batch_size, sequence_length or batch_size, max_length).

• rngs – A JAX PRNG key for sampling the first token.

Returns

• generation_state: The initial GenerationState for the decode loop.

• result: A ResultTokens object containing the first generated token.

Return type

A tuple containing

property prefill_lengths: list[int]

Returns the configured list of max prefill length buckets for the engine.

property prng_key: PRNGKey

Provides a new PRNG key split from the internal state for sampling.

Each call to this property consumes the current key and returns a new, unique key, ensuring that subsequent sampling operations use different randomness.

Returns

A new JAX PRNGKey.

property samples_per_slot: int

Number of samples generated per inference slot.

This determines how many independent generation results are produced for each logical slot managed by the engine. It’s often 1, but could be higher for techniques like parallel sampling.

class easydel.inference.vsurge.__init__.vSurge(driver: Union[vDriver, oDriver], vsurge_name: str | None = None)

Bases: object

Orchestrates the interaction between client requests and the vDriver.

compile()

async complete(request: vSurgeRequest) → tp.AsyncGenerator[tp.List[ReturnSample]]

Initiates and streams the results of a text completion request.

Creates an ActiveRequest using the plain prompt from the vSurgeRequest, places it on the driver’s prefill queue, and then asynchronously iterates through the results provided by the ActiveRequest’s return_channel.

It handles both client-side and server-side tokenization scenarios, buffering and processing results appropriately before yielding them.

Parameters

request – The vSurgeRequest containing the prompt and generation parameters.

Yields

Processed generation results, similar to the decode method. The format depends on the tokenization mode.

Raises

RuntimeError – If the prefill queue is full when trying to place the request.

count_tokens(text_or_conversation: Union[str, list]) → int

Counts the number of tokens in a given string or conversation list.

Uses the underlying driver’s processor to tokenize the input and returns the count of tokens.

Parameters

text_or_conversation – Either a single string or a list of message dictionaries (like OpenAI chat format).

Returns

The total number of tokens in the input.

Raises

ValueError – If the input type is invalid or tokenization fails.

classmethod create_odriver(model: Any, processor: Any, storage: Optional[PagedAttentionCache] = None, manager: Optional[HBMPageManager] = None, page_size: int = 128, hbm_utilization: float = 0.6, max_concurrent_prefill: int | None = None, max_concurrent_decodes: int | None = None, prefill_lengths: int | None = None, max_prefill_length: int | None = None, max_length: int | None = None, seed: int = 894, vsurge_name: str | None = None) → vSurge

classmethod create_vdriver(model: Any, processor: Any, max_concurrent_decodes: int | None = None, prefill_lengths: int | None = None, max_prefill_length: int | None = None, max_length: int | None = None, seed: int = 894, vsurge_name: str | None = None) → vSurge

Creates a new instance of vSurge with configured vDriver and vEngines.

This class method provides a convenient way to instantiate the vSurge by setting up the necessary prefill and decode engines with the provided model, processor, and configuration parameters.

Parameters

• model – The EasyDeLBaseModule instance representing the model.

• processor – The tokenizer/processor instance.

• max_concurrent_decodes – Maximum number of concurrent decode requests the decode engine can handle.

• prefill_lengths – A list of prefill lengths to compile for the prefill engine.

• max_prefill_length – The maximum prefill length for the prefill engine.

• max_length – The maximum total sequence length for both engines.

• seed – The random seed for reproducibility.

• vsurge_name – An optional name for the vsurge.

Returns

A new instance of vSurge.

property driver

Provides access to the underlying vDriver instance.

async generate(prompts: tp.Union[str, tp.Sequence[str]], sampling_params: tp.Optional[tp.Union[SamplingParams, tp.Sequence[SamplingParams]]] = None, stream: bool = False) → tp.Union[tp.List[ReturnSample], tp.AsyncGenerator[tp.List[ReturnSample]]]

Generates text completions concurrently for the given prompts.

Parameters

• prompts – A single prompt string or a list of prompt strings.

• sampling_params – A single SamplingParams object or a list of SamplingParams objects. If None, default SamplingParams will be used. If a single SamplingParams object is provided with multiple prompts, it will be applied to all prompts. If a list is provided, it must have the same length as the prompts list.

• stream – If True, yields results (List[ReturnSample]) from any request as they become available. The list corresponds to one generation step from one request. If False, waits for all requests to complete and returns a list containing one aggregated ReturnSample per prompt.

Returns

An async generator yielding lists of ReturnSample as

steps complete across concurrent requests.

If stream is False: A list of aggregated ReturnSample objects, one for

each input prompt, after all requests have finished.

Return type

If stream is True

Raises

• ValueError – If the lengths of prompts and sampling_params lists mismatch.

• RuntimeError – If the underlying driver’s queue is full.

process_client_side_tokenization_response(response: list[easydel.inference.vsurge.utils.ReturnSample])

Processes responses when tokenization is handled client-side.

In this case, the response items (ReturnSample) are typically yielded directly without further server-side processing like detokenization or buffering.

Parameters

response – A list of ReturnSample objects from a single generation step.

Returns

The input list of ReturnSample objects, unchanged.

process_server_side_tokenization_response(response: list[easydel.inference.vsurge.utils.ReturnSample], buffered_response_list: list[list[easydel.inference.vsurge.utils.ReturnSample]]) → list[easydel.inference.vsurge.utils.ReturnSample]

Processes responses when tokenization/detokenization is server-side.

Combines the text and token IDs from the current response and any buffered previous responses for each sample. It then uses the metrics (TPS, generated token count) from the latest response in the sequence for the final output.

Parameters

• response – The list of ReturnSample objects from the current step.

• buffered_response_list – A list containing lists of ReturnSample objects from previous steps that were buffered.

Returns

A list of tuples, where each tuple represents a completed sample and contains: (decoded_string, all_token_ids, latest_tps, latest_num_generated_tokens).

property processor: Any

Returns the processor/tokenizer associated with the underlying driver.

should_buffer_response(response: list[easydel.inference.vsurge.utils.ReturnSample]) → bool

Determines if a response needs buffering for server-side detokenization.

Buffering is needed if any sample in the response ends with a byte token (e.g., “<0xAB>”), as this indicates an incomplete multi-byte character that requires subsequent tokens for proper decoding.

Parameters

response – A list of ReturnSample objects from a single generation step.

Returns

True if buffering is required, False otherwise.

start()

stop()

property vsurge_name

class easydel.inference.vsurge.__init__.vSurgeApiServer(vsurge_map: Union[Dict[str, vSurge], vSurge] = None, max_workers: int = 10, oai_like_processor: bool = True)

Bases: object

FastAPI server for serving vEngine instances.

This server provides endpoints mimicking the OpenAI API structure for chat completions, liveness/readiness checks, token counting, and listing available models. It handles both streaming and non-streaming requests asynchronously using a thread pool.

async available_inference()

Lists available models (GET /v1/models).

async chat_completions(request: ChatCompletionRequest)

Handles chat completion requests (POST /v1/chat/completions).

Validates the request, retrieves the appropriate vEngine model, tokenizes the input, and delegates to streaming or non-streaming handlers.

Parameters

request (ChatCompletionRequest) – The incoming request data.

Returns

The generated response, either

a complete JSON object or a streaming event-stream.

Return type

Union[JSONResponse, StreamingResponse]

async completions(request: CompletionRequest)

Handles completion requests (POST /v1/completions).

Processes the prompt for completion and returns generated text.

Parameters

request (CompletionRequest) – The incoming request data.

Returns

The generated response.

Return type

Union[JSONResponse, StreamingResponse]

async count_tokens(request: CountTokenRequest)

Token counting endpoint (POST /v1/count_tokens).

fire(host='0.0.0.0', port=11556, metrics_port: Optional[int] = None, log_level='info', ssl_keyfile: Optional[str] = None, ssl_certfile: Optional[str] = None)

Starts the uvicorn server to run the FastAPI application.

Parameters

• host (str) – The host address to bind to. Defaults to “0.0.0.0”.

• port (int) – The port to listen on. Defaults to 11556.

• metrics_port (tp.Optional[int]) – The port for the Prometheus metrics server. If None, defaults to port + 1. Set to -1 to disable.

• log_level (str) – The logging level for uvicorn. Defaults to “info”.

• ssl_keyfile (tp.Optional[str]) – Path to the SSL key file for HTTPS.

• ssl_certfile (tp.Optional[str]) – Path to the SSL certificate file for HTTPS.

async liveness()

Liveness check endpoint (GET /liveness).

async readiness()

Readiness check endpoint (GET /readiness).

class easydel.inference.vsurge.__init__.vSurgeRequest(prompt: str, max_tokens: int, top_p: float = 1.0, top_k: int = 0, min_p: float = 0.0, temperature: float = 0.7, presence_penalty: float = 0.0, frequency_penalty: float = 0.0, repetition_penalty: float = 1.0, metadata: easydel.inference.vsurge.vsurge.vSurgeMetadata | None = None, is_client_side_tokenization: bool = False)

Bases: object

Represents a request specifically for text completion.

frequency_penalty: float = 0.0

classmethod from_sampling_params(prompt: str, sampling_params: SamplingParams)

is_client_side_tokenization: bool = False

max_tokens: int

metadata: easydel.inference.vsurge.vsurge.vSurgeMetadata | None = None

min_p: float = 0.0

presence_penalty: float = 0.0

prompt: str

repetition_penalty: float = 1.0

temperature: float = 0.7

top_k: int = 0

top_p: float = 1.0