Skip to content

vllm_omni.worker_v2.omni_ar_model_runner

OmniARModelRunner — autoregressive stage runner on MR V2.

Extends OmniGPUModelRunner with:

  • OmniOutput post-processing in sample_tokens
  • Per-request pooler_output construction (hidden + multimodal slices)
  • Async D2H copy via OmniAsyncOutput for non-blocking output transfer
  • Cross-stage KV extraction before state cleanup

logger module-attribute 

logger = init_logger(__name__)

OmniARModelRunner

Bases: OmniGPUModelRunner

AR stage runner. Produces per-request hidden states + multimodal outputs.

kv_transfer_manager instance-attribute 

kv_transfer_manager: OmniKVTransferManager | None = None

execute_model 

execute_model(
    scheduler_output: SchedulerOutput,
    intermediate_tensors: Any | None = None,
    dummy_run: bool = False,
    skip_attn_for_dummy_run: bool = False,
    is_profile: bool = False,
    context_len: int = 0,
    valid_dummy_state_slots: bool = False,
) -> Any

sample_tokens 

sample_tokens(
    grammar_output: GrammarOutput | None,
) -> (
    OmniAsyncOutput
    | OmniModelRunnerOutput
    | ModelRunnerOutput
    | None
)

OmniAsyncOutput

Bases: AsyncModelRunnerOutput

Async D2H copy for Omni AR model outputs.

Mirrors upstream AsyncOutput but additionally handles pooler_output (hidden states + multimodal outputs) via non-blocking copies on the copy stream.

copy_event instance-attribute 

copy_event = (
    copy_event
    if copy_event is not None
    else torch.cuda.Event(blocking=True)
)

logprobs_tensors instance-attribute 

logprobs_tensors = None

model_runner_output instance-attribute 

model_runner_output = model_runner_output

num_nans instance-attribute 

num_nans: ndarray | None = None

num_sampled_tokens instance-attribute 

num_sampled_tokens = num_sampled_tokens

num_sampled_tokens_np instance-attribute 

num_sampled_tokens_np = _async_copy_to_np(
    num_sampled_tokens,
    copy_stream=copy_stream,
    pin_memory=pin_memory,
)

prompt_logprobs_dict instance-attribute 

prompt_logprobs_dict = {
    k: v.to_cpu_nonblocking()
    if v is not None
    else None
    for k, v in self.model_runner_output.prompt_logprobs_dict.items()
}

routed_experts instance-attribute 

routed_experts = routed_experts

routed_experts_cpu instance-attribute 

routed_experts_cpu = None

sampled_token_ids_np instance-attribute 

sampled_token_ids_np = _async_copy_to_np(
    sampler_output.sampled_token_ids,
    copy_stream=copy_stream,
    pin_memory=pin_memory,
)

sampler_output instance-attribute 

sampler_output = sampler_output

sampling_mask_tensors instance-attribute 

sampling_mask_tensors = None

get_output 

get_output() -> OmniModelRunnerOutput