LLM Rank Runtime#

class LLMRankRuntime#

Internal one-rank LLM execution runtime.

Manages inference pipeline for vanilla and speculative decoding modes (EAGLE, MTP, etc.). When constructed without a drafting config, operates as a pure vanilla decoding runtime with zero draft-model memory overhead. Coordinates base model, optional draft model, and multimodal processing (vision + audio).

Borrows the shared tokenizer and owns the engine, decoder, KV cache, multimodal runners, sampler, and request execution state for exactly one resolved runtime rank. Public SD/MD APIs should go through LLMInferenceRuntime and RuntimeCoordinator instead.

Note

This class is not thread-safe. Callers must externally serialize every method invocation and the object lifetime. handleRequest() defensively rejects accidental overlapping calls before mutating runtime state, but that gate does not authorize concurrent use of this object.

Public Types

using TokenBroadcastFn = std::function<bool(void *buffer, int32_t count, cudaStream_t stream)>#: Callback type for broadcasting GPU int32 buffers across parallel ranks. Rank 0 sends; peer ranks receive.

Public Functions

LLMRankRuntime( std::string const &engineDir, std::string const &multimodalEngineDir, std::unordered_map<std::string, std::string> const &loraWeightsMap, std::optional<SpecDecodeDraftingConfig> const &draftingConfig, cudaStream_t stream, ParallelMapping const &mapping, tokenizer::Tokenizer &tokenizer, ContextCacheConfig const &contextCacheConfig, std::string const &checkpointDir, std::string const &draftCheckpointDir )#: Construct one rank-local runtime from a fully-resolved parallel mapping. Preferred entry point — carries tensor/context/expert coordinates so future CP/EP support needs no further constructor changes.

LLMRankRuntime( ModelArtifacts &&artifacts, std::string const &engineDir, std::string const &multimodalEngineDir, std::unordered_map<std::string, std::string> const &loraWeightsMap, std::optional<SpecDecodeDraftingConfig> const &draftingConfig, cudaStream_t stream, ParallelMapping const &mapping, tokenizer::Tokenizer &tokenizer, ContextCacheConfig const &contextCacheConfig )#

~LLMRankRuntime()#: Destructor.

bool captureDecodingCUDAGraph(cudaStream_t stream)#

Capture CUDA graphs for decoding stages to optimize performance.

When draft model is present, captures graphs for draft proposal, draft accept token, base verification, and base vanilla decoding. Without draft model, captures only vanilla decoding graphs.

Note

If capture fails for any stage, the inference can proceed without CUDA graph capture, but at cost of performance degradation.

Parameters:: stream – CUDA stream
Throws:: std::runtime_error – if a tensor reshape operation fails
Returns:: True if all stage captures succeed, false otherwise

bool handleRequest( LLMGenerationRequest const &request, LLMGenerationResponse &response, cudaStream_t stream, bool outputThinkerEmbeddings = false, TokenBroadcastFn tokenBroadcast = nullptr, int32_t parallelRank = -1 )#

Handle generation request.

Note

Calls on the same runtime must be externally serialized. An accidental overlap with another handleRequest() is rejected before runtime or response state is mutated; this is not a general thread-safety guarantee.

Parameters:

request – Generation request with prompts and parameters
response – Output response with generated tokens and text
stream – CUDA stream

Throws:

std::runtime_error – if an LLM or CUDA operation fails

Returns:

True on success, false on failure

std::vector<int32_t> countPromptTokens( LLMGenerationRequest const &request ) const#: Return the input size for an explicit text token-count request.

bool genAndSaveSystemPromptKVCache( std::string const &prompt, std::string const &loraWeightsName, cudaStream_t stream )#

Generate and save system prompt KV cache (public API matching standard runtime signature)

Parameters:

prompt – The system prompt to generate the KVCache
loraWeightsName – The name of the LoRA weights
stream – The CUDA stream used for the generation

Throws:

std::runtime_error – if a CUDA operation fails

Returns:

True if the KVCache is generated and saved successfully, false otherwise

void setActionNoiseSeed(int32_t seed) noexcept#

Set the random seed used when initializing the action diffusion noise trajectory.

Parameters:: seed – Random seed value; has no effect if no action runner is loaded

void setVisualPrunerConfig(VisualPrunerConfig const &config)#: Enable visual-token pruning for supported VLM prefill execution.

inline metrics::LLMPrefillMetrics const &getPrefillMetrics( ) const noexcept#: Get LLM prefill stage metrics.

inline metrics::SpecDecodeGenerationMetrics const &getSpecDecodeGenerationMetrics( ) const noexcept#: Get speculative decoding generation stage metrics (only meaningful when draft model is present)

inline char const *getSpeculativeDecodingStrategyName( ) const noexcept#

inline metrics::LLMGenerationMetrics const &getGenerationMetrics( ) const noexcept#: Get vanilla generation stage metrics (only meaningful when no draft model / vanilla path)

std::optional<ContextCacheMetrics> getContextCacheMetrics( ) const noexcept#: Get context-cache metrics, or nullopt when the runtime cache is disabled.

inline metrics::MultimodalMetrics getMultimodalMetrics( ) const noexcept#: Get multimodal metrics (returns empty metrics if no multimodal runner)

inline rt::Tensor const &getEmbeddingTable() const#: Get the embedding table (for Talker streaming pipeline)

inline rt::Tensor const *getBaseModelHiddenStates( int32_t layerIdx ) const noexcept#

Get a base model hidden-states buffer for the requested layer index.

Buffers are owned by the runtime and reused across requests. Layer 0 corresponds to the post-multimodal input embeddings (backed up before the decode loop reshapes them); other layer indices correspond to engine-output hidden states (e.g. acceptHiddenLayer for the Qwen3-Omni Talker, or future MTP layers).

Lifetime contract:

Buffers are sized to {maxRuntimeBatchSize, maxSupportedInputLength, hiddenSize}.
Contents are cleared (overwritten) at the start of each handleRequest() call and remain valid until the next handleRequest() begins. The buffer is reshaped to {activeBatchSize, prefillLength, hiddenSize} for the most recent request — use getBaseModelPrefillLength() to query the valid prefill length.
The caller is responsible for consuming the data within that window.

Parameters:: layerIdx – Layer index. 0 = input embeddings (post-multimodal); other indices are model-specific (e.g. acceptHiddenLayer for Qwen3-Omni Talker).
Returns:: Pointer to the buffer, or nullptr if no buffer is registered for that layer.

inline int32_t getBaseModelPrefillLength() const noexcept#: Number of valid prefill tokens in the hidden-states buffers from the most recent handleRequest() call. Returns 0 if no hidden-states output was requested.

inline std::vector<std::vector<int32_t>> const &getBaseModelInputTokenIds( ) const noexcept#: Per-batch input token IDs from the most recent handleRequest() call. Cleared at the start of each handleRequest(); valid until the next one begins.

inline bool hasDraftModel() const noexcept#: Check if draft model is loaded and spec-decode is available.