This document is relevant for: Trn2, Trn3

Async Scheduling and Async Execution

Table of Contents

• Overview

• When Does Async Scheduling Help Most?

• Custom Async Execution Integration with Async Scheduling

• Normal Flow (Synchronous Baseline)

• Async Flow (Steady-State)

• What Does “Breaking the Async Flow” Mean?

• How Do We Break the Async Flow Safely?

• Development and Debugging

Overview

Async scheduling overlaps CPU work (scheduling, input preparation, output processing) with Neuron device execution (NEFF forward pass) to reduce per-step latency. In synchronous mode, the CPU idles while the device runs the forward pass, and the device idles while the CPU prepares the next step. Async scheduling eliminates this idle time by:

1. Using a batch queue (depth 2) in the engine core to schedule one step ahead, so the worker always has work ready.

2. Feeding the previous step’s sampled token future directly as input_ids into the next forward pass (device-to-device, no CPU roundtrip), using the Neuron async runtime to enqueue NEFF executions ahead of time.

3. Materializing tokens to CPU on a separate async output thread, overlapped with the engine core scheduling the next step.

When the batch composition changes (requests finish, new requests join), or when the previous sampled-token future cannot be reused for the next decode input shape, the async flow must be “broken” — the pending device future is materialized on the critical path, batch state is updated, and CPU-built input_ids is used for that step. If the request set is unchanged but the batch order changed, the future can be remapped into the new order; because the remap clones/forces the future on the worker main thread, that step is accounted as a sync fallback. The async flow automatically reinstates on the next steady-state decode step.

Async scheduling is on by default and requires on-device sampling. To disable, use --no-async-scheduling.

When Does Async Scheduling Help Most?

All workloads benefit from async scheduling, but these scenarios see the largest gains:

• Smaller models (faster NEFF execution): The CPU overhead (scheduling, input preparation, output processing) is roughly constant regardless of model size. For smaller models where the NEFF forward pass is fast, the CPU overhead becomes a larger fraction of the total step time. Async scheduling hides this overhead by overlapping it with device execution, so the relative speedup is greater.

• Longer average response length: Async scheduling only engages during steady-state decode (consecutive decode steps with no batch composition change). Longer responses mean more consecutive decode steps, so the async path is active for a larger fraction of the total generation. Short responses that finish quickly cause frequent batch composition changes, forcing sync fallbacks that reduce the benefit.

Custom Async Execution Integration with Async Scheduling

The upstream vLLM GPU implementation uses CUDA streams for async execution — the GPU runtime provides stream semantics where operations on the same stream are ordered and the CPU can enqueue work without blocking. The CPU simply enqueues kernels and continues.

The Neuron runtime does not currently provide torch.cuda.Stream-style semantics. Once native PyTorch support for Neuron is available, it will provide stream semantics and this custom implementation can be simplified. Until then, the runtime offers an async execution mode where:

• A NEFF (compiled model) can be enqueued for execution with device tensor futures as inputs. The runtime guarantees the NEFF will wait for the input futures to resolve before executing.

• Output tensors are also futures — accessing them (e.g., .to("cpu")) blocks until the NEFF completes.

This means we must explicitly manage the future lifecycle:

1. Hold futures across steps: Store the sampled token future in async_execution_buffer and pass it as input_ids to the next NEFF execution, avoiding the device→host→device roundtrip.

2. Intentionally break the async flow: When the batch composition changes, or when a shape/dtype mismatch forces CPU-built inputs, we must materialize the future on the critical path (blocking .to("cpu")) to update CPU-side state before proceeding. The GPU implementation handles this implicitly via stream synchronization; we must do it explicitly. When the request set is unchanged but the request order changes, we can remap the future on-device, but the remap still forces the future and is counted as a sync fallback step.

3. Thread-safe materialization: Both the async output thread and the worker main thread may call get_output() on the same future. A per-object lock prevents double-materialization.

Normal Flow (Synchronous Baseline)

In synchronous mode (--no-async-scheduling), each inference step runs as a strictly sequential pipeline. The engine core orchestrates the flow:

1. Scheduler selects which requests to run and produces a SchedulerOutput.

2. Engine core calls worker.execute_model(scheduler_output).

3. Inside the worker, NeuronModelRunner.execute_model() does the heavy lifting:

   a. _update_states() — applies the scheduler’s decisions (add/remove requests, update token counts, manage sequence IDs).

   b. _prepare_model_input() — builds input_ids, positions, logits_indices, and attn_metadata from the current InputBatch state.

   c. _execute_model_forward() — moves tensors to the Neuron device, runs model(**kwargs), and blocks on model_output_tensor.to("cpu") to copy the result back to host memory.

   d. Stores the result in execute_model_state and returns None.

4. Engine core sees None and calls worker.sample_tokens(grammar_output).

5. NeuronModelRunner.sample_tokens() unpacks the stored state, runs the sampler (_sample()), updates batch state with the sampled tokens (_update_batch_state_with_samples()), and returns a ModelRunnerOutput.

6. Engine core passes the output to scheduler.update_from_output() which processes finished requests, frees blocks, and prepares for the next step.

The critical property of this flow is that every tensor transfer is blocking: the CPU waits for the device-to-host copy of logits/sampled tokens to complete before proceeding. This means the CPU is idle during the entire model forward pass, and the device is idle during scheduling and input preparation.

Step N                                          Step N+1
─────────────────────────────────────────────    ──────────────────────

┌─────────────────────────────────────────────────────────────────────┐
│ Engine Core                                                        │
│                                                                    │
│  schedule()  ──►  execute_model()  ──►  sample_tokens()  ──►  ...  │
│      │                  │                     │                     │
└──────┼──────────────────┼─────────────────────┼─────────────────────┘
       │                  │                     │
       ▼                  ▼                     ▼
┌─────────────────────────────────────────────────────────────────────┐
│ NeuronModelRunner                                                  │
│                                                                    │
│  ┌──────────────┐  ┌──────────────────────────────┐  ┌──────────┐  │
│  │ _update_     │  │ _execute_model_forward()     │  │ _sample()│  │
│  │  states()    │  │                              │  │          │  │
│  │              │  │  ┌────────────────────────┐  │  │          │  │
│  │ _prepare_    │  │  │ model(**kwargs)        │  │  │          │  │
│  │  model_      │  │  │   (Neuron device)      │  │  │          │  │
│  │  input()     │  │  └────────────────────────┘  │  │          │  │
│  │              │  │  output.to("cpu")  ◄─ BLOCK  │  │          │  │
│  └──────────────┘  └──────────────────────────────┘  └──────────┘  │
│                                                                    │
└────────────────────────────────────────────────────────────────────┘

Timeline (single step):

CPU:   ┃ schedule │ prepare_input │░░░░░ IDLE ░░░░░│ sample │ update ┃
Device:┃░░░░░░░░░░░░░░░░░░░░░░░░░│ model forward  │░░░░░░░░░░░░░░░░┃
                                   ▲               ▲
                                   │               │
                               .to(device)    .to("cpu") ← blocking

The ░░░░░ IDLE regions are the latency cost of synchronous execution: the CPU stalls while the device runs the forward pass, and the device sits idle while the CPU schedules and prepares the next step. Async scheduling exists to overlap these two phases.

Async Flow (Steady-State)

Steady-state means consecutive decode steps where the batch composition does not change — no requests finish, no new requests join. When the batch changes, we leave steady-state and enter the “breaking” path described in later sections.

Key Idea

In synchronous mode the CPU blocks on .to("cpu") to copy sampled tokens back from the device before the next step can begin. Async scheduling eliminates that blocking transfer from the critical path by:

1. Keeping the sampled token tensor on the device and feeding it directly as input_ids into the next forward pass (device-to-device, no CPU roundtrip). This is made possible by the Neuron async runtime: we can pass the future of a device tensor to a NEFF and enqueue the NEFF execution ahead of time. The runtime guarantees that the NEFF will wait for the tensor to be ready before executing.

2. Materializing the tokens to CPU on a separate thread — the multiproc executor’s async_output_busy_loop thread calls AsyncNeuronModelRunnerOutput.get_output(), which performs the .to("cpu") transfer and updates batch state. This runs concurrently with the engine core’s main thread scheduling the next step.

So the CPU does eventually see every token — but the materialization happens off the critical path, overlapped with scheduling and input preparation for the next step. This also guarantees that output tokens are streamed back to the client: the async output thread materializes and enqueues each step’s tokens as soon as they’re ready, so the client receives tokens incrementally rather than waiting for the full sequence to complete.

Step-by-Step Flow

Consider two consecutive decode steps, Step N and Step N+1, with the same batch of requests.

Step N (produces futures):

1. execute_model() calls _update_states() and _prepare_model_input() as usual. On the very first async step, input_ids comes from CPU (there are no futures yet).

2. _execute_model_forward() runs model(**kwargs) on the Neuron device. The model returns sampled_token_ids as a device tensor future.

3. Unlike sync mode, the .to("cpu") call is skipped — the output stays on device as a future.

4. sample_tokens() takes a snapshot of req_ids via _bookkeeping_sync() (so the next step’s _update_states() can freely mutate InputBatch without corrupting the current step’s output mapping), runs _sample(), and calls _generate_model_runner_output(). Note: there is no race condition here — execute_model() and sample_tokens() are dispatched sequentially on the same worker process, so the next step’s _update_states() is guaranteed not to run before the current step’s _bookkeeping_sync() completes.

5. _generate_model_runner_output() stores the device-side sampled token future into async_execution_buffer["futures_sampled_token_ids"] and wraps the output in an AsyncNeuronModelRunnerOutput.

6. The AsyncNeuronModelRunnerOutput is returned to the engine core. In the multiproc executor, handle_output() puts it on the async_output_queue. A separate async_output_busy_loop thread picks it up and calls get_output(), which blocks on .to("cpu") and writes the sampled tokens into token_ids_cpu and req_state.output_token_ids via _update_batch_state_with_samples(). This materialization runs concurrently with the engine core scheduling Step N+1.

7. With async scheduling, the engine core uses step_with_batch_queue() with a batch queue of size 2 (max_concurrent_batches). The engine core is always one step ahead: by the time Step N’s get_output() materializes, Step N+1 has already been scheduled and enqueued. Once Step N’s output is processed, the engine core schedules Step N+2 and then blocks waiting for Step N+1’s output to materialize.

Step N+1 (consumes futures):

1. execute_model() runs _update_states() and _prepare_model_input(). _prepare_model_input() builds input_ids from token_ids_cpu as usual — but these are stale (they don’t contain Step N’s sampled tokens, which are still on device).

2. execute_model() checks whether the async swap is eligible: the buffer is populated, this is a decode step, and the batch composition did not change.

3. It reads async_execution_buffer["futures_sampled_token_ids"] — the device tensor future from Step N — and checks that its shape and dtype match the current padded decode input.

4. Since the batch hasn’t changed, the shape and dtype match. If the request order is unchanged, the stale CPU-built input_ids is replaced with the device-side future.

5. If the request set is the same but condense() changed slot order, the future is remapped into the current request order before it is used. This produces correct input_ids, but the clone/remap forces the future on the worker main thread, so this step is counted as sync fallback rather than steady-state async.

6. The future (still on device) is passed directly to model(**kwargs) — no device→host→device roundtrip.

7. The cycle repeats: new sampled token futures stay on device, get stored in the buffer, and feed into Step N+2.

Where the Overlap Happens

While the device is executing Step N’s forward pass, the CPU is free to run the scheduler for Step N+1, call _update_states(), and prepare inputs. The CPU-built input_ids will be thrown away in favor of the device future, but the scheduling and state bookkeeping work is done in parallel with the device computation.

Steady-state async (batch unchanged):

              Step N                              Step N+1
              ───────────────────────────         ───────────────────────────

Worker        ┃ execute_model(N) │ sample(N) │    │ execute_model(N+1) │ sample(N+1) │
main          ┃ (input_ids =     │           │    │ (input_ids =        │             │
thread:       ┃  future(N-1))    │           │    │  future(N))         │             │
              ┃       │          │           │    │       │              │             │
Device:       ┃ ══ forward(N) ═══════════════╪════╪══ forward(N+1) ══════════════════ ...
              ┃    waits for ◄───────┐       │    │    waits for ◄───────┐
              ┃    future(N-1)       │       │    │    future(N)         │
              ┃                      │       │    │                      │
              ┃                      │       │    │                      │
Worker        ┃    get_output(N-1)   │       │    │    get_output(N)     │
async         ┃    waits for ◄───────┘       │    │    waits for ◄───────┘
thread:       ┃    future(N-1)               │    │    future(N)
              ┃    .to("cpu") + update_batch │    │    .to("cpu") + update_batch
              ┃                              │    │
Engine        ┃    ◄── block ──► update(N-1) │    │    ◄── block ──► update(N)
core:         ┃                  schedule(N+1)    │                   schedule(N+2)

Both the device (NEFF execution) and the async thread (.to("cpu"))
wait on the same future — the sampled token future from the previous step.

The async_execution_buffer

This dict is the bridge between consecutive async steps. In steady-state, only futures_sampled_token_ids is consumed:

Key	Value	Purpose
futures_sampled_token_ids	Device tensor future [num_reqs]	Sampled token future from the previous step, fed as input_ids into the next forward pass
prev_req_ids_ordered	Ordered list[str] of request IDs	Request order that produced the future. Its set is compared against the next step’s scheduled requests to detect composition changes; its ordering is used to remap same-set reordered futures.

The buffer also stores async_output (AsyncNeuronModelRunnerOutput), but this is only accessed when breaking the async flow — see How Do We Break the Async Flow Safely?.

What Does “Breaking the Async Flow” Mean?

“Breaking the async flow” means forcing the materialization of a previous step’s device tensor future on the worker main thread (the critical path), rather than letting it be consumed only by the async output thread and the next NEFF execution.

Why Does It Happen?

Consider the steady-state async flow: step N is a decode step, and step N+1 has already been scheduled ahead of time (batch queue depth = 2). During step N, one or more requests complete (EOS or max output tokens). We only discover this after step N’s output is materialized — when the engine core’s scheduler.update_from_output(N) checks stop conditions and marks the request as finished. At that point, step N+1 is already in flight on the worker.

After processing step N’s output, the engine core schedules step N+2. This time, the SchedulerOutput reflects the updated state: the finished request is gone, and the scheduler may have added a new prefill, filled the decode batch to max batch size, or simply continued decode with fewer requests. The batch composition has changed.

Why Can’t We Stay Async?

In steady-state, the worker discards the CPU-built input_ids and replaces them with the device future from the previous step. For step N+2, we cannot always do that:

• If the request set changed, the scheduler’s input_ids must be respected. Step N+2’s input_ids come from the scheduler’s updated view of the world (with the finished request removed, possibly a new prefill added, etc.). We must use them as-is, not replace them with a stale device future.

• If the request set is unchanged but the padded input shape or dtype differs, the future cannot be passed to the next NEFF input. We must materialize the pending output, rebuild CPU inputs from refreshed state, and use those rebuilt inputs.

• If the request set is unchanged but the request order changed, the future has the right shape but wrong slot ownership. We remap it before use; this keeps inputs correct but forces the future, so it is not counted as a steady-state async step.

For composition changes and shape/dtype mismatches, “breaking” means: materialize step N+1’s future on the main thread (get_output()), update all batch state, rebuild input_ids from the now-correct CPU state, and proceed synchronously for this step. For same-set order changes, the worker remaps the future into the current order; this does not rebuild CPU inputs, but it still forces the future on the main thread and is accounted as sync fallback. The async flow can resume once the batch stabilizes again.

How Do We Break the Async Flow Safely?

Breaking the async flow happens on the worker main thread, on the critical path of step N+2’s execution. After breaking, the async flow is automatically reinstated at the end of the same step.

Note: the scheduler always schedules one step ahead — it is unaware of the async/sync distinction. It is the worker (specifically NeuronModelRunner) that detects when the async flow must be broken and is responsible for materializing futures, adjusting batch state, and reinstating the async pipeline.

Pre-_update_states Guard in execute_model()

This is the first check, before any state mutation happens for step N+2.

# In execute_model(), before _update_states():
if self.use_async_scheduling:
    if not self.async_execution_buffer:
        self._batch_composition_changed = False
    else:
        curr_req_ids = frozenset(scheduler_output.num_scheduled_tokens.keys())
        prev_req_ids_ordered = self.async_execution_buffer.get(
            "prev_req_ids_ordered"
        )
        prev_req_ids = (
            frozenset(prev_req_ids_ordered)
            if prev_req_ids_ordered is not None
            else None
        )
        composition_changed = (
            prev_req_ids is None or curr_req_ids != prev_req_ids
        )
        if composition_changed:
            self._materialize_pending_async_output("batch composition changed")
            self._batch_composition_changed = True
        else:
            self._batch_composition_changed = False

What it does:

• Detects batch composition changes by comparing the current step’s scheduled request IDs (curr_req_ids) against the previous step’s (prev_req_ids, derived from prev_req_ids_ordered, which is stored in the buffer when the future was produced). If the sets differ for ANY reason — requests finished, new requests joined, requests became unscheduled — it’s a composition change. This single set comparison is simpler and more robust than checking individual signals (finished_req_ids, scheduled_new_reqs, etc.) because it catches all cases including the max_tokens pipelining gap where a request disappears from num_scheduled_tokens one step before appearing in finished_req_ids.

• If composition changed, it calls get_output() on the previous step’s AsyncNeuronModelRunnerOutput, which:

  • Blocks until the device tensor future is ready

  • Copies sampled tokens to CPU via .to("cpu")

  • Writes them into token_ids_cpu and req_state.output_token_ids via _update_batch_state_with_samples()

• Sets _batch_composition_changed = True so that the swap logic later knows not to replace input_ids with the stale device future.

• This ensures _update_states() (called immediately after) sees correct token data when it calls input_batch.add_request() for resumed/new requests — add_request copies from req_state.output_token_ids into token_ids_cpu, so the data must be up to date.

• After materialization, _update_states() and _prepare_model_input() run against correct state, so input_ids is built correctly from CPU and does not need rebuilding.

The prev_req_ids_ordered list is stored in _generate_model_runner_output() as list(req_ids_output_copy) — the exact order of request IDs that produced the future. The same list supports both composition detection (via set comparison) and same-set reorder remapping (via index lookup).

Async/Sync Swap Decision in execute_model()

After _update_states() and _prepare_model_input() have run, the swap logic decides whether to use the device future or the CPU-built input_ids:

# In execute_model(), after _prepare_model_input():
if self.use_async_scheduling:
    if (
        self.async_execution_buffer
        and self._is_decode()
        and not self._batch_composition_changed
    ):
        future = self.async_execution_buffer["futures_sampled_token_ids"]
        if future.shape[0] == input_ids.shape[0] and future.dtype == input_ids.dtype:
            input_ids, broke_async = self._maybe_remap_async_future(
                future, input_ids
            )
            if broke_async:
                self._sync_fallback_steps += 1
            else:
                self._async_steps += 1
        else:
            # Shape/dtype mismatch: materialize pending output, rebuild inputs,
            # and use refreshed CPU-built input_ids.
            if self._materialize_pending_async_output(
                "async future shape/dtype mismatch"
            ):
                input_ids, positions, logits_indices, attn_metadata, spec_decode_metadata = (
                    self._prepare_model_input(scheduler_output)
                )
            self._sync_fallback_steps += 1
    else:
        # Broken: use CPU-built input_ids as-is.
        self._sync_fallback_steps += 1

The async path is used only when all four conditions are met:

• async_execution_buffer is non-empty (not the first iteration)

• _is_decode() returns True (all requests are in decode phase — note: this must be checked after _update_states() so it reflects the current step’s batch, not the previous one)

• _batch_composition_changed is False (no requests finished/added/resumed)

• The future shape and dtype match input_ids — _prepare_model_input() applies decode batch padding (e.g., 1 request padded to bucket size 4), and the future must match this padded shape. On the first decode step after a prefill, the future has shape [num_reqs] (from the prefill output) while input_ids has shape [padded_bucket_size], causing a mismatch and a one-step sync fallback. The fallback first materializes the pending async output, then rebuilds model inputs from refreshed CPU state. This prevents stale CPU input_ids from being used after the latest sampled token still lived only in a device future. From the second decode step onward, the model output (and thus the future) already has the padded shape, so the shapes match and async resumes.

• If the request set is unchanged but the current batch order differs from prev_req_ids_ordered, _maybe_remap_async_future() remaps real request slots into the current order and leaves padding slots unchanged. This keeps token ownership correct after condense() reorder. Because remap uses clone()/gather and forces the future on the main thread, the step is counted as _sync_fallback_steps, not _async_steps.

When any condition fails, input_ids from _prepare_model_input() is used as-is — it was already built from correct CPU state (either because the pre-_update_states guard materialized, because the shape/dtype mismatch path materialized and rebuilt it, or because this is a prefill/first iteration where no stale futures exist).

Reinstatement

The async flow is automatically reinstated at the end of every step, regardless of whether it was broken:

1. _execute_model_forward() always skips .to("cpu") when use_async_scheduling is True — the model output stays on device as a future.

2. _generate_model_runner_output() always stores sampler_output.sampled_token_ids (the new device future) into async_execution_buffer["futures_sampled_token_ids"].

3. sample_tokens() always wraps the output in AsyncNeuronModelRunnerOutput and stores it in async_execution_buffer["async_output"].

So even after a sync fallback step, the buffer is repopulated with fresh futures. The next step will use the async path if it’s a decode step with no composition change and matching shapes, and the async flow resumes.

Development and Debugging

get_output() Call Patterns Across TP Ranks

Only the output rank (TP0) has its AsyncNeuronModelRunnerOutput processed by the async output thread. This is because the multiproc executor’s worker_busy_loop only calls handle_output(output) when self.rank == output_rank — non-output ranks discard their output.

This means:

• TP0: get_output() is called on every step by the async output thread (via enqueue_output). Additionally, the worker main thread calls it when breaking the async flow. With the threading lock, only one thread materializes; the other hits the “already actualized” else branch.

• TP1–TP7: get_output() is only called by the worker main thread when breaking the async flow (the async output thread never sees their objects). In steady-state, their AsyncNeuronModelRunnerOutput objects are created and stored in the buffer but never materialized via get_output() — they are simply overwritten by the next step’s output.

This document is relevant for: Trn2, Trn3