Skip to content

fix: prepare DeepSeek-V4 draft sliding tables #17178

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zcxGGmu wants to merge 1 commit into
base: main
Choose a base branch
from
+68 −0
Open

fix: prepare DeepSeek-V4 draft sliding tables #17178

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -681,6 +681,19 @@ def prepare(self):
self.request_ids,
self.num_contexts,
)
# MTP uses a separate draft KV cache manager; prepare its DeepSeek-V4
# sliding-window tables before the base path copies draft block offsets.
draft_kv_cache_manager = getattr(self, "draft_kv_cache_manager", None)
draft_compute_sliding_block_tables = getattr(
draft_kv_cache_manager,
"compute_sliding_block_tables",
None,
)
if draft_compute_sliding_block_tables is not None:
draft_compute_sliding_block_tables(
self.request_ids,
self.num_contexts,
)

TrtllmAttentionMetadata.prepare(self)

Expand Down
55 changes: 55 additions & 0 deletions tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from unittest.mock import MagicMock

import pytest

from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import (
DeepseekV4TrtllmAttentionMetadata,
)
from tensorrt_llm._torch.attention_backend.trtllm import TrtllmAttentionMetadata


def test_prepare_computes_draft_sliding_block_tables_before_base_prepare(monkeypatch):
"""DeepSeek-V4 MTP draft KV managers need their sliding tables prepared.

The base TRT-LLM metadata prepare path copies block offsets from both the
target and draft managers. DeepSeek-V4's copy path consumes precomputed
sliding-window tables, so the draft manager must compute them before the
base prepare reaches copy_batch_block_offsets().
"""
metadata = object.__new__(DeepseekV4TrtllmAttentionMetadata)
metadata.kv_cache_manager = MagicMock()
metadata.draft_kv_cache_manager = MagicMock()
metadata.request_ids = [11, 12, 13]
metadata.num_contexts = 2

def stop_at_base_prepare(self):
raise RuntimeError("base prepare reached")

monkeypatch.setattr(TrtllmAttentionMetadata, "prepare", stop_at_base_prepare)
Comment on lines +26 to +43

@coderabbitai coderabbitai Bot Aug 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -e
printf '%s\n' '--- target file ---'
sed -n '1,180p' tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py
printf '%s\n' '--- relevant guidelines ---'
if [ -f CODING_GUIDELINES.md ]; then
  rg -n -A4 -B4 'Annotate every function|Prefix non-public|pytest|test' CODING_GUIDELINES.md | head -160
fi
printf '%s\n' '--- nearby test-list references ---'
rg -n 'deepseek_v4|test_deepseek_v4_metadata' tests/integration/test_lists tests 2>/dev/null | head -160

Repository: NVIDIA/TensorRT-LLM

Length of output: 27204

🏁 Script executed:

#!/bin/bash
set -e
sed -n '1,180p' tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py
printf '%s\n' '--- guidelines ---'
if [ -f CODING_GUIDELINES.md ]; then
  rg -n -A4 -B4 'Annotate every function|Prefix non-public|pytest|test' CODING_GUIDELINES.md | head -160
fi
printf '%s\n' '--- test-list references ---'
rg -n 'deepseek_v4|test_deepseek_v4_metadata' tests/integration/test_lists tests 2>/dev/null | head -160

Repository: NVIDIA/TensorRT-LLM

Length of output: 26941

🏁 Script executed:

#!/bin/bash
set -e
python3 - <<'PY'
import ast
from pathlib import Path

path = Path("tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py")
tree = ast.parse(path.read_text())
for node in ast.walk(tree):
    if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
        args = [*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs]
        missing_args = [arg.arg for arg in args if arg.annotation is None]
        print(
            f"{node.name}: return_annotation={node.returns is not None}, "
            f"missing_argument_annotations={missing_args}"
        )
        if node.name == "test_prepare_computes_draft_sliding_block_tables_before_base_prepare":
            for child in node.body:
                if isinstance(child, ast.FunctionDef):
                    print(
                        f"nested_helper={child.name}: "
                        f"self_annotated={bool(child.args.args and child.args.args[0].annotation)}, "
                        f"return_annotation={child.returns is not None}"
                    )
PY
printf '%s\n' '--- exact guideline matches ---'
rg -n -i 'annotat(e|ion).*every function|non-public|private names|leading underscore' CODING_GUIDELINES.md
printf '%s\n' '--- tracked-file status and focused diff ---'
git status --short -- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py tests/integration/test_lists
git diff --stat -- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py tests/integration/test_lists
git diff --unified=3 -- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py
printf '%s\n' '--- exact test-list membership ---'
rg -n -F 'unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py' tests/integration/test_lists || true

Repository: NVIDIA/TensorRT-LLM

Length of output: 496

🏁 Script executed:

#!/bin/bash
set -e
python3 - <<'PY'
import ast
from pathlib import Path

path = Path("tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py")
tree = ast.parse(path.read_text())
for node in ast.walk(tree):
    if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
        args = [*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs]
        missing_args = [arg.arg for arg in args if arg.annotation is None]
        print(
            f"{node.name}: return_annotation={node.returns is not None}, "
            f"missing_argument_annotations={missing_args}"
        )
        if node.name == "test_prepare_computes_draft_sliding_block_tables_before_base_prepare":
            for child in node.body:
                if isinstance(child, ast.FunctionDef):
                    print(
                        f"nested_helper={child.name}: "
                        f"self_annotated={bool(child.args.args and child.args.args[0].annotation)}, "
                        f"return_annotation={child.returns is not None}"
                    )
PY
printf '%s\n' '--- exact guideline matches ---'
rg -n -i 'annotat(e|ion).*every function|non-public|private names|leading underscore' CODING_GUIDELINES.md
printf '%s\n' '--- focused status and diff ---'
git status --short -- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py tests/integration/test_lists
git diff --stat -- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py tests/integration/test_lists
git diff --unified=3 -- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py
printf '%s\n' '--- test-list membership ---'
rg -n -F 'unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py' tests/integration/test_lists || true

Repository: NVIDIA/TensorRT-LLM

Length of output: 496

🏁 Script executed:

#!/bin/bash
set -e
printf '%s\n' '--- focused status and diff ---'
git status --short -- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py tests/integration/test_lists
git diff --stat -- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py tests/integration/test_lists
git diff --unified=5 -- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py
printf '%s\n' '--- exact test-list membership ---'
rg -n -F 'unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py' tests/integration/test_lists || true
printf '%s\n' '--- nearby CI list entries ---'
sed -n '68,90p' tests/integration/test_lists/test-db/l0_b200.yml

Repository: NVIDIA/TensorRT-LLM

Length of output: 2045

🏁 Script executed:

#!/bin/bash
set -e
python3 - <<'PY'
import ast
from pathlib import Path

source = Path(
    "tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py"
).read_text()
tree = ast.parse(source)
test = next(
    node for node in tree.body
    if isinstance(node, ast.FunctionDef)
    and node.name == "test_prepare_computes_draft_sliding_block_tables_before_base_prepare"
)
helper = next(node for node in test.body if isinstance(node, ast.FunctionDef))
print("test_function_present:", test.name)
print("helper_present:", helper.name)
print("test_annotation_state:", test.returns is not None, [a.annotation is not None for a in test.args.args])
print("helper_annotation_state:", helper.returns is not None, [a.annotation is not None for a in helper.args.args])

list_path = Path("tests/integration/test_lists/test-db/l0_b200.yml")
list_text = list_path.read_text()
print("parent_directory_covers_file:",
      "  - unittest/_torch/attention\n" in list_text)
print("exact_file_entry:",
      "unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py"
      in list_text)
PY
python3 - <<'PY'
import pytest
print("pytest.MonkeyPatch_available:", hasattr(pytest, "MonkeyPatch"))
print("pytest.MonkeyPatch:", pytest.MonkeyPatch)
PY

Repository: NVIDIA/TensorRT-LLM

Length of output: 540

Annotate the new test functions and mark the local helper as non-public.

Annotate test_prepare_computes_draft_sliding_block_tables_before_base_prepare() and _stop_at_base_prepare() with pytest.MonkeyPatch, TrtllmAttentionMetadata, and None as appropriate. Rename stop_at_base_prepare to _stop_at_base_prepare.

Test coverage: test_prepare_computes_draft_sliding_block_tables_before_base_prepare() was added and is covered by unittest/_torch/attention in tests/integration/test_lists/test-db/l0_b200.yml. Verdict: sufficient.

🤖 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/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py`
around lines 26 - 43, Update
test_prepare_computes_draft_sliding_block_tables_before_base_prepare with
explicit pytest.MonkeyPatch and TrtllmAttentionMetadata parameter annotations
and a None return annotation. Rename the local stop_at_base_prepare helper to
_stop_at_base_prepare, annotate its self parameter with TrtllmAttentionMetadata
and its return type as None, and use the renamed helper in the monkeypatch.

Source: Coding guidelines


with pytest.raises(RuntimeError, match="base prepare reached"):
DeepseekV4TrtllmAttentionMetadata.prepare(metadata)

metadata.kv_cache_manager.compute_sliding_block_tables.assert_called_once_with(
metadata.request_ids,
metadata.num_contexts,
)
metadata.draft_kv_cache_manager.compute_sliding_block_tables.assert_called_once_with(
metadata.request_ids,
metadata.num_contexts,
)
Comment on lines +26 to +55

@coderabbitai coderabbitai Bot Aug 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Locate candidate CI and QA test-list files.
fd -t f . tests/integration/test_lists | sort

# Check whether this test module or DeepSeek-V4 test group is registered.
rg -n -i -C 2 \
  'test_deepseek_v4_metadata|deepseek_v4' \
  tests/integration/test_lists || true

# Locate available CBTS scope evidence.
fd -t f 'cbts_touchmap.sqlite' . || true
fd -t f '*coverage*' . | rg -i 'cbts|touchmap' || true

Repository: NVIDIA/TensorRT-LLM

Length of output: 12813

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- changed files ---'
git diff --name-status

printf '%s\n' '--- test file diff ---'
git diff -- tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py

printf '%s\n' '--- matching test-list entries ---'
rg -n -F \
  'unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py' \
  tests/integration/test_lists || true

printf '%s\n' '--- DeepSeek-V4 test-list entries ---'
rg -n -i -C 1 \
  'unittest/_torch/attention/sparse/deepseek_v4' \
  tests/integration/test_lists || true

printf '%s\n' '--- CBTS artifacts tracked or present ---'
git ls-files | rg -i '(^|/)(cbts_touchmap\.sqlite|.*cbts.*coverage.*|.*coverage.*cbts.*)$' || true
find . -type f \( -name 'cbts_touchmap.sqlite' -o -iname '*cbts*coverage*' -o -iname '*coverage*cbts*' \) -print

Repository: NVIDIA/TensorRT-LLM

Length of output: 2077

Add the test module to tests/integration/test_lists/test-db/l0_b200.yml. test_deepseek_v4_metadata.py is absent from the CI and QA lists. Coverage summary: one test added; coverage verdict: insufficient.

🤖 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/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py`
around lines 26 - 55, Add
tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_metadata.py
to the appropriate CI and QA entries in
tests/integration/test_lists/test-db/l0_b200.yml so
test_prepare_computes_draft_sliding_block_tables_before_base_prepare is included
in automated coverage.

Source: Path instructions