[unified-memory] Drop the vacated 'dense' qualifier and the restating comments - #37170
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7375b050cb
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| f"{cfg.cuda_graph_backend_prefill!r}; pass " | ||
| "--cuda-graph-backend-prefill=disabled." | ||
| ) | ||
| _cg_cfg.prefill.backend = Backend.DISABLED |
There was a problem hiding this comment.
When unified memory uses the default prefill CUDA-graph backend and mem_fraction_static is not explicitly set, this disables capture only after handle_gpu_memory_settings has already run in arg_groups/pipeline.py. Consequently, reserve_for_graph_mb still charges the prefill-graph reserve—1.5 GiB for MLA models—and permanently lowers the inferred KV-cache fraction even though those graphs are never created, reducing unified-pool capacity. Apply this default before the memory hook or recompute the memory setting after changing the backend.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f2d9561d43
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| max_seq = int(slc.max()) | ||
| else: | ||
| max_seq = self.req_to_token.shape[1] | ||
| max_pages = max(-(-max_seq // self.page_size), 1) |
There was a problem hiding this comment.
When unified memory is used with FA3/FA4 prefill context parallelism and prepare_mlp_sync_batch adds padding, seq_lens_cpu.max() excludes the padding even though cache_seqlens_int32 includes it. The FA backend explicitly widens its raw page table by pad_delta at flashattention_backend.py:998-1019, but the unified path subsequently replaces that table with this narrower translated table. FA then reads beyond the translated row width for padded queries, potentially crossing into another row or outside the allocation. Size the translated table using the CP-adjusted metadata bound, or preserve the widened sink columns during replacement.
Useful? React with 👍 / 👎.
f2d9561 to
2656f18
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2656f18c86
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| @@ -827,6 +827,8 @@ def init_new( | |||
|
|
|||
| device = model_runner.device | |||
|
|
|||
| model_runner.kv_index_translator.rebind_write_loc(ret) | |||
There was a problem hiding this comment.
When unified memory is enabled with the explicitly supported DSPARK path, this rebind only runs for ForwardBatch.init_new, but dspark_draft.py:415-421 constructs the real draft ForwardBatch directly using cache locations produced from req_to_token in dspark_planner.py:836-845, so those locations remain virtual. The updated Triton/TRTLLM MLA paths now assume every real batch already carries kernel-facing write IDs and copy these values directly into their write buffers, causing draft KV writes to target incorrect rows after virtual and physical page IDs diverge. Ensure the DSPARK constructor invokes the translator or move this conversion to a boundary shared by directly constructed batches.
Useful? React with 👍 / 👎.
172e669 to
06bd62d
Compare
… KV vocabulary 'dense' named a per-layer KV view only in contrast to the strided MHA view, which no longer exists, so the word marks nothing. Where it qualified a view it is dropped; where it named the id space the views index it becomes 'kernel-facing', the term the translator already uses. Untouched, because the word still carries meaning there: MoE dense layers and first_k_dense_replace, dense-vs-sparse attention (DSA fallback, minimax dense_layer_ids, dual-chunk), trtllm_mla's dense [bs, draft] query indptr, and the mamba conv/SSM 'dense pitch' -- that pool really is envelope-strided, so its contrast is live.
06bd62d to
22e8a43
Compare
1 participant
Stacked on #34613.
Two cleanups the unified-memory stack left behind.
1. Drop the vacated
densequalifier from the KV vocabulary.densenamed aper-layer KV view only in contrast to the strided MHA view, which no longer
exists, so the word marks nothing. Where it qualified a view it is dropped; where
it named the id space the views index it becomes
kernel-facing, the term thetranslator already uses.
Deliberately untouched, because the word still carries meaning there: MoE dense
layers and
first_k_dense_replace, dense-vs-sparse attention (DSA fallback,minimax
dense_layer_ids, dual-chunk), trtllm_mla's dense[bs, draft_token_num]query indptr, and the mamba conv/SSM "dense pitch" -- that pool really is
envelope-strided, so its contrast is live.
2. Cut the comments that restate the line below them. Eight copies of a
mock-runner comment that narrated its own call, four copies of the argument for
seeding a field to
None(the house rule already says to), a duplicated"check the fast path" line detached from the branch it described, and one
ungrammatical restatement. Also replaces the em-dashes this stack introduced in
comments and docstrings with
--, per.claude/rules/comment-style.md.No behavior change. Identifier renames are file-local (verified) apart from the
test-class and test-method names, which nothing outside their own files reads.
CI States
Latest PR Test (Base): ⏳ Run #33366381968
Latest PR Test (Extra): ❌ Run #33366381821
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33366382047