Skip to content

[Perf] Avoid O(n) re-allocation per token in Req.append_host - #143

Open
joshfuery wants to merge 3 commits into
sgl-project:mainfrom
joshfuery:main
Open

[Perf] Avoid O(n) re-allocation per token in Req.append_host#143
joshfuery wants to merge 3 commits into
sgl-project:mainfrom
joshfuery:main

Conversation

@joshfuery

Copy link
Copy Markdown

Motivation

Currently, Req.append_host uses torch.cat to append tokens. This requires a fresh tensor allocation and O(n) memory copy on every single decode step. On the scheduler thread, this CPU overhead scales quadratically per request, eventually causing the GPU to stall during overlap scheduling at larger context lengths.

Changes

  • Lazily allocate a max_device_len buffer on the first append.
  • Write subsequent tokens in place (worst-case O(1)).
  • Fall back to input_ids = _host_buf[:new_len] to preserve views/references.

Microbenchmark

(64 requests, CPU append cost per step)

  • 4k context: 0.18ms (old) vs 0.22ms (new)
  • 16k context: 0.34ms (old) vs 0.22ms (new)
  • 64k context: 0.99ms (old) vs 0.22ms (new)
  • 128k context: 2.73ms (old) vs 0.22ms (new)

Testing

Added test_req_append.py to verify:

  • Token-for-token equivalence with old torch.cat behavior.
  • Dtype preservation.
  • Write-once guarantee (previously written positions are never mutated, keeping radix-tree cache keys safe).
  • Correct overflow handling.

Note: All CPU unit tests pass locally, but I do not have a local GPU set up to run the full end-to-end scheduler tests.

joshfuery added 3 commits July 9, 2026 10:21
Refactor append_host to use a pre-allocated buffer for efficiency.
Corrected a typo in the comment for clarity.
Verifies the behavior of the append_host method in the Req class, ensuring it matches expected behavior in various scenarios, including dtype preservation and overflow handling.
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.

1 participant