refactor(runtime): make ForwardContext tensor-free (restore "Do not contain Tensor")'#745
Draft
rjzhb wants to merge 2 commits into
Draft
refactor(runtime): make ForwardContext tensor-free (restore "Do not contain Tensor")'#745rjzhb wants to merge 2 commits into
rjzhb wants to merge 2 commits into
Conversation
rjzhb
force-pushed
the
rjzhb/refactor/ctx-drop-draft-tensors
branch
from
July 22, 2026 04:31
32c9be9 to
166ae35
Compare
Signed-off-by: rjzhb <rjzhb222@163.com> Co-authored-by: Cursor <cursoragent@cursor.com>
rjzhb
force-pushed
the
rjzhb/refactor/ctx-drop-draft-tensors
branch
from
July 22, 2026 04:34
166ae35 to
337bfb2
Compare
Pass the drafter's mutable draft_seq_lens_buf via the existing seq_lens forward kwarg instead of a separate draft_seq_lens parameter. Only forward seq_lens into the model when accept_lengths is present so target models are unchanged. Signed-off-by: rjzhb <rjzhb222@163.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Goal
ForwardContextdocuments itself as "Do not contain Tensor", yet fiveloose
torch.Tensorfields have crept in over time:req_to_page,gather_ids,draft_seq_lens_buf,accept_lengths,dsa_swa_slot_mapping.This (draft) PR is the umbrella for making the context tensor-free again —
routing each tensor through a persistent buffer or an object reference
so nothing loose rides on the per-forward context.
Not a bugfix: today these are only read on eager paths or are persistent
buffers, so CUDA-graph replay is unaffected. This is invariant / hygiene
cleanup — keeping the context from silently becoming a tensor grab-bag.
Approach
Draft-related tensors are bundled into a drafter-owned
SpecForwardState,referenced once via
ctx.spec(an object ref, same category asattn_backend/
token_to_kv_pool). The rest move to their natural owner (backend / inputbuffers) or have their read site restructured.
Status
draft_seq_lens_buf→ctx.spec.draft_seq_lensaccept_lengths→ctx.spec.accept_lengthsreq_to_page— persistent page-table buffer; only 1 ctx read left (glm5),the rest already flow as explicit args. Small.
gather_ids— per-forward; ~16 reads across 6 models + logits. Needsdeciding where the first-step row-narrowing lives. Structural.
dsa_swa_slot_mapping— deepseek_v4 intra-forward scratch (set/clearedwithin one forward); different mechanism.