Repository navigation
[None][feat] Rubin merge back: MLA, DSpark, DWDP, Helix - #21067
mingyangHao wants to merge 21 commits into
Conversation
The mocked target-only manager must expose an empty draft-layer list so layer-id mapping does not iterate an unconstrained Mock. Keep pool-layout and role assertions unchanged. Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
Avoid duplicate dummy-cache allocation when target and draft share a manager. Restore the SM107 allocator mapping present in the GitLab source and retain the independent RMSNorm/RoPE support gate. Mark the inherited Mamba quota override and complete partial Helix/Mamba manager fixtures for the unified draft cache interface. Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
| draft_num_layers: int = 0, | ||
| draft_num_kv_heads: int = 0, | ||
| draft_head_dim: int = 0, | ||
| draft_dtype: DataType = DataType.BF16, | ||
| draft_kv_factor: int = 2, | ||
| draft_total_num_kv_heads: int | None = None, | ||
| draft_scratch_tokens: int = 0, | ||
| draft_window_size: int | None = None, | ||
| draft_attention_backend: str | None = None, | ||
| draft_max_position_embeddings: int | None = None, |
There was a problem hiding this comment.
🟠 Please drop 272f7bf274 (my port of the GitLab standalone-DSpark draft-KV transfer, !10774 + !10802) from this merge-back. Main is redesigning the same feature on KVCacheLayerSpec(owner=DRAFT) in #20486 / #20664, while this commit adds the older draft_* kwargs path to KVCacheManagerV2.__init__ (kv_cache_manager_v2.py:1856) next to layer_specs, so main would carry two unified-draft-cache implementations; git merge-tree against #20486 and #20664 already conflicts in _util.py, kv_cache_manager_v2.py, dflash.py and dspark.py.
That also covers the follow-up pieces of f189f03403 / b386e7472a / f4587daabc that only adapt it (managed draft history, _uses_unified_standalone_draft_cache, the num_draft_layers peer check whose error points at a TRTLLM_DISAGG_SHARED_DRAFT_KV nothing in this PR reads).
The DSpark kernel/perf parts in the same commit (FP8 projection fusions, Rubin CuteDSL GEMMs, paged attention / RMSNorm+RoPE page write) can stay if they split cleanly; the standalone draft-KV transfer will be redone on top of #20664.
Remove the superseded standalone draft KV ownership, lifecycle and CTX-to-GEN transfer path, including its later integration adaptations. Restore the main worker/cache-manager contracts in preparation for the layer_specs based design. Keep the independent FP8 projections, Rubin GEMMs, paged attention and page-write kernels. Let the restored rolling-window path reuse the loaded projection modules, with CPU regressions for attention and context-window writes. Preserve the original GitLab replay commits, including 272f7bf, unchanged. Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
Keep main's _admit_run_context call, which validates the context and already waits on x_ready_event before dispatch. This supersedes the duplicate direct wait from the earlier mergeback adaptation without dropping synchronization. Preserve the source replay history and the legacy draft-KV removal. Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #81417 [ run ] triggered by Bot. Commit: |
|
PR_Github #81417 [ run ] completed with state
|
Retain MiniMax DSpark support from main together with the existing projection and Rubin imports. The conflict is limited to the typing and attention imports. Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
Initialize the selected linear backend in the confidence-loader fixture, exercise the rolling-window stage entry point in V4.1 tests, and expect the retained SM107 allocator support. Preserve the existing behavioral assertions. Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
Keep the converted query when dispatching through the new windowed helper, and round the future-block K/V to the cache precision before merging attention. This preserves the retained FP8 cache path after main refactored windowed decode. Cover BF16 and FP8 causal/noncausal calls at the CPU kernel boundary and compare against a dense oracle over the rounded cache values. Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
|
/bot run --disable-fail-fast |
mingyangHao
changed the title
|
PR_Github #81595 [ run ] triggered by Bot. Commit: |
Backport cc6b46096defce52488a27f06c9be25a43ad55f4. Define the join on the fused FC12 implementation, which does not inherit the grouped-GEMM implementation, and consume the input-ready event on empty batches too. Cover repeated dispatch within a forward, fresh events on the next forward, empty inputs, absent events, and output allocation/reuse with CPU tests. Signed-off-by: Jonas Li <jonasl@nvidia.com> Signed-off-by: Mingyang Hao <mingyangh@nvidia.com>
|
/bot kill |
|
/bot run --disable-fail-fast |
|
PR_Github #81642 [ kill ] triggered by Bot. Commit: |
|
PR_Github #81643 [ run ] triggered by Bot. Commit: |
|
PR_Github #81642 [ kill ] completed with state |
|
PR_Github #81595 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #81676 [ run ] triggered by Bot. Commit: |
|
PR_Github #81643 [ run ] completed with state |
|
PR_Github #81676 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #81701 [ run ] triggered by Bot. Commit: |
|
PR_Github #81701 [ run ] completed with state
|
Dev Engineer Review
QA Engineer Review
unittest/disaggregated/test_extractor.pyin bothl0_a10.ymlandl0_h100.yml. It found no exact matching entries for the other searched changed tests.l0_dgx_b200.ymllists a different autotuner test,test_autotuner_distributed_strategy. No matching manual-QA entries were found.Per-File QA Perspective
Source files
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h— VerifyCacheDomainvalues and the documented storage order.cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h— Verify thatAttentionLayerConfigdefaults to the target cache domain.cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.cpp— Verify lifecycle construction forwards the configured cache domain.cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h— Verify domain-aware comparison and scratch-range behavior.cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp— Verify pool grouping separates cache domains while preserving intended reuse.cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h— Verify the new fusion and quantization enum values match dispatch consumers.cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu— Verify FP8 block quantization, scale packing, padding, and launch validation.cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h— Verify callers setsfLeadingDimcorrectly.cpp/tensorrt_llm/kernels/customAllReduceKernels.h— Verify the fusion operation name and value match kernel and Python dispatch.cpp/tensorrt_llm/kernels/helixAllToAll.cu— Verify the asynchronous shared-memory fence protects FIFO data before reuse.cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp— Verify binding defaults and properties match the C++ API.cpp/tensorrt_llm/thop/allreduceOp.cpp— Verify output padding and the returned scale-view shape.tensorrt_llm/_torch/attention/backends/dspark.py— Verify captured-context writes respect masks, capacities, and page-table bounds.tensorrt_llm/_torch/attention/backends/fp4_mla/cache_manager.py— Verify draft-layer dtype and cache-size accounting across Helix and non-Helix configurations.tensorrt_llm/_torch/attention/backends/fp4_mla/config.py— Verify the wave-plan environment setting and default.tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_cutedsl_mufu16.py— Verify packed-task bounds, split-output merging, and pipeline synchronization.tensorrt_llm/_torch/attention/backends/fp4_mla/fp4_mla_packed_reduce.py— Verify candidate limits, makespan selection, tie-breaking, and disabled-planner behavior.tensorrt_llm/_torch/attention/backends/fp4_mla/packed_plan.py— Verify environment settings reach the planner launch.tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py— Verify draft-cache quota and fixed-reservation calculations.tensorrt_llm/_torch/attention/backends/sparse/dsa/cache_manager.py— Verify local-layer pool sizing and draft-cache accounting.tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention.py— Verify paged-attention arguments and supported geometry match callers.tensorrt_llm/_torch/attention/kernels/blackwell/dspark/attention_kernel.py— Verify history-page masking, capacity handling, and FP8 output scales.tensorrt_llm/_torch/attention/kernels/blackwell/dspark_rmsnorm_rope.py— Verify quantization scales and bounds checks on managed page writes.tensorrt_llm/_torch/attention/mla.py— Verify chunked-prefill dispatch and Helix V-projection behavior against fallback paths.tensorrt_llm/_torch/autotuner.py— Verify the full tactic list is used when it is shorter than the tuning group.tensorrt_llm/_torch/custom_ops/dspark_attention_custom_op.py— Verify custom-op validation, kernel caching, and the removed cache/slot interface.tensorrt_llm/_torch/custom_ops/dspark_rmsnorm_rope_custom_op.py— Verify device-specific compilation, page-write validation, and FP8 precompile gating.tensorrt_llm/_torch/disaggregation/native/auxiliary.py— Verify draft-history serialization and recycled-slot clearing.tensorrt_llm/_torch/disaggregation/native/bounce/impl.py— Verify byte accounting forcp_as_tppaged groups.tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py— Verify peer checks reject incompatible overlapping attention groups.tensorrt_llm/_torch/disaggregation/native/peer.py— Verify draft-layer compatibility and rank mapping forcp_as_tp.tensorrt_llm/_torch/disaggregation/native/rank_info.py— Verify the draft-layer count reflects the active cache manager.tensorrt_llm/_torch/disaggregation/native/transfer.py— Verify auxiliary-slot requirements and draft-history transfer behavior.tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py— Verify block-ordinal conversion for sharedcp_as_tppages.tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py— Verify draft-page expansion and pool separation in extracted page tables.tensorrt_llm/_torch/disaggregation/resource/page.py— Verify deserialization remains compatible when the new group fields are absent.tensorrt_llm/_torch/disaggregation/transceiver.py— Verify draft-history validation and restoration.tensorrt_llm/_torch/distributed/ops.py— Verify the new fusion operation reaches supported MNNVL dispatch.tensorrt_llm/_torch/models/modeling_dflash.py— Verify FP8 page writes and block-query layout conversion.tensorrt_llm/_torch/models/modeling_dspark.py— Verify managed-page and legacy paths, projection fallback, and required paged metadata.tensorrt_llm/_torch/models/modeling_kimi_linear.py— Verify the Helix V-projection weight path when enabled.tensorrt_llm/_torch/moe/fused_moe/cutedsl/fused_fc12.py— Verify the input-ready event is joined before fused execution.tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py— Verify event ordering before either execution path.tensorrt_llm/_torch/pyexecutor/_util.py— Verify unified draft-cache selection and rejection of unsupported manager combinations.tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py— Verify draft-cache preparation and history snapshots occur only in applicable decode steps.tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py— Verify draft allocation, sizing, history lifecycle, and target/draft pool ratios.tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py— Verify draft storage is included in hybrid-cache cost estimates.tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py— Verify draft pools use base-manager roles and free-block counts.tensorrt_llm/_torch/pyexecutor/py_executor.py— Verify pending history updates publish after sampler updates and handle absent state.tensorrt_llm/_torch/pyexecutor/sampler/sampler.py— Verify the optional update field remains compatible with sample-state consumers.tensorrt_llm/_torch/pyexecutor/standalone_draft_cache.py— Verify history validation and stale-cache protection during publication.tensorrt_llm/_torch/speculative/dflash.py— Verify unified-cache preparation, page-table fallback, and overflow handling.tensorrt_llm/_torch/speculative/dspark.py— Verify managed-page allocation, context seeding, and progress-slot handling.tensorrt_llm/_torch/speculative/interface.py— Verify managed-cache hooks and padding-slot validation across worker implementations.tensorrt_llm/functional.py— Verify residual-input validation permits only the intended new fusion.tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py— Verify the publicCacheDomainexport.tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi— Verify stub signatures and defaults match the bound API.Test files
cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp— Covers cache-domain defaults, lifecycle comparison, and pool grouping. No test-list entry was established.tests/unittest/_torch/attention/multi_gpu/test_mla_helix.py— Covers V-up projection placement and Helix all-to-all behavior in CUDA graphs. No matching test-list entry was found.tests/unittest/_torch/attention/test_fp4_mla_packed.py— Covers planner wave simulation, ordering, and disabled-planning behavior. No matching test-list entry was found.tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py— Covers standalone draft-cache sizing and dtype selection. No matching test-list entry was found.tests/unittest/_torch/misc/test_autotuner.py— Covers short tactic lists across Helix CP ranks. The list search found a different autotuner test, not this test by name.tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py— Covers managed-cache block-table decoding and FP8 context-page writes. No matching test-list entry was found.tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py— Covers projection paths, backend selection, and page writes. No matching test-list entry was found.tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py— The reported change removes comments only; test behavior is unchanged. No matching test-list entry was found.tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py— Covers managed draft pages, request slots, context seeding, and generation progress. No matching test-list entry was found.tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py— Covers paged decode parity for BF16 and FP8 histories. No matching test-list entry was found.tests/unittest/disaggregated/region/test_aux.py— Covers draft-history serialization round trips. No matching test-list entry was found.tests/unittest/disaggregated/test_extractor.py— Updates the fake-manager fixture. Listed inl0_a10.ymlandl0_h100.yml.Description
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.