[None][perf] Add DeepSeek V4 block table custom op#15821
Conversation
|
/bot run |
|
PR_Github #56957 [ run ] triggered by Bot. Commit: |
|
PR_Github #56957 [ run ] completed with state
|
Add CUDA custom ops for DeepSeek V4 sliding block table generation so the sparse cache manager no longer relies on the torch.compile block-table implementations that caused large host-memory pressure during compile. Changes: - Add no-scratch and scratch CUDA kernels plus thop bindings for trtllm::deepseek_v4_compute_sliding_block_tables and trtllm::deepseek_v4_compute_sliding_block_tables_with_scratch. - Use a tiled no-scratch row kernel with int4 vectorization when maxBlocksPerSeq is 4-aligned and scalar handling otherwise, so the no-scratch path always uses the new kernel. - Replace the DeepSeek V4 cache-manager torch.compile block-table functions with direct torch.ops.trtllm custom-op calls. - Add CUDA unit tests that compare both custom ops against the removed torch reference implementation for scalar and vectorized max-block sizes. - Add a standalone microbenchmark script for comparing custom-op latency against eager or torch.compile references across no-scratch and scratch cases. Validation: - python3 -m py_compile tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py tests/unittest/_torch/custom_ops/test_deepseek_v4_block_table.py - git diff --cached --check - CMH build_sqsh job 778723: COMPLETED exit=0. - CMH GPU unit test job 781707: 4 passed in 4.78s. - CMH no-scratch perf sweep job 780480: comparable production-ish cases averaged +0.15% latency delta versus the previous 2D tiled build, with 16/17 cases within +/-5%. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
b8a085c to
45fbc25
Compare
|
Rebased to pick up #15713 |
|
Oh sorry, I see this is the feature branch. I think we need to cherry-pick #15713 into the feature branch first. |
|
Looks like it has already been cherry-picked here, so the existing rebase should resolve the issue: #15818 |
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #57083 [ run ] triggered by Bot. Commit: |
|
PR_Github #57083 [ run ] completed with state
|
|
/bot run |
|
PR_Github #57131 [ run ] triggered by Bot. Commit: |
|
PR_Github #57131 [ run ] completed with state
|
|
/bot run |
|
PR_Github #57380 [ run ] triggered by Bot. Commit: |
|
PR_Github #57380 [ run ] completed with state |
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>
4 participants
Add CUDA custom ops for DeepSeek V4 sliding block table generation so the sparse cache manager no longer relies on the torch.compile block-table implementations that caused large host-memory pressure during compile.
Changes:
Add no-scratch and scratch CUDA kernels plus thop bindings for trtllm::deepseek_v4_compute_sliding_block_tables and trtllm::deepseek_v4_compute_sliding_block_tables_with_scratch.
Use a tiled no-scratch row kernel with int4 vectorization when maxBlocksPerSeq is 4-aligned and scalar handling otherwise, so the no-scratch path always uses the new kernel.
Replace the DeepSeek V4 cache-manager torch.compile block-table functions with direct torch.ops.trtllm custom-op calls.
Add CUDA unit tests that compare both custom ops against the removed torch reference implementation for scalar and vectorized max-block sizes.
Add a standalone microbenchmark script for comparing custom-op latency against eager or torch.compile references across no-scratch and scratch cases.
Validation:
python3 -m py_compile tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py tests/unittest/_torch/custom_ops/test_deepseek_v4_block_table.py
git diff --cached --check
CMH build_sqsh job 778723: COMPLETED exit=0.
CMH GPU unit test job 781707: 4 passed in 4.78s.
CMH no-scratch perf sweep job 780480: comparable production-ish cases averaged +0.15% latency delta versus the previous 2D tiled build, with 16/17 cases within +/-5%.
@coderabbitai summary
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.