Skip to content

feat(engine): overlap decode steps with asynchronous token handoff - #524

Open
qinyiqun wants to merge 2 commits into
InfiniTensor:mainfrom
qinyiqun:decode
Open

feat(engine): overlap decode steps with asynchronous token handoff#524
qinyiqun wants to merge 2 commits into
InfiniTensor:mainfrom
qinyiqun:decode

Conversation

@qinyiqun

@qinyiqun qinyiqun commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

Summary

This PR reduces the CPU-side gap between consecutive decode steps by
overlapping sampled-token handoff with CUDA graph submission.

The sampled token is first copied into a stable GPU relay buffer. An
asynchronous D2H transfer into pinned host memory is then queued before the
submission thread starts the next graph. The scheduler waits only for the
corresponding output event when it needs to materialize the token on the host.

No chunked-prefill behavior is introduced by this PR.

Motivation

The original decode path performs host token materialization between two model
forwards. Python scheduling, synchronization, and D2H transfer can therefore
leave the GPU idle between consecutive decode graphs.

The asynchronous handoff path moves graph submission onto a dedicated thread
and preserves stream ordering:

  1. Run the current decode graph.
  2. Copy the sampled token into a stable GPU relay buffer.
  3. Queue an asynchronous D2H copy into pinned memory.
  4. Record the output-ready event.
  5. Submit the next decode graph.
  6. Materialize the previous token only when required by the scheduler.

Configuration

The server exposes:

  • --async-token-handoff=auto
  • --async-token-handoff=on
  • --async-token-handoff=off

auto is the default.

It enables the optimization only when all validated requirements are met:

  • CUDA backend;
  • CUDA graph enabled;
  • paged KV cache;
  • processor support for GPU decode inputs;
  • no speculative draft model;
  • no KV transfer connector;
  • no Mamba state cache;
  • no multimodal request in the scheduled batch.

Explicit on reports the unsupported conditions instead of silently falling
back.

Changes

  • Add the asynchronous model submission and token handoff path.
  • Use double-buffered GPU relay and pinned-host token buffers.
  • Queue D2D relay and asynchronous D2H before the next graph submission.
  • Use per-output device events instead of synchronizing the complete stream.
  • Release the Python GIL while waiting for output events.
  • Use I64 for public token IDs and compiled graph input IDs.
  • Retain input tensors until their device work is complete.
  • Support stable paged decode batches with multiple requests.
  • Add processor capability declarations.
  • Auto-disable unsupported multimodal and specialized processor paths.
  • Update lookahead scheduling for the current prefix-cache interface.
  • Harden startup, shutdown, exception, rollback, and partial-initialization
    handling.

Validation

Build

  • InfiniCore rebuilt with --graph=y, SM80, CCL, ATen, cuDNN, and
    FlashAttention.
  • InfiniCore libraries were reinstalled.
  • InfiniCore and InfiniLM editable Python packages were reinstalled.
  • Python compile checks and Git diff checks passed.

Runtime

  • Explicit asynchronous D2H copy and dtype validation passed.
  • 8B TP=1 service inference passed.
  • 8B concurrent-request service inference passed.
  • 70B TP=8 service inference passed.
  • All eight GPUs participated in the 70B run.
  • Both models produced correct Chinese output.
  • Service logs confirmed:
    • enable_graph=True
    • async_token_handoff=auto
    • Async GPU token handoff enabled
    • asynchronous handoff fast-path activation

A 70B TP=8 functional run produced approximately 41.5 output tokens/s for a
256-token response. This is reported as a smoke-test observation rather than a
formal performance claim because its prompt length and repetition methodology
do not exactly match the existing baseline suite.

Compatibility

Unsupported models and hardware platforms retain the existing synchronous
decode path. The optimization is not enabled globally unless the runtime and
processor explicitly satisfy the compatibility checks.

Dependency

Depends on the InfiniCore asynchronous tensor copy PR: <InfiniCore PR link>.

The first commit preserves the original optimization work and authorship of
@GodKickMyAss. The follow-up integration, compatibility, and lifecycle changes
are authored by @qinyiqun.

@qinyiqun
qinyiqun requested a review from a team August 3, 2026 07:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants