[None][perf] Port remaining DeepSeek V4 optimizations to main#16028
Conversation
ac761c4 to
151fbd5
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #57919 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR adds new CUDA kernels and Torch custom ops for DeepSeek V4 sliding block table computation (with and without scratch buffers), integrates them into the cache manager replacing compiled Python helpers, adds tests/microbenchmarks, introduces DSA indexer query-tiling and chunk-size heuristics, and adds base64-encoded prompt-token-id relay support across disaggregated serving components. ChangesDeepSeek V4 Sliding Block Table Kernels and Integration
Estimated code review effort: 4 (Complex) | ~60 minutes DSA Indexer Logits Tiling and Chunk Size Heuristic
Estimated code review effort: 3 (Moderate) | ~25 minutes Disaggregated Serving Base64 Prompt Token ID Relay
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant OpenAIServer as openai_server
participant DisaggService as OpenAIDisaggregatedService
participant ContextWorker
Client->>OpenAIServer: chat completion request
OpenAIServer->>DisaggService: forward request
DisaggService->>ContextWorker: _get_ctx_request(return_prompt_token_ids_b64)
ContextWorker-->>DisaggService: response with prompt_token_ids_b64
DisaggService->>DisaggService: _get_gen_request relays base64 or int-array ids
DisaggService-->>OpenAIServer: generation request/response
OpenAIServer->>OpenAIServer: decode/encode base64 prompt_token_ids
OpenAIServer-->>Client: chat completion response
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: 1
🧹 Nitpick comments (6)
tensorrt_llm/_torch/attention_backend/sparse/dsa.py (1)
470-479: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPrefix this helper if it is not part of the module API.
select_indexer_chunk_sizeappears to be an internal helper. Rename it to_select_indexer_chunk_sizeunless tests or callers intentionally import it. As per coding guidelines, “Variables and functions not part of a class’s or module’s public interface should be prefixed with an underscore.”🤖 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 `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py` around lines 470 - 479, The helper select_indexer_chunk_size in dsa.py looks internal and should not be part of the module API unless it is intentionally imported elsewhere. Rename it to _select_indexer_chunk_size, and update any local references or tests that call it so the sparse attention backend keeps the helper clearly private.Source: Coding guidelines
cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu (1)
238-309: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMinor duplication: inline scale/offset math duplicates
applyScaleAndOffset.Lines 300-307 of
computeSlidingBlockTablesWithScratchRowsKernelreimplement the samebase == kBadPageIndex ? kBadPageIndex : base * scale + layerOffsetternary already factored intoapplyScaleAndOffset(line 66-69). Reusing the helper here would remove the risk of the two implementations silently diverging in the future.♻️ Proposed refactor
- int32_t const base = blockOffsets[blockOffsetsOffset + blockId]; - output[outputOffset + blockId] = base == kBadPageIndex ? kBadPageIndex : base * scale + layerOffset; + int32_t const base = blockOffsets[blockOffsetsOffset + blockId]; + output[outputOffset + blockId] = applyScaleAndOffset(base, scale, layerOffset);🤖 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 `@cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu` around lines 238 - 309, The end of computeSlidingBlockTablesWithScratchRowsKernel reimplements the same scale/offset and bad-page handling already centralized in applyScaleAndOffset. Replace the inline base-to-output mapping in the non-scratch path with a call to applyScaleAndOffset so the kernel uses the shared helper consistently and avoids duplicated logic drifting apart.tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py (1)
554-558: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
self._device_block_positionsappears to be dead code now.This staging tensor was almost certainly consumed by the removed
@maybe_compile-decorated Python helpers (per the line-range change details, those were removed around line 168) that computed sliding block tables. The new CUDA kernels compute block positions internally viathreadIdx/loop index rather than taking an explicitblock_positionstensor (confirmed by the op signatures indeepseekV4BlockTableOp.cppand by the microbenchmark's_custom_sliding_block_tables_with_scratch, which explicitlydelsblock_positionsbefore calling the op). Scanning the rest of this file,self._device_block_positionsis never read again after this allocation. Recommend removing it to avoid confusing future readers and the small unnecessary allocation.🤖 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 `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py` around lines 554 - 558, Remove the unused staging tensor allocation in the cache manager setup: `self._device_block_positions` is no longer referenced anywhere in `CacheManager` after the removed Python helpers, and the CUDA block-table ops now compute positions internally. Delete the `torch.arange(...)` assignment in the initialization path and make sure no remaining methods in `cache_manager.py` still depend on `self._device_block_positions`.cpp/tensorrt_llm/kernels/deepseekV4BlockTable.h (1)
29-39: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing Doxygen documentation on public function prototypes.
Neither
invokeDeepseekV4ComputeSlidingBlockTablesnorinvokeDeepseekV4ComputeSlidingBlockTablesWithScratchhas any//!documentation describing parameters (e.g. tensor layouts/shapes forblockOffsets,poolIds, semantics ofscratchBegs/scratchEnds/scratchSlots). As per coding guidelines,**/*.{h,hpp,hxx}files should "Follow Doxygen rules for documenting class interfaces and function prototypes; use//!for C++ style comments and//!<for member comments."📝 Example doc-comment addition
+//! Computes per-layer sliding-window block tables (no scratch remapping). +//! blockOffsets: [numPools, copyIdxCapacity, 2, maxBlocksPerSeq] int32 device tensor. +//! output: [numLayers, numAttnTypes, numTables, maxBlocksPerSeq] int32 device tensor. void invokeDeepseekV4ComputeSlidingBlockTables(int32_t const* blockOffsets, int32_t const* copyIdx, int64_t const* poolIds, bool const* validPool, int32_t const* scales, int32_t const* layerOffsets, int32_t* output, int32_t numPools, int32_t copyIdxCapacity, int32_t numLayers, int32_t numAttnTypes, int32_t numTables, int32_t maxBlocksPerSeq, cudaStream_t stream);🤖 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 `@cpp/tensorrt_llm/kernels/deepseekV4BlockTable.h` around lines 29 - 39, Add Doxygen-style `//!` comments for the public prototypes `invokeDeepseekV4ComputeSlidingBlockTables` and `invokeDeepseekV4ComputeSlidingBlockTablesWithScratch` in `deepseekV4BlockTable.h`, documenting each parameter’s meaning, expected tensor/layout shapes, and any scratch-related semantics. Keep the comments directly above the declarations and make sure they follow the project’s Doxygen rules for C++ function prototypes.Source: Coding guidelines
cpp/tensorrt_llm/thop/deepseekV4BlockTableOp.cpp (1)
60-99: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winNo validation that
scalesare strictly positive before they're used as a divisor.Both scratch CUDA kernels (
computeSlidingBlockTablesWithScratchKernelandcomputeSlidingBlockTablesWithScratchRowsKernelindeepseekV4BlockTable.cu) computetotalOffset / scaleandtotalOffset % scalewithout any zero-check.checkCommonInputsvalidates dtype/shape/contiguity forscalesbut never checks its values. Todayscalesalways originates fromconverter.scaleincache_manager.pyand is presumably always ≥ 1, but since this is a public Torch op boundary, a cheap value check would prevent undefined behavior/crash if that invariant is ever violated by a future caller.🛡️ Proposed guard
checkInt32Tensor(scales, "scales"); checkInt32Tensor(layerOffsets, "layer_offsets"); checkInt32Tensor(output, "output"); + TORCH_CHECK(scales.min().item<int32_t>() > 0, "scales must be strictly positive");🤖 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 `@cpp/tensorrt_llm/thop/deepseekV4BlockTableOp.cpp` around lines 60 - 99, `checkCommonInputs` currently validates `scales`’s dtype and shape but not its value range, while the CUDA kernels it feeds divide and mod by `scale`. Add an explicit value check in `checkCommonInputs` to reject non-positive `scales` (at minimum zero, ideally strictly positive) before the kernels run, using the existing `scales` validation block alongside `checkInt32Tensor` and `checkLayerAttnShape`. Make sure the error message clearly names `scales` and the positivity requirement so callers of the `deepseekV4BlockTable` op get an early, safe failure instead of undefined behavior.tests/unittest/_torch/custom_ops/test_deepseek_v4_block_table.py (1)
151-219: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGood baseline coverage, but scratch-capacity edge cases are untested.
Both parametrizations always set
scratch_capacity = copy_idx.shape[0](i.e., scratch capacity == num_tables), andnum_contexts = scratch_capacity - 1. This exercises the "scratch fully covers all tables, almost-all-contexts-active" case but not:
scratch_capacity < num_tables(tables wheretableId >= scratchCapacityshould fall back to the base page index — thecanUseScratch = tableId < scratchCapacitybranch in both scratch kernels).num_contexts == scratch_capacity(boundary where all context slots are active) ornum_contexts == 0.scratch_capacity == 0(should degrade to plain base-page behavior).Consider adding a parametrized case (e.g. in
test_deepseek_v4_block_table.py) coveringscratch_capacity in [0, num_tables // 2, num_tables]andnum_contexts in [0, scratch_capacity]to close this gap. Current coverage is otherwise sufficient for the two kernel-selection paths (element-linear vs. row-based, exercised viamax_blocks in [17, 256]).🤖 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/custom_ops/test_deepseek_v4_block_table.py` around lines 151 - 219, Add coverage for scratch-capacity boundary cases in the DeepSeek V4 block-table tests: the current `test_deepseek_v4_compute_sliding_block_tables_with_scratch_matches_reference` only exercises `scratch_capacity == num_tables` and `num_contexts == scratch_capacity - 1`. Extend this test (or add a parametrized variant) to cover `scratch_capacity < copy_idx.shape[0]`, `scratch_capacity == 0`, and `num_contexts` at both `0` and `scratch_capacity`, so the `canUseScratch = tableId < scratchCapacity` path and the zero/fully-active context cases are validated against `_reference_sliding_block_tables_with_scratch` and `torch.ops.trtllm.deepseek_v4_compute_sliding_block_tables_with_scratch`.Source: Path instructions
🤖 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/sparse/dsa.py`:
- Around line 58-59: The module-level initialization of
_INDEXER_MQA_LOGITS_ELEM_BUDGET in dsa.py can crash on import if
TLLM_INDEXER_MQA_LOGITS_ELEM_BUDGET is malformed. Move the parsing into a safe
path in the same initialization block, validate the environment value, and fall
back to the default budget when parsing fails; optionally emit a warning instead
of raising so the sparse attention backend can still load.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cu`:
- Around line 238-309: The end of computeSlidingBlockTablesWithScratchRowsKernel
reimplements the same scale/offset and bad-page handling already centralized in
applyScaleAndOffset. Replace the inline base-to-output mapping in the
non-scratch path with a call to applyScaleAndOffset so the kernel uses the
shared helper consistently and avoids duplicated logic drifting apart.
In `@cpp/tensorrt_llm/kernels/deepseekV4BlockTable.h`:
- Around line 29-39: Add Doxygen-style `//!` comments for the public prototypes
`invokeDeepseekV4ComputeSlidingBlockTables` and
`invokeDeepseekV4ComputeSlidingBlockTablesWithScratch` in
`deepseekV4BlockTable.h`, documenting each parameter’s meaning, expected
tensor/layout shapes, and any scratch-related semantics. Keep the comments
directly above the declarations and make sure they follow the project’s Doxygen
rules for C++ function prototypes.
In `@cpp/tensorrt_llm/thop/deepseekV4BlockTableOp.cpp`:
- Around line 60-99: `checkCommonInputs` currently validates `scales`’s dtype
and shape but not its value range, while the CUDA kernels it feeds divide and
mod by `scale`. Add an explicit value check in `checkCommonInputs` to reject
non-positive `scales` (at minimum zero, ideally strictly positive) before the
kernels run, using the existing `scales` validation block alongside
`checkInt32Tensor` and `checkLayerAttnShape`. Make sure the error message
clearly names `scales` and the positivity requirement so callers of the
`deepseekV4BlockTable` op get an early, safe failure instead of undefined
behavior.
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py`:
- Around line 554-558: Remove the unused staging tensor allocation in the cache
manager setup: `self._device_block_positions` is no longer referenced anywhere
in `CacheManager` after the removed Python helpers, and the CUDA block-table ops
now compute positions internally. Delete the `torch.arange(...)` assignment in
the initialization path and make sure no remaining methods in `cache_manager.py`
still depend on `self._device_block_positions`.
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py`:
- Around line 470-479: The helper select_indexer_chunk_size in dsa.py looks
internal and should not be part of the module API unless it is intentionally
imported elsewhere. Rename it to _select_indexer_chunk_size, and update any
local references or tests that call it so the sparse attention backend keeps the
helper clearly private.
In `@tests/unittest/_torch/custom_ops/test_deepseek_v4_block_table.py`:
- Around line 151-219: Add coverage for scratch-capacity boundary cases in the
DeepSeek V4 block-table tests: the current
`test_deepseek_v4_compute_sliding_block_tables_with_scratch_matches_reference`
only exercises `scratch_capacity == num_tables` and `num_contexts ==
scratch_capacity - 1`. Extend this test (or add a parametrized variant) to cover
`scratch_capacity < copy_idx.shape[0]`, `scratch_capacity == 0`, and
`num_contexts` at both `0` and `scratch_capacity`, so the `canUseScratch =
tableId < scratchCapacity` path and the zero/fully-active context cases are
validated against `_reference_sliding_block_tables_with_scratch` and
`torch.ops.trtllm.deepseek_v4_compute_sliding_block_tables_with_scratch`.
🪄 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: 45bd3805-00cf-4b0d-bbba-19f44b2f83b9
📒 Files selected for processing (12)
cpp/tensorrt_llm/kernels/deepseekV4BlockTable.cucpp/tensorrt_llm/kernels/deepseekV4BlockTable.hcpp/tensorrt_llm/thop/CMakeLists.txtcpp/tensorrt_llm/thop/deepseekV4BlockTableOp.cpptensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/dsa.pytensorrt_llm/llmapi/disagg_utils.pytensorrt_llm/serve/openai_disagg_service.pytensorrt_llm/serve/openai_protocol.pytensorrt_llm/serve/openai_server.pytests/microbenchmarks/dsv4_block_table_perf.pytests/unittest/_torch/custom_ops/test_deepseek_v4_block_table.py
|
PR_Github #57919 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58000 [ run ] triggered by Bot. Commit: |
|
PR_Github #58000 [ run ] completed with state
|
…ces (NVIDIA#15459) Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com> Co-authored-by: Fanrong Li <lfr-0531@users.noreply.github.com> Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
…NVIDIA#15569) Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-authored-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
…t_token_ids (NVIDIA#15698) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Co-authored-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
151fbd5 to
a534c93
Compare
lfr-0531
added
the
api-compatible
|
/bot run --disable-fail-fast |
|
PR_Github #58027 [ run ] triggered by Bot. Commit: |
|
PR_Github #58027 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58161 [ run ] triggered by Bot. Commit: |
|
PR_Github #58161 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58332 [ run ] triggered by Bot. Commit: |
|
PR_Github #58332 [ run ] completed with state |
JunyiXu-nv
left a comment
JunyiXu-nv
left a comment
There was a problem hiding this comment.
LGTM for llm_api part.
juney-nvidia
left a comment
juney-nvidia
left a comment
There was a problem hiding this comment.
Approved from API perspective
9 participants
Summary by CodeRabbit
New Features
Bug Fixes
Description
Port the four remaining optimizations from
feat/deepseek_v4onto the latest GitHubmain:prompt_token_ids.The cherry-picks were reconciled with current-main interfaces, including the current conversation protocol fields and KV cache manager v2 import path. The changes remain opt-in where applicable and add no new dependency.
Source commits:
acbcb93,d6d1b48,dfb0609, andb7a8a8f.Test Coverage
scripts/build_wheel.py --benchmarks --use_ccache --cuda_architectures "90-real;100-real"; the generated wheel was installed successfully.1.3.0rc21package contains both DSA changes and that its builtlibth_common.soregisters both DeepSeek V4 block-table operators.max_blocks=17andmax_blocks=256on an idle B300 GPU (four cases total).python -m py_compilefor all modified Python modules.main.tests/unittest/_torch/custom_ops/test_deepseek_v4_block_table.pyis blocked before test execution by the local FlashInfer/CUTLASS Python mismatch (cutlass.cute.nvgpu.OperandMajorModeis unavailable). The equivalent operator/reference cases above pass by loading the installed custom-op library directly; CI will run the committed pytest coverage in its pinned environment.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.