[None][feat] Support Helix CP with MTP - #16003
Conversation
94405a6 to
8e54404
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57845 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR adds Helix context-parallel (CP) support for one-model speculative decoding (MTP/Eagle3). Changes span flattened generation attention metadata, KV cache decode-group ownership/reserve/rewind, request bookkeeping fields, sampler mapping repurposing (CP-to-TP), overlap-scheduler compatibility gating, model wiring, and associated tests/docs/CI updates. ChangesHelix CP speculative decoding support
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant ModelEngine as PyTorchModelEngine
participant AttnMetadata as TrtllmAttentionMetadata
participant KVCacheManager as KVCacheManager
participant MLA as MLA._attn_forward_gen
ModelEngine->>AttnMetadata: prepare() -> _maybe_prepare_helix_flatten()
AttnMetadata->>AttnMetadata: fill flattened per-row KV/prompt buffers
ModelEngine->>ModelEngine: _preprocess_inputs -> apply overlap corrections (sign=1)
ModelEngine->>KVCacheManager: prepare_resources -> _helix_prepare_generation_kv
KVCacheManager->>KVCacheManager: reserve tokens only on owning rank
ModelEngine->>MLA: forward generation attention
MLA->>AttnMetadata: enter helix_flattened_generation()
AttnMetadata-->>MLA: flattened cu_q_seqlens/cu_kv_seqlens
MLA->>MLA: attn_backend.forward(kwargs)
AttnMetadata-->>MLA: restore original runtime metadata
ModelEngine->>KVCacheManager: update_resources -> _helix_rewind_generation_kv (on rejection)
KVCacheManager->>KVCacheManager: rewind_kv_cache only if owner
ModelEngine->>ModelEngine: _postprocess_inputs -> undo overlap corrections (sign=-1)
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
tests/unittest/_torch/speculative/hw_agnostic/test_helix_eagle3_ownership.py (3)
25-26: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDuplicate ownership formula defined twice.
owns_decode_index(Line 25-26) andowns_decode_group(Line 117-119) are identical implementations of(idx // tpb) % cp_size == cp_rank. Consider defining a single shared helper and aliasing/reusing it for both the verify-token and decode-group test sections to avoid drift between the two copies.♻️ Proposed consolidation
-def owns_decode_index(decode_index, tpb, cp_size, cp_rank): - return (decode_index // tpb) % cp_size == cp_rank +def owns_decode_index(decode_index, tpb, cp_size, cp_rank): + """Mirror of the shared Helix CP ownership formula.""" + return (decode_index // tpb) % cp_size == cp_rank + + +# Verify-group and decode-group ownership use the same formula in production +# (see resource_manager._helix_owns_decode_group / model_engine helpers). +owns_decode_group = owns_decode_indexAlso applies to: 117-119
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/speculative/hw_agnostic/test_helix_eagle3_ownership.py` around lines 25 - 26, The ownership check is duplicated in both owns_decode_index and owns_decode_group, so consolidate the shared `(idx // tpb) % cp_size == cp_rank` logic into one helper and reuse or alias it from both the verify-token and decode-group test sections. Update the tests to call the single shared function instead of maintaining two identical implementations, keeping the existing names if needed for readability.
25-46: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftTests validate mirror formulas, not the actual production implementations.
owns_decode_index/owns_decode_group/verify_token_paramsare explicitly labeled as mirrors ofresource_manager._helix_owns_decode_groupandmodel_engine._helix_verify_token_params(Line 39, Line 118), but the tests only check internal consistency of these local re-implementations. If the production formulas change or diverge (e.g., a future edit to_helix_owns_decode_groupinresource_manager.py), this "hardware-agnostic" suite would keep passing without detecting the drift, per <ai_summary> and the provided context snippet fromresource_manager.py:738-750showing the real_helix_owns_decode_group.Given this is meant to guard the ownership/FIFO invariants underpinning Helix CP correctness, consider either:
- Importing and testing the real
_helix_owns_decode_group(with a lightweight stubMapping/selfobject) directly instead of a hand-rolled mirror, or- Adding at least one assertion elsewhere (e.g., in the resource_manager/model_engine test suites) that pins the production formula's behavior against this mirror, so regressions in either file are caught.
As per path instructions for
tests/**, flagging coverage as insufficient given the risk of silent drift between mirror and production logic; a follow-up outside this PR that exercises the real implementations would close the gap.Also applies to: 117-166
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/speculative/hw_agnostic/test_helix_eagle3_ownership.py` around lines 25 - 46, The test helpers here are only re-implementing production logic, so they can drift without catching regressions. Update the ownership coverage in test_helix_eagle3_ownership to exercise the real resource_manager._helix_owns_decode_group and model_engine._helix_verify_token_params (using a lightweight stub object if needed) instead of only asserting the local mirror functions agree with themselves. If direct import is not practical, add a production-facing assertion in the resource_manager/model_engine test suite that pins the real formulas against these expectations.Source: Path instructions
25-26: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing return-type/parameter annotations on module-level functions and helper class.
Functions such as
owns_decode_index,verify_token_params,owns_decode_group, and_FifoRankState.__init__/reserve/rewindlack type annotations. As per coding guidelines, "Always annotate functions with return types (useNoneif no return); annotate class members and variables when necessary."🏷️ Example annotations
-def owns_decode_index(decode_index, tpb, cp_size, cp_rank): +def owns_decode_index(decode_index: int, tpb: int, cp_size: int, cp_rank: int) -> bool: return (decode_index // tpb) % cp_size == cp_rank-def verify_token_params(total_input_len, g, num_draft, tpb, cp_size, cp_rank): +def verify_token_params( + total_input_len: int, g: int, num_draft: int, tpb: int, cp_size: int, cp_rank: int +) -> tuple[list[int], list[bool], int]:- def __init__(self, tpb, cp_size, cp_rank): + def __init__(self, tpb: int, cp_size: int, cp_rank: int) -> None: self.tpb = tpb self.cp_size = cp_size self.cp_rank = cp_rank - self.group_index = 0 - self.pending = [] - self.reserve_log = [] - self.rewind_log = [] + self.group_index: int = 0 + self.pending: list[bool] = [] + self.reserve_log: list[bool] = [] + self.rewind_log: list[bool] = []Also applies to: 38-46, 117-119, 144-166
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/speculative/hw_agnostic/test_helix_eagle3_ownership.py` around lines 25 - 26, Add explicit type annotations to the untyped module-level helpers and FIFO state methods in the Helix Eagle3 ownership test module. Update owns_decode_index, verify_token_params, and owns_decode_group with parameter and return types, and annotate _FifoRankState’s __init__, reserve, and rewind signatures (including any instance attributes they set). Keep the fixes localized to those named symbols so the test helpers follow the project’s annotation guidelines.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Around line 513-523: The `_helix_flat_block_offsets` allocation is using
`kv_cache_manager.num_pools`, but the flatten/copy path follows
`kv_cache_block_offsets` and may use `num_attention_op_pools` instead. Update
the buffer shape in the `kv_cache_manager is not None` branch of `trtllm.py` to
size `_helix_flat_block_offsets` from `num_attention_op_pools` so it matches the
source tensor and remains valid when the two pool counts differ.
In `@tensorrt_llm/_torch/speculative/interface.py`:
- Around line 874-892: The new cache field and accessor in sampler_mapping need
explicit type annotations so mypy does not infer None-only or Any. Annotate
self._sampler_mapping_cache with the concrete optional mapping type used by
mapping.repurpose_helix_cp_to_tp(), and add an explicit return type to the
sampler_mapping property consistent with the surrounding members in this class.
Use the existing model_config/mapping logic to determine the precise annotation
and keep the cache behavior unchanged.
---
Nitpick comments:
In
`@tests/unittest/_torch/speculative/hw_agnostic/test_helix_eagle3_ownership.py`:
- Around line 25-26: The ownership check is duplicated in both owns_decode_index
and owns_decode_group, so consolidate the shared `(idx // tpb) % cp_size ==
cp_rank` logic into one helper and reuse or alias it from both the verify-token
and decode-group test sections. Update the tests to call the single shared
function instead of maintaining two identical implementations, keeping the
existing names if needed for readability.
- Around line 25-46: The test helpers here are only re-implementing production
logic, so they can drift without catching regressions. Update the ownership
coverage in test_helix_eagle3_ownership to exercise the real
resource_manager._helix_owns_decode_group and
model_engine._helix_verify_token_params (using a lightweight stub object if
needed) instead of only asserting the local mirror functions agree with
themselves. If direct import is not practical, add a production-facing assertion
in the resource_manager/model_engine test suite that pins the real formulas
against these expectations.
- Around line 25-26: Add explicit type annotations to the untyped module-level
helpers and FIFO state methods in the Helix Eagle3 ownership test module. Update
owns_decode_index, verify_token_params, and owns_decode_group with parameter and
return types, and annotate _FifoRankState’s __init__, reserve, and rewind
signatures (including any instance attributes they set). Keep the fixes
localized to those named symbols so the test helpers follow the project’s
annotation guidelines.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 257c5863-2838-49e8-ac88-f120133f34cf
📒 Files selected for processing (20)
docs/source/features/feature-combination-matrix.mdtensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/models/modeling_deepseekv3.pytensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/modules/mla.pytensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/_torch/pyexecutor/request_utils.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/_torch/speculative/eagle3.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/mtp.pytensorrt_llm/_torch/speculative/utils.pytests/integration/defs/accuracy/accuracy_core.pytests/integration/defs/accuracy/test_disaggregated_serving.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/test-db/l0_dgx_b200.ymltests/integration/test_lists/waives.txttests/unittest/_torch/speculative/hw_agnostic/test_helix_eagle3_ownership.py
|
/bot run --disable-fail-fast |
|
PR_Github #57875 [ run ] triggered by Bot. Commit: |
|
PR_Github #57845 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #57892 [ run ] triggered by Bot. Commit: |
|
PR_Github #57875 [ run ] completed with state |
| py_helix_is_inactive_rank, set at KV reserve time. Overlap corrections | ||
| run on-device in _preprocess_inputs. | ||
| """ | ||
| del tokens_per_block |
There was a problem hiding this comment.
Why is this deleted right away?
| kv_len_offsets_device[previous_slots], | ||
| non_blocking=True) | ||
| if _has_cp_helix: | ||
| # Owner-gated kv-length correction for overlap scheduling. |
There was a problem hiding this comment.
This is called at every forward pass, right? Could we pre-allocate helix_prev_owned_cuda once?
| seg_starts = torch.cumsum(seg_lens, 0) - seg_lens | ||
| row_incl = (torch.arange(total_q, dtype=torch.int64) - | ||
| torch.repeat_interleave(seg_starts, seg_lens) + 1) | ||
| owned_incl = row_incl * torch.repeat_interleave(active, seg_lens) |
There was a problem hiding this comment.
These might create new buffers in a hot path, could we pre-allocate them?
| for i in range(runtime_draft_len): | ||
| # Set per-token global positions and ownership before draft forward. | ||
| helix_owner_mask = self._helix_draft_owner_mask( | ||
| attn_metadata, inputs["position_ids"], batch_size) |
There was a problem hiding this comment.
Could be casted to attn_metadata.kv_lens_cuda.dtype directly here (or done by helix_draft_owner_mask)
| with launch_disaggregated_llm(disaggregated_server_config, | ||
| ctx_server_config, gen_server_config, | ||
| self.MODEL_PATH) as llm: | ||
| run_accuracy_test(llm, self.MODEL_NAME, ["MMLU", "GSM8K"]) |
There was a problem hiding this comment.
Is the removal of MMLU intended here?
| # Optional GSM8K debug env vars: TRTLLM_GSM8K_NUM_SAMPLES, TRTLLM_GSM8K_OUTPUT_DIR. | ||
| # A custom sample count skips the accuracy hypothesis test. | ||
| gsm8k_num_samples = None | ||
| gsm8k_output_dir = None | ||
| if self.DATASET == "gsm8k": | ||
| gsm8k_output_dir = os.getenv("TRTLLM_GSM8K_OUTPUT_DIR") or None | ||
| num_samples_env = os.getenv("TRTLLM_GSM8K_NUM_SAMPLES") | ||
| if num_samples_env is not None: | ||
| gsm8k_num_samples = int(num_samples_env) | ||
| if gsm8k_num_samples <= 0: | ||
| raise ValueError( | ||
| "TRTLLM_GSM8K_NUM_SAMPLES must be a positive integer, " | ||
| f"got {num_samples_env!r}.") | ||
|
|
||
| if gsm8k_num_samples is not None: | ||
| logger.info( | ||
| f"Running GSM8K on a deterministic subset of {gsm8k_num_samples} " | ||
| "sample(s) (TRTLLM_GSM8K_NUM_SAMPLES) and skipping accuracy " | ||
| "verification.") | ||
| hypothesis_testing_params = HypothesisTestingParams( | ||
| ref_accuracy=0 if self.HIGHER_IS_BETTER else math.inf, | ||
| num_samples=gsm8k_num_samples, | ||
| metric_name=self.METRIC_NAME, | ||
| higher_is_better=self.HIGHER_IS_BETTER) |
There was a problem hiding this comment.
Nit: could be in a separate PR
| return | ||
| off_rows = per_request_kv_offset[self._helix_flat_row_to_seq[:n].to( | ||
| torch.long)].to(self._helix_flat_kv_lens_cuda.dtype) | ||
| self._helix_flat_kv_lens_cuda[:n] += sign * off_rows |
There was a problem hiding this comment.
we might also need to update _helix_flat_kv_lens_cpu ?
syuoni
left a comment
syuoni
left a comment
There was a problem hiding this comment.
Reviewed the DeepSeek modeling part. Thanks!
|
PR_Github #57892 [ run ] completed with state
|
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
45bc291 to
3b9f67b
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58048 [ run ] triggered by Bot. Commit: |
| | Disaggregated Serving | Yes | Yes | Yes | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | | | ||
| | Chunked Prefill | Yes | Yes | Yes | Untested | Yes | Yes | Yes | Yes | --- | | | | | | | | | | | | ||
| | Speculative Decoding — Linear | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | --- | | | | | | | | | | | ||
| | Speculative Decoding — Linear | Yes | Yes | Yes | No | Yes | Yes (Eagle3 one-model) | Yes | Yes | Yes | --- | | | | | | | | | | |
There was a problem hiding this comment.
does the PR support Eagle3+Helix as well?
| - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp1cp4] TIMEOUT (60) | ||
| - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp2cp2-mtp_nextn=0] TIMEOUT (60) | ||
| - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1tp1cp4-mtp_nextn=0] TIMEOUT (60) | ||
| - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v2-cudagraph:with_padding-pp1tp1cp4-mtp_nextn=3] TIMEOUT (60) |
There was a problem hiding this comment.
this test doesn't cover TP + Helix + MTP. could it be changed to pp1tp2cp2-mtp_nextn=3 since it's the only one that would actually catch the allgather bug that motivated the sampler_mapping change?
| with launch_disaggregated_llm(disaggregated_server_config, | ||
| ctx_server_config, gen_server_config, | ||
| self.MODEL_PATH) as llm: | ||
| run_accuracy_test(llm, self.MODEL_NAME, ["MMLU", "GSM8K"]) |
There was a problem hiding this comment.
any particular reason to drop MMLU?
| is_separate_draft_engine: bool = False, | ||
| mapping_with_cp: Optional[Mapping] = None): | ||
| # Thread mapping_with_cp into attention so MLA keeps the full head count. | ||
| super().__init__(model_config, |
There was a problem hiding this comment.
I'm wondering if we have compared acceptance rate with and without Helix offline.
|
PR_Github #58048 [ run ] completed with state
|
|
[by Codex] @SimengLiu-nv Could you review this PR? Thanks! |
Shixiaowei02
left a comment
Shixiaowei02
left a comment
There was a problem hiding this comment.
Hi @brb-nv , given that some Helix CP issues have recently only been caught at the dis-agg E2E test stage, would it be possible to add more modularized test coverage? For example, adding tests between the agg layer and individual modules may help catch these issues earlier. Thank you!
| accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp1cp4] | ||
| accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp1tp2cp2] | ||
| accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[nccl-cudagraph:with_padding-pp2tp1cp2] | ||
| accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype_with_helix[fifo_v1-cudagraph:with_padding-pp1dp2cp2-mtp_nextn=0] |
There was a problem hiding this comment.
To help reduce resource consumption, could we kindly consider reducing the number of DeepSeek-V3 test cases mentioned above? In addition, since several DeepSeek-V3 issues are currently being tracked as P1, it would be greatly appreciated if the issues related to the newly added tests could be addressed (or appropriately reassigned) before introducing additional test coverage. Thank you for your support!
|
[by Codex] @SimengLiu-nv Friendly reminder to review this PR as the primary KV-cache-manager reviewer when you have a chance. Thanks! |
|
[by Codex] @SimengLiu-nv Could you please review this PR? Thank you! |
10 participants
Description
This MR adds support for Helix CP with MTP. Some implementation details:
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Summary by CodeRabbit
New Features
Bug Fixes
Tests