[None][perf] DeepSeek-V4: cut host overhead in generation-phase metadata preparation - #17257
Open
hyukn wants to merge 5 commits into
Open
[None][perf] DeepSeek-V4: cut host overhead in generation-phase metadata preparation#17257hyukn wants to merge 5 commits into
hyukn wants to merge 5 commits into
Conversation
Collaborator
Author
|
/bot run --disable-fail-fast |
…ata preparation `_prepare_inputs` on the DeepSeek-V4 generation worker is pure host work that is rebuilt from scratch every decode iteration. Profiling a disaggregated GB300 run (1 generation server, tp16/ep16, attention DP, MTP-3, batch 64, fp8 KV cache) showed 12.5 ms per iteration spent there, spread over ~1480 ATen dispatches of which only ~5% launch any GPU work: the rest is interpreter, dispatcher and allocator cost building index and block-table metadata. This change removes that work in three ways. 1. Reuse the request->slot mapping within one prepare() step. `IndexMapper.get_copy_index` was walked three times per step -- once by `compute_sliding_block_tables`, once per compression ratio by `copy_batch_compress_block_tables`, and once by `copy_batch_indexer_compress_block_tables` -- always with the same (request_ids, num_contexts, beam_width). The mapping is deterministic in those arguments and returns a view of one shared pinned buffer, so the later walks recomputed a result equal to the one already there. Each walk costs O(num_seqs) dispatches (select + as_strided + fill_ per entry). It is now computed once and memoized for the step, keyed on the request-id list contents rather than its identity because callers reuse and mutate one list object. The memo is reset at the top of every step. 2. Stop pre-filling block-table regions that are immediately overwritten. `copy_batch_sliding_block_tables` and `copy_batch_block_offsets` filled their whole destination with BAD_PAGE_INDEX and then overwrote the live rows. The destinations are sized for the mapper capacity, not the batch, so the fill covered a [layers, sliding_types, capacity, max_blocks] tensor to keep only the padding. Only the regions the following copy does not touch are filled now. The padding still has to carry BAD_PAGE_INDEX: padded CUDA-graph token slots can index those rows through req_idx_per_token. 3. Build the per-iteration index and compressed-KV metadata on the device. Six new CUDA ops replace element-wise ATen chains that ran on the decode critical path: sparse-attention indices, token positions with the cu_seq_lens/req_idx_per_token pair, per-compression-ratio KV lengths, the compressed-token mask, context and generation compressed position ids, and the shared compression block table. The block-table op also removes a host gather of a few hundred KB out of host_kv_cache_block_offsets plus its staging copy, per compression ratio. host_kv_cache_block_offsets is now uploaded at most once per iteration and shared by all consumers instead of once per table. Together these cut the ATen dispatch count in the region from ~1480 to ~1290. Measured on the configuration above, torch profiler, 60 consecutive decode iterations, 4 generation ranks, comparing against the same build with the python call sites unchanged so the kernels are compiled into both arms: _prepare_inputs 11.82 ms -> 8.96 ms (-24.2%) compressed KV metadata 0.756 -> 0.112 ms (-85%) sparse-attention indices 0.682 -> 0.112 ms (-84%) block tables 1.508 -> 0.986 ms (-35%) compressed-token mask 0.250 -> 0.049 ms (-80%) indices conversion 0.155 -> 0.071 ms (-55%) cu_seq_lens 0.086 -> 0.045 ms (-48%) Regions the change does not touch move by -1.09% +/- 4.01% over 15 samples, which sets the noise floor; every region above is 8 to 21 sigma outside it and reproduces on all 4 ranks. Per-rank spread of the region also collapses from 0.111 ms to 0.016 ms. The two host-side items alone (1 and 2), measured separately on a build without the kernels, account for 0.771 ms (-6.2%): block tables 1.526 -> 0.929 ms and indexer K-cache 0.696 -> 0.417 ms. End-to-end serving throughput and latency are unchanged. The generation loop in this configuration is 92-95% GPU-busy, so the overlap scheduler absorbs a host saving of this size; a paired comparison over 6921 identical conversation turns put TTFT, inter-token latency and output speed all within noise. The value here is host headroom on the decode critical path, not a throughput number. Tests: `tests/unittest/_torch/custom_ops/test_deepseek_v4_metadata_ops.py` pins each new op against the python reference it replaces, reproduced verbatim in the test file, and covers padded buffers, per-ratio list mismatches and empty batches. `tests/unittest/_torch/attention/sparse/ test_deepseek_v4_block_table_host_path.py` covers the mapping memo (value equality with the uncached path, recomputation on changed arguments, in-place mutation of a reused request-id list, per-step invalidation) and asserts the partially-filled block-table buffers are bit-identical to the previous fill-everything behaviour, including the no-tail case. Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>
…epSeek-V4 namespace
Two of the kernels added in the previous commit do not depend on any DeepSeek-V4
concept, and naming them `deepseek_v4_*` hid that they are reusable:
- `compute_token_positions` takes only (seq_lens, cached_tokens, batch_size,
num_tokens). It is the device-side form of the generic
`pad(cumsum(seq_lens))` + `repeat_interleave(arange(batch), seq_lens)` +
`searchsorted` triple. The shared DSA metadata path builds exactly this by
hand today, and that path also serves the separate `dsa` algorithm, so there
is already a second consumer.
- `compute_shared_block_table` takes (pool_id, scale) as plain ints and computes
`where(base == BAD_PAGE_INDEX, BAD_PAGE_INDEX, base * scale)` over
`block_offsets[pool_id, copy_idx, 0, :]`. Any KV-cache manager with a
shared-page pool can use it.
Both now live in `kernels/attentionMetadataKernels.{h,cu}` with the prefix
dropped, and the registered ops are `trtllm::compute_token_positions` and
`trtllm::compute_shared_block_table`. No behaviour change: same kernels, same
launch configuration, same schemas apart from the names. The bound guards and
the size ladder of the single-block scan are unchanged.
The remaining five kernels stay under the DeepSeek-V4 name because their
signatures carry the per-compression-ratio concept (per-ratio KV lengths,
compressed-token mask, context/generation compressed position ids) or DeepSeek-V4
specifics (`window_size`, `max_comp_128`, `sparse_mla_topk` for the sparse
indices).
Also extends the op tests, which previously stopped at batch 300 and therefore
never exercised the scan's template ladder: added cases at 511/512/513/2047/
2048/2049/4096 for both scan-based ops, and two cases asserting a batch above the
compile-time bound is rejected rather than silently overrunning shared memory.
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com>
- Apply the repo's clang-format (v16) and ruff / ruff-format (v0.9.4) to the files this branch touches. - Drop a dead `device` local in prepare_for_deepseek_v4_indices: it was only needed by the tensor-building code the fused kernel replaced. The CPU fallback branch further down still binds its own `device` from `token_positions` and is unaffected. No functional change. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-Authored-By: Claude <noreply@anthropic.com>
hyukn
force-pushed
the
feat/dsv4-gen-metadata-host-overhead
branch
from
August 4, 2026 16:58
b7f18b5 to
c04e1a0
Compare
Collaborator
|
PR_Github #63801 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63801 [ run ] completed with state
|
…e lists CI collected `tests/unittest/_torch/attention/sparse/ test_deepseek_v4_block_table_host_path.py` automatically (the B300 list pulls in `unittest/_torch/attention` as a directory) but not `tests/unittest/_torch/custom_ops/test_deepseek_v4_metadata_ops.py`, because the `custom_ops` directory is enumerated file by file. Add it next to the existing `test_deepseek_v4_q_norm.py` entry on B200 and to the DeepSeek-V4 block on B300, so the ops these tests cover are actually exercised pre-merge. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-Authored-By: Claude <noreply@anthropic.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63839 [ run ] triggered by Bot. Commit: |
The release license check requires the complete Apache-2.0 notice on C++/CUDA sources, not the short SPDX-identifier form. My six new/rewritten kernel files used the short form, which produced `errcount: 6` in [Release-Check] Run while the same stage passes on other PRs. Copy the header verbatim from the existing `deepseekV4BlockTable.h` in the same directory. The two new test files keep the short SPDX form, matching the surrounding test sources (for example `test_compressor_module.py`), and were not among the reported errors. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-Authored-By: Claude <noreply@anthropic.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63841 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63839 [ run ] completed with state |
Collaborator
|
PR_Github #63841 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #63909 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63909 [ run ] completed with state |
hyukn
requested review from
JunyiXu-nv,
QiJune,
asfiyab-nvidia,
dpitman-nvda,
schetlur-nv,
xinhe-nv,
yiqingy0,
yuxianq and
zongfeijing
August 5, 2026 05:53
Contributor
WalkthroughAdded CUDA kernels and Torch operators for DeepSeek V4 token positions, indices, compressed metadata, and shared block tables. Updated Python metadata preparation with device buffers and fused operations. Added validation tests and pre-merge test coverage. ChangesDeepSeek V4 metadata pipeline
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant DeepseekV4Metadata
participant TorchOperator
participant CUDAKernels
participant DeviceBuffers
DeepseekV4Metadata->>TorchOperator: request fused metadata computation
TorchOperator->>CUDAKernels: dispatch CUDA kernels
CUDAKernels->>DeviceBuffers: write positions, indices, masks, and block tables
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (6)
cpp/tensorrt_llm/kernels/attentionMetadataKernels.h (1)
37-63: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new public interfaces with Doxygen.
The new public CUDA interfaces use regular
//comments. Use Doxygen comments so generated API documentation includes these contracts.
cpp/tensorrt_llm/kernels/attentionMetadataKernels.h#L37-L63: Convert the launch-interface documentation to Doxygen comments.cpp/tensorrt_llm/kernels/deepseekV4Indices.h#L29-L35: Convert the launch-interface documentation to Doxygen comments.cpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.h#L29-L79: Add Doxygen documentation for the public parameter structures and launch interfaces.As per coding guidelines, “document new interfaces with Doxygen.”
🤖 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/attentionMetadataKernels.h` around lines 37 - 63, Convert the launch-interface comments in cpp/tensorrt_llm/kernels/attentionMetadataKernels.h lines 37-63 and cpp/tensorrt_llm/kernels/deepseekV4Indices.h lines 29-35 from regular comments to Doxygen comments, preserving their documented contracts. Add Doxygen documentation for the public parameter structures and launch interfaces in cpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.h lines 29-79, covering the visible public symbols and their parameters without changing behavior.Source: Coding guidelines
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py (3)
591-591: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the new signature.
prepare_for_deepseek_v4_indicesdeclarestoken_positions=Nonewith no type and no return type. The coding guidelines require an annotation on every function andNonefor procedures.♻️ Suggested change
- def prepare_for_deepseek_v4_indices(self, token_positions=None): + def prepare_for_deepseek_v4_indices( + self, token_positions: Optional[torch.Tensor] = None + ) -> None:As per coding guidelines: "Annotate every function, use
Nonefor procedures".🤖 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/deepseek_v4.py` at line 591, Annotate the prepare_for_deepseek_v4_indices signature with an appropriate type for token_positions and an explicit None return type, preserving its existing behavior.Source: Coding guidelines
923-934: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
gen_output_offsetsforces a device-to-host sync per ratio.Line 920 calls
.item()onself.cu_new_comp_kv_cuda[r][num_contexts]for every ratio whennum_contexts > 0. Each call is an implicit D2H read plus a stream synchronization on the critical path. The comment at lines 794-796 states that this pattern is what the host-sidectx_output_sizesexists to avoid.
ctx_output_sizes[r]already holdscu_new_comp_kv[num_contexts]for ratioras a Python int. Whenctx_output_sizesis notNone, reuse it instead.on_update_kv_lens()can passNone, so keep the.item()path as the fallback.The
.item()line itself is outside the changed range. The new fused call consumes its result, so the fix belongs with this block.♻️ Suggested change
gen_output_offsets = { - r: self.cu_new_comp_kv_cuda[r][num_contexts].item() if num_contexts > 0 else 0 + r: ( + 0 + if num_contexts == 0 + else ( + int(ctx_output_sizes[r]) + if ctx_output_sizes is not None + else self.cu_new_comp_kv_cuda[r][num_contexts].item() + ) + ) for r in self._compress_ratios_sorted }🤖 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/deepseek_v4.py` around lines 923 - 934, Update the gen_output_offsets construction immediately before deepseek_v4_compute_gen_compressed_position_ids to reuse ctx_output_sizes[r] when ctx_output_sizes is available, avoiding the per-ratio .item() device-to-host synchronization. Preserve the existing cu_new_comp_kv_cuda[r][num_contexts].item() fallback when ctx_output_sizes is None, and keep the resulting offsets passed unchanged to the fused operation.
581-589: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winRemove the unused host
cu_seq_lensbuffer and staging. TheTruecall computescu_seq_lens_cuda, and the laterFalsecall reuses it. Remove the host cumsum, H2D copy, and host-buffer 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/deepseek_v4.py` around lines 581 - 589, Remove the host cu_seq_lens allocation, host cumulative-sum staging, and H2D copy from the token-position preparation flow around compute_token_positions. Keep the True invocation to compute cu_seq_lens_cuda, and preserve the later False invocation’s reuse of that device buffer without referencing the removed host buffer.cpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.cu (1)
138-152: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winClamp
reqIdxin the context kernel, or validatecounts[i]againstcu[numContexts].
upperBoundRequest(cu, numContexts, t)can returnnumContexts. This happens whent >= cu[numContexts]. The kernel then readspastKv[numContexts], which is one element past the validatedbatchSizeregion for the context entry point (batchSize == numContextsthere).The current callers pass
counts[i] == cu[numContexts], so the branch is unreachable today. The operator layer only checkscount >= 0, so nothing enforces that invariant.computeGenCompressedPositionIdsKernelalready appliesmin(max(...))for the same reason.♻️ Suggested defensive clamp
for (int32_t t = blockIdx.x * blockDim.x + threadIdx.x; t < total; t += gridDim.x * blockDim.x) { - int32_t const reqIdx = upperBoundRequest(cu, numContexts, t); + int32_t const reqIdx = min(upperBoundRequest(cu, numContexts, t), numContexts - 1); out[t] = (pastKv[reqIdx] + (t - cu[reqIdx])) * ratio; }🤖 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/deepseekV4CompressedMeta.cu` around lines 138 - 152, Prevent out-of-bounds access in computeCtxCompressedPositionIdsKernel by defensively clamping the upperBoundRequest result to the valid context range before indexing pastKv. Preserve the existing position calculation while ensuring reqIdx never reaches numContexts, or alternatively validate counts against cu[numContexts] before launching the kernel.tests/unittest/_torch/custom_ops/test_deepseek_v4_metadata_ops.py (1)
219-223: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe zero-count branches assert nothing.
test_compute_compressed_maskskips the ratio whentotals[r] == 0.test_compute_ctx_compressed_position_ids(line 246) andtest_compute_gen_compressed_position_ids(line 292) do the same. The PR description lists empty batches as covered, but thesecontinuestatements mean the zero-count case only checks that the op does not raise.The launch helpers return early when
maxTotalTokens <= 0ormaxCount <= 0. That early return is worth pinning: assert that the output buffer is unchanged instead of skipping.♻️ Suggested change for the mask test
for r in ratios: if totals[r] == 0: + # The launch helper returns early; the buffer must stay untouched. + assert not masks[r].any() continue🤖 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_metadata_ops.py` around lines 219 - 223, Update test_compute_compressed_mask, test_compute_ctx_compressed_position_ids, and test_compute_gen_compressed_position_ids so their totals[r] == 0 branches assert that the corresponding output buffer remains unchanged, matching the launch helpers’ early-return behavior, rather than continuing without assertions. Preserve the existing expected-value comparisons for nonzero counts.
🤖 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/deepseek_v4/cache_manager.py`:
- Around line 1120-1127: Initialize _block_offsets_uploaded, _copy_idx_memo_key,
and _copy_idx_memo_val in the constructor path _prepare_page_table_tensor with
their appropriate initial values, then replace the getattr defaults in
_ensure_device_block_offsets and the related memo-copy logic with direct
attribute access. Make the per-step ordering requirement explicit by asserting
that the step was initialized before _ensure_device_block_offsets proceeds,
while preserving compute_sliding_block_tables as the reset point.
- Around line 1372-1392: Update the shared block-table construction flow around
_compute_shared_block_table_device and _compute_shared_block_table so
host_indexer_k_cache_block_offsets remains synchronized when device_block_table
is provided. Populate the host table as well, or consistently redirect DSA and
speculative replay consumers to the device table, including MTP expansion and
slot-mapping conversion.
In
`@tests/unittest/_torch/attention/sparse/test_deepseek_v4_block_table_host_path.py`:
- Around line 123-182: The three padding tests must invoke the production
methods instead of duplicating their logic inline. Update
test_sliding_block_table_padding_keeps_bad_page_index,
test_sliding_block_table_padding_when_full, and
test_block_offsets_padding_keeps_bad_page_index to construct a
DeepseekV4CacheManager.__new__ instance with _num_tables and
_precomputed_sliding_block_tables configured as _manager_with_mapper does, pass
CUDA tensors to copy_batch_sliding_block_tables or copy_batch_block_offsets, and
apply the existing CUDA skip marker.
---
Nitpick comments:
In `@cpp/tensorrt_llm/kernels/attentionMetadataKernels.h`:
- Around line 37-63: Convert the launch-interface comments in
cpp/tensorrt_llm/kernels/attentionMetadataKernels.h lines 37-63 and
cpp/tensorrt_llm/kernels/deepseekV4Indices.h lines 29-35 from regular comments
to Doxygen comments, preserving their documented contracts. Add Doxygen
documentation for the public parameter structures and launch interfaces in
cpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.h lines 29-79, covering the
visible public symbols and their parameters without changing behavior.
In `@cpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.cu`:
- Around line 138-152: Prevent out-of-bounds access in
computeCtxCompressedPositionIdsKernel by defensively clamping the
upperBoundRequest result to the valid context range before indexing pastKv.
Preserve the existing position calculation while ensuring reqIdx never reaches
numContexts, or alternatively validate counts against cu[numContexts] before
launching the kernel.
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py`:
- Line 591: Annotate the prepare_for_deepseek_v4_indices signature with an
appropriate type for token_positions and an explicit None return type,
preserving its existing behavior.
- Around line 923-934: Update the gen_output_offsets construction immediately
before deepseek_v4_compute_gen_compressed_position_ids to reuse
ctx_output_sizes[r] when ctx_output_sizes is available, avoiding the per-ratio
.item() device-to-host synchronization. Preserve the existing
cu_new_comp_kv_cuda[r][num_contexts].item() fallback when ctx_output_sizes is
None, and keep the resulting offsets passed unchanged to the fused operation.
- Around line 581-589: Remove the host cu_seq_lens allocation, host
cumulative-sum staging, and H2D copy from the token-position preparation flow
around compute_token_positions. Keep the True invocation to compute
cu_seq_lens_cuda, and preserve the later False invocation’s reuse of that device
buffer without referencing the removed host buffer.
In `@tests/unittest/_torch/custom_ops/test_deepseek_v4_metadata_ops.py`:
- Around line 219-223: Update test_compute_compressed_mask,
test_compute_ctx_compressed_position_ids, and
test_compute_gen_compressed_position_ids so their totals[r] == 0 branches assert
that the corresponding output buffer remains unchanged, matching the launch
helpers’ early-return behavior, rather than continuing without assertions.
Preserve the existing expected-value comparisons for nonzero counts.
🪄 Autofix
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: 867ba1a2-6123-4542-9b12-6c656c819eab
📒 Files selected for processing (14)
cpp/tensorrt_llm/kernels/attentionMetadataKernels.cucpp/tensorrt_llm/kernels/attentionMetadataKernels.hcpp/tensorrt_llm/kernels/deepseekV4BlockTable.cucpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.cucpp/tensorrt_llm/kernels/deepseekV4CompressedMeta.hcpp/tensorrt_llm/kernels/deepseekV4Indices.cucpp/tensorrt_llm/kernels/deepseekV4Indices.hcpp/tensorrt_llm/thop/deepseekV4BlockTableOp.cpptensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_b300.ymltests/unittest/_torch/attention/sparse/test_deepseek_v4_block_table_host_path.pytests/unittest/_torch/custom_ops/test_deepseek_v4_metadata_ops.py
Comment on lines
+1120
to
+1127
| def _ensure_device_block_offsets(self) -> None: | ||
| """Upload host_kv_cache_block_offsets to the device at most once per iteration.""" | ||
| if getattr(self, "_block_offsets_uploaded", False): | ||
| return | ||
| self._device_kv_cache_block_offsets_input.copy_( | ||
| self.host_kv_cache_block_offsets, non_blocking=True | ||
| ) | ||
| self._block_offsets_uploaded = True |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Initialize _block_offsets_uploaded in the constructor and make the ordering requirement explicit.
Two problems share one cause here. _block_offsets_uploaded is never declared. _ensure_device_block_offsets reads it through getattr(self, ..., False), and only compute_sliding_block_tables resets it to False (line 1230). The same pattern applies to _copy_idx_memo_key and _copy_idx_memo_val at lines 1179-1193.
The correctness of every consumer therefore depends on compute_sliding_block_tables running first in each step. DeepseekV4TrtllmAttentionMetadata.prepare does call it first, so the current code is correct. If any future path builds a compress or indexer table without that call, _ensure_device_block_offsets returns immediately and the kernels read block offsets from the previous iteration. The failure is silent wrong block tables, not a crash.
Declare all three attributes in _prepare_page_table_tensor and drop the getattr defaults. Consider asserting in _ensure_device_block_offsets that the step has been started.
♻️ Suggested change
# in _prepare_page_table_tensor, next to the other device buffers
+ self._block_offsets_uploaded = False
+ self._copy_idx_memo_key = None
+ self._copy_idx_memo_val = None def _ensure_device_block_offsets(self) -> None:
"""Upload host_kv_cache_block_offsets to the device at most once per iteration."""
- if getattr(self, "_block_offsets_uploaded", False):
+ if self._block_offsets_uploaded:
return- if getattr(self, "_copy_idx_memo_key", None) == key:
+ if self._copy_idx_memo_key == key:
return self._copy_idx_memo_valAs per coding guidelines: "initialize externally visible class members in the constructor".
🤖 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 1120 - 1127, Initialize _block_offsets_uploaded,
_copy_idx_memo_key, and _copy_idx_memo_val in the constructor path
_prepare_page_table_tensor with their appropriate initial values, then replace
the getattr defaults in _ensure_device_block_offsets and the related memo-copy
logic with direct attribute access. Make the per-step ordering requirement
explicit by asserting that the step was initialized before
_ensure_device_block_offsets proceeds, while preserving
compute_sliding_block_tables as the reset point.
Source: Coding guidelines
Comment on lines
+1372
to
+1392
| device_block_table: Optional[torch.Tensor] = None, | ||
| ) -> None: | ||
| """Build the shared INDEXER_COMPRESS compatibility block table.""" | ||
| """Build the shared INDEXER_COMPRESS compatibility block table. | ||
|
|
||
| When `device_block_table` is given the table is produced directly on the | ||
| device and `host_block_table` is left untouched. | ||
| """ | ||
| assert beam_width == 1, "DSV4 only supports beam width 1 now" | ||
| copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) | ||
| copy_idx = self._get_copy_index_cached(request_ids, num_contexts, beam_width) | ||
| pool_id = self._csa_indexer_compress_pool_id | ||
| scale = self._csa_indexer_compress_scale | ||
| if pool_id is None or scale is None: | ||
| raise RuntimeError("Missing INDEXER_COMPRESS pool metadata") | ||
| host_block_table[:num_seqs] = self._compute_shared_block_table(pool_id, scale, copy_idx) | ||
| if device_block_table is not None: | ||
| # Same rationale as copy_batch_compress_block_tables: build on device | ||
| # and skip both the host gather and the staging copy. | ||
| self._compute_shared_block_table_device( | ||
| pool_id, scale, copy_idx, device_block_table[:num_seqs] | ||
| ) | ||
| else: | ||
| host_block_table[:num_seqs] = self._compute_shared_block_table(pool_id, scale, copy_idx) |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find every reader of the host indexer block-offset buffer.
rg -nP -C4 'host_indexer_k_cache_block_offsets' --type=py
# Find any other caller of the table builder.
rg -nP -C4 'copy_batch_indexer_compress_block_tables\s*\(' --type=pyRepository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- target file symbols ---'
ast-grep outline tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py
printf '%s\n' '--- identifier references across the repository ---'
rg -n -C3 'host_indexer_k_cache_block_offsets|copy_batch_indexer_compress_block_tables|prepare_for_indexer_k_cache' .
printf '%s\n' '--- target implementation ---'
sed -n '1325,1410p' tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py
printf '%s\n' '--- metadata caller context ---'
rg -l 'class DeepseekV4TrtllmAttentionMetadata|prepare_for_indexer_k_cache' tensorrt_llmRepository: NVIDIA/TensorRT-LLM
Length of output: 29245
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- DeepseekV4 metadata caller ---'
sed -n '490,555p' tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py
printf '%s\n' '--- DSA base preparation and host-table consumers ---'
sed -n '1035,1080p' tensorrt_llm/_torch/attention_backend/sparse/dsa.py
sed -n '1495,1530p' tensorrt_llm/_torch/attention_backend/sparse/dsa.py
sed -n '1565,1600p' tensorrt_llm/_torch/attention_backend/sparse/dsa.py
sed -n '2095,2130p' tensorrt_llm/_torch/attention_backend/sparse/dsa.py
sed -n '2208,2235p' tensorrt_llm/_torch/attention_backend/sparse/dsa.py
printf '%s\n' '--- DeepseekV4 class declaration and table attributes ---'
rg -n -C5 'class DeepseekV4|indexer_k_cache_block_offsets|host_indexer_k_cache_block_offsets' \
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py \
tensorrt_llm/_torch/attention_backend/sparse/dsa.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 33430
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- metadata preparation order ---'
sed -n '700,765p' tensorrt_llm/_torch/attention_backend/sparse/dsa.py
rg -n -C5 'prepare_for_(indexer_k_cache|spec_decode|indices_conversion)|create_buffers_for_indexer|prepare_for_mtp' \
tensorrt_llm/_torch/attention_backend/sparse/dsa.py \
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py
printf '%s\n' '--- slot-mapping method declarations and callers ---'
rg -n -C4 'def (prepare_for_indices_conversion|prepare_for_spec_decode|prepare_for_block_tables)|_compute_slot_mappings\(' \
tensorrt_llm/_torch/attention_backend/sparse/dsa.py \
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py
printf '%s\n' '--- speculative state handling ---'
sed -n '135,225p' tensorrt_llm/_torch/speculative/interface.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 33002
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
from pathlib import Path
import re
caller = Path("tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py").read_text()
manager = Path("tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py").read_text()
dsa = Path("tensorrt_llm/_torch/attention_backend/sparse/dsa.py").read_text()
spec = Path("tensorrt_llm/_torch/speculative/interface.py").read_text()
call = re.search(
r"copy_batch_indexer_compress_block_tables\((.*?)\n\s*\)",
caller,
re.S,
)
assert call and "device_block_table=self.indexer_k_cache_block_offsets" in call.group(1)
assert "self.host_indexer_k_cache_block_offsets" in call.group(1)
method = re.search(
r"def copy_batch_indexer_compress_block_tables\(.*?(?=\n @|\n def )",
manager,
re.S,
)
assert method and "host_block_table[:num_seqs] = " in method.group(0)
assert "if device_block_table is not None:" in method.group(0)
required_readers = {
"MTP expansion": "self.host_indexer_k_cache_block_offsets[" in dsa,
"context slot mapping": "metadata.host_indexer_k_cache_block_offsets,\n" in dsa,
"full-KV slot mapping": dsa.count("metadata.host_indexer_k_cache_block_offsets,") >= 2,
"speculative replay": "m.host_indexer_k_cache_block_offsets[:m.num_seqs" in spec,
}
for name, present in required_readers.items():
print(f"{name}: {'present' if present else 'absent'}")
assert present
print("DeepseekV4 passes device_block_table and the cache-manager device branch skips the host write.")
print("The host table still has DSA and speculative consumers.")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 419
Keep host_indexer_k_cache_block_offsets synchronized. The device path skips the host write, but DSA and speculative replay still read this table for MTP expansion and slot-mapping conversion. Populate the host table or update all consumers to use the device table.
🤖 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 1372 - 1392, Update the shared block-table construction flow around
_compute_shared_block_table_device and _compute_shared_block_table so
host_indexer_k_cache_block_offsets remains synchronized when device_block_table
is provided. Populate the host table as well, or consistently redirect DSA and
speculative replay consumers to the device table, including MTP expansion and
slot-mapping conversion.
Comment on lines
+123
to
+182
| def test_sliding_block_table_padding_keeps_bad_page_index(): | ||
| """Only [:_num_tables] is overwritten; the tail must stay BAD_PAGE_INDEX.""" | ||
| layers, types, capacity, max_blocks, num_tables = 2, 5, 9, 4, 3 | ||
| src = torch.arange(layers * types * num_tables * max_blocks, dtype=torch.int32).reshape( | ||
| layers, types, num_tables, max_blocks | ||
| ) | ||
|
|
||
| # Reference: the original code filled the whole tensor first. | ||
| reference = torch.empty((layers, types, capacity, max_blocks), dtype=torch.int32) | ||
| reference.fill_(BAD_PAGE_INDEX) | ||
| reference[:, :, :num_tables, :].copy_(src) | ||
|
|
||
| # Optimized: pre-poison everything, fill only the tail, then copy the head. | ||
| optimized = torch.full((layers, types, capacity, max_blocks), -999999, dtype=torch.int32) | ||
| if num_tables < optimized.size(2): | ||
| optimized[:, :, num_tables:, :].fill_(BAD_PAGE_INDEX) | ||
| optimized[:, :, :num_tables, :].copy_(src) | ||
|
|
||
| torch.testing.assert_close(optimized, reference, rtol=0, atol=0) | ||
| assert not (optimized == -999999).any(), "head was not fully overwritten" | ||
|
|
||
|
|
||
| def test_sliding_block_table_padding_when_full(): | ||
| """With num_tables == capacity there is no tail to fill.""" | ||
| layers, types, capacity, max_blocks = 1, 2, 6, 3 | ||
| src = torch.arange(layers * types * capacity * max_blocks, dtype=torch.int32).reshape( | ||
| layers, types, capacity, max_blocks | ||
| ) | ||
|
|
||
| reference = torch.empty((layers, types, capacity, max_blocks), dtype=torch.int32) | ||
| reference.fill_(BAD_PAGE_INDEX) | ||
| reference[:, :, :capacity, :].copy_(src) | ||
|
|
||
| optimized = torch.full((layers, types, capacity, max_blocks), -999999, dtype=torch.int32) | ||
| if capacity < optimized.size(2): # false: guard must skip the fill | ||
| optimized[:, :, capacity:, :].fill_(BAD_PAGE_INDEX) | ||
| optimized[:, :, :capacity, :].copy_(src) | ||
|
|
||
| torch.testing.assert_close(optimized, reference, rtol=0, atol=0) | ||
|
|
||
|
|
||
| def test_block_offsets_padding_keeps_bad_page_index(): | ||
| """copy_batch_block_offsets writes only beam 0 of [:_num_tables].""" | ||
| layers, capacity, beams, max_blocks, num_tables = 2, 7, 2, 3, 4 | ||
| src = torch.arange(layers * num_tables * max_blocks, dtype=torch.int32).reshape( | ||
| layers, num_tables, max_blocks | ||
| ) | ||
|
|
||
| reference = torch.empty((layers, capacity, beams, max_blocks), dtype=torch.int32) | ||
| reference.fill_(BAD_PAGE_INDEX) | ||
| reference[:, :num_tables, 0, :].copy_(src) | ||
|
|
||
| optimized = torch.full((layers, capacity, beams, max_blocks), -999999, dtype=torch.int32) | ||
| if num_tables < optimized.size(1): | ||
| optimized[:, num_tables:, :, :].fill_(BAD_PAGE_INDEX) | ||
| optimized[:, :num_tables, 1:, :].fill_(BAD_PAGE_INDEX) | ||
| optimized[:, :num_tables, 0, :].copy_(src) | ||
|
|
||
| torch.testing.assert_close(optimized, reference, rtol=0, atol=0) | ||
| assert not (optimized == -999999).any(), "some region was left unwritten" |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
These three tests never call the code they claim to cover.
test_sliding_block_table_padding_keeps_bad_page_index, test_sliding_block_table_padding_when_full, and test_block_offsets_padding_keeps_bad_page_index reimplement the fill-and-copy logic inline and compare it against a local reference. Neither copy_batch_sliding_block_tables nor copy_batch_block_offsets is imported or invoked.
The tests prove that the arithmetic in the test file is self-consistent. They cannot detect a regression in cache_manager.py. If someone removes the dst_tensor[:, :_num_tables, 1:, :].fill_(BAD_PAGE_INDEX) line at cache_manager.py line 1291, every test here still passes. The file docstring states these cover the partial-overwrite behaviour, so the gap is easy to miss.
Call the real methods on a DeepseekV4CacheManager.__new__ instance with _num_tables and _precomputed_sliding_block_tables injected, the same pattern _manager_with_mapper already uses. Both methods assert a CUDA destination, so gate these three with the CUDA skip marker.
🤖 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/test_deepseek_v4_block_table_host_path.py`
around lines 123 - 182, The three padding tests must invoke the production
methods instead of duplicating their logic inline. Update
test_sliding_block_table_padding_keeps_bad_page_index,
test_sliding_block_table_padding_when_full, and
test_block_offsets_padding_keeps_bad_page_index to construct a
DeepseekV4CacheManager.__new__ instance with _num_tables and
_precomputed_sliding_block_tables configured as _manager_with_mapper does, pass
CUDA tensors to copy_batch_sliding_block_tables or copy_batch_block_offsets, and
apply the existing CUDA skip marker.
Collaborator
Author
|
/bot run --disable-fail-fast --add-multi-gpu-test |
Collaborator
|
PR_Github #64006 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64006 [ run ] completed with state
|
xinhe-nv
approved these changes
Aug 6, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
_prepare_inputson the DeepSeek-V4 generation worker is pure host work that is rebuilt from scratch on every decode iteration. Profiling a disaggregated GB300 run showed 12.5 ms per iteration spent there, spread over ~1480 ATen dispatches of which only ~5% launch any GPU work — the rest is interpreter, dispatcher and allocator cost building index and block-table metadata.This PR removes that work in three ways and brings
_prepare_inputsto 8.96 ms (-24.2%).What changed
1. Reuse the request→slot mapping within one
prepare()step.IndexMapper.get_copy_indexwas walked three times per step — bycompute_sliding_block_tables, once per compression ratio bycopy_batch_compress_block_tables, and bycopy_batch_indexer_compress_block_tables— always with the same(request_ids, num_contexts, beam_width). The mapping is deterministic in those arguments and returns a view of one shared pinned buffer, so the later walks recomputed a result equal to the one already there. Each walk costs O(num_seqs) dispatches (select+as_strided+fill_per entry). It is now computed once and memoized for the step, keyed on the request-id list contents rather than its identity, because callers reuse and mutate a single list object. The memo is reset at the top of every step.2. Stop pre-filling block-table regions that are immediately overwritten.
copy_batch_sliding_block_tablesandcopy_batch_block_offsetsfilled their whole destination withBAD_PAGE_INDEXand then overwrote the live rows. The destinations are sized for the mapper capacity, not the batch, so the fill covered a[layers, sliding_types, capacity, max_blocks]tensor in order to keep only the padding. Only the regions the following copy does not touch are filled now. The padding still has to carryBAD_PAGE_INDEX— padded CUDA-graph token slots can index those rows throughreq_idx_per_token, so this is load-bearing, not dead.3. Build the per-iteration index and compressed-KV metadata on the device.
Six new CUDA ops replace element-wise ATen chains on the decode critical path: sparse-attention indices, token positions with the
cu_seq_lens/req_idx_per_tokenpair, per-compression-ratio KV lengths, the compressed-token mask, context and generation compressed position ids, and the shared compression block table. The block-table op also removes a host gather of a few hundred KB out ofhost_kv_cache_block_offsetsplus its staging copy, per compression ratio; that buffer is now uploaded at most once per iteration and shared by all consumers instead of once per table. Together these cut the dispatch count in the region from ~1480 to ~1290.Performance
Config: disaggregated serving, 1 generation server, tp16/ep16, attention DP, MTP-3, max batch 64, fp8 KV cache, GB300. torch profiler, 60 consecutive decode iterations, 4 generation ranks. Both arms use the same build with the kernels compiled in, differing only in whether the python call sites use them, so the comparison isolates this change.
_prepare_inputscu_seq_lensRegions this change does not touch move by -1.09% ± 4.01% over 15 samples, which sets the noise floor; every region above is 8–21σ outside it and reproduces on all 4 ranks. Per-rank spread of the region also collapses from 0.111 ms to 0.016 ms.
The two host-side items alone (1 and 2) were measured separately on a build without the kernels: 0.771 ms (-6.2%) — block tables 1.526 → 0.929 ms, indexer K-cache 0.696 → 0.417 ms. So they stand on their own if the kernels are ever disabled.
Two of the eight regions targeted showed no effect and are reported as such:
cached_token_lensH2D (+1%) and the per-ratio host scalars (-0.3%, already only 0.016 ms).On end-to-end numbers
End-to-end serving throughput and latency are unchanged, and I would not claim otherwise. The generation loop in this configuration is 92–95% GPU-busy, so the overlap scheduler absorbs a host saving of this size. A paired comparison over 6921 identical conversation turns (same prompt, same output length, matched within 2%/5%) put TTFT, inter-token latency and output speed all at ~50/50 — no detectable effect in either direction. An unpaired comparison appeared to show -2% throughput, but that was a load-mismatch artifact: the two servers held different request mixes at the same iteration index (GPU-busy differed by +1.6%, which a host-side python change cannot cause), and the deficit vanishes once turns are paired.
The value of this change is host headroom on the decode critical path, not a throughput number. It matters where the host is the constraint — lower GPU occupancy, larger batches, or after GPU-side work lands.
Testing
tests/unittest/_torch/custom_ops/test_deepseek_v4_metadata_ops.py— pins each new op against the python reference it replaces, with the references reproduced verbatim in the test file so the tests fix the semantics rather than the current implementation. Covers all compression-ratio subsets, batch sizes from 1 to 300, padded output buffers (asserting the padding tail is left untouched), per-ratio list-length mismatches, and empty batches.tests/unittest/_torch/attention/sparse/test_deepseek_v4_block_table_host_path.py— covers the mapping memo (value equality with the uncached path, recomputation when any argument changes, in-place mutation of a reused request-id list, per-step invalidation) and asserts the partially-filled block-table buffers are bit-identical to the previous fill-everything behaviour, including thenum_tables == capacityno-tail case and the beam/row padding ofcopy_batch_block_offsets. These are host-only and need no GPU.The correctness of the kernels was also verified on GB200/GB300 SBSA hardware during development (221 parametrized cases passing).
Notes for reviewers
main. Requesting CI to confirm the build across architectures before de-drafting._compute_compressed_maskand_compute_gen_compressed_position_idsare no longer called fromprepare()now that the ops replace them, but they should not be deleted:tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.pycallsmetadata._compute_compressed_mask(...)directly (pre-existing test, signature unchanged by this PR). Correcting an earlier note in this description that suggested they were dead code.copy_batch_compress_block_tablesis gated ondst_tensor.is_cuda, so the host fallback remains for non-CUDA destinations.🤖 Generated with Claude Code
Update: reuse boundary made explicit
Reviewing which of these kernels are actually DeepSeek-V4-specific, two of them are not, and the
deepseek_v4_prefix was hiding that. They now live inkernels/attentionMetadataKernels.{h,cu}with the prefix dropped:trtllm::compute_token_positions(seq_lens, cached_tokens, batch_size, num_tokens). It is the device-side form ofpad(cumsum(seq_lens))+repeat_interleave(arange(batch), seq_lens)+searchsorted. The shared DSA metadata path builds exactly this by hand today, and that path also serves the separatedsaalgorithm — so a second consumer already exists.trtllm::compute_shared_block_table(pool_id, scale)as plain ints; computeswhere(base == BAD_PAGE_INDEX, BAD_PAGE_INDEX, base * scale)overblock_offsets[pool_id, copy_idx, 0, :]. Any KV manager with a shared-page pool can use it.dsabut not beyondwindow_size,max_comp_128,sparse_mla_topkNo behaviour change in that commit: same kernels, same launch configuration, same schemas apart from the names.
I have not converted the DSA path to call the shared op in this PR — that is a separate change with its own correctness surface, and it should be measured on a prefill-heavy workload rather than inherited from this one. Flagging it as the obvious follow-up.
Note the two python-level items are generic in pattern but not in benefit: the base
KVCacheManagerV2also callsget_copy_index, but only once per invocation and it passescopy_idxdown rather than re-walking, so the 3x redundancy is specific to DeepSeek-V4 having three independent table builders. And the basecopy_batch_block_offsetsalready delegates to a device op, so it never had the fill-everything-then-overwrite pattern.Test coverage added in the same update
The op tests previously stopped at batch 300 and so never exercised the single-block scan's template ladder (
<=512,<=2048, else the 4096 compile-time bound). Added cases at 511/512/513/2047/2048/2049/4096 for both scan-based ops, plus two asserting that a batch above the bound is rejected rather than silently overrunning shared memory (the bound is enforced byTORCH_CHECKin the op layer). Also added coverage for the shared block table, which had none despite being registered and called — including an assertion that rows pastnum_tablesare left untouched.Verification status
Run on GB300 against a build with these kernels compiled in:
test_deepseek_v4_metadata_ops.py: 262 passed, 3 skippedtest_deepseek_v4_block_table_host_path.py: 8 passedThe post-rename suite was re-run and passes identically (the prebuilt image predates the rename, so the new names were aliased onto the old registrations for that run; the rename itself is covered by the schema/impl pairing and by CI's build). Still draft pending a CI build against current
main.Adjacent gap noticed, deliberately left alone
While checking that my own tests get collected I found that
tests/unittest/_torch/custom_ops/test_deepseek_v4_block_table.py(pre-existing, covers the twoupstream
deepseek_v4_compute_sliding_block_tables*ops) is not referenced by any test listand therefore never runs in CI — I confirmed it is absent from
tests/integration/test_lists/entirely and absent from the testReport of my earlier pipeline. My PR only adds a license-header
line to the file those ops live in, so registering that test is out of scope here; flagging it
rather than silently expanding this PR.
CI status
Latest pipeline on
d63bc157fa. Both builds pass on both architectures, and the tests this PRadds are now actually collected and green:
custom_ops/test_deepseek_v4_metadata_ops.pyattention/sparse/test_deepseek_v4_block_table_host_path.pyThe 5 remaining failures are outside this PR's surface —
git diff --name-only <base>...HEAD | grep -iE 'executor|overlap|qwen|scheduler'is empty:test_modeling_qwen_moe—torch.OutOfMemoryErroron a shared A30 with another processalready resident (
Process 892308 has 246 MiB…), plus 1 aggregator assert rolling those upunittest/_torch/executor/test_overlap_scheduler.py— "Test terminated unexpectedly"[Test-SBSA-Single-GPU]shows as failed but is an infrastructure abort, not a test result:ABORTEDat exactly 60.03 min with "Queue task was cancelled" and 0 test failures(9143 passed / 353 skipped) — the Slurm agent never got an allocation. The same job aborted the
same way on three other PRs' builds in the same window.
An earlier run had 5 different failures (llmapi / scaffolding / NanoV3Omni); those did not recur.
Happy to re-trigger once the queue clears if reviewers want a fully green pipeline before review.
Dev Engineer Review
BAD_PAGE_INDEX.test_deepseek_v4_metadata_ops.pyto the B200 and B300 pre-merge test lists._prepare_inputsby 24.2%, with no reported serving throughput or latency regression.QA Engineer Review
Added test functions in
tests/unittest/_torch/attention/sparse/test_deepseek_v4_block_table_host_path.py:test_copy_index_memo_reuses_one_walk_per_step()test_copy_index_memo_matches_uncached_result()test_copy_index_memo_recomputes_on_changed_arguments()test_copy_index_memo_keys_on_contents_not_identity()test_copy_index_memo_is_reset_each_step()test_sliding_block_table_padding_keeps_bad_page_index()test_sliding_block_table_padding_when_full()test_block_offsets_padding_keeps_bad_page_index()These host-path tests are not listed in the modified CI or manual QA test-list files.
Added CUDA-gated coverage in
tests/unittest/_torch/custom_ops/test_deepseek_v4_metadata_ops.pyfor index computation, KV lengths, masks, position IDs, token positions, shared block tables, boundary sizes, padding, optional outputs, invalid inputs, scan limits, and large batches. This test file is covered byl0_b200.ymlandl0_b300.yml.Verdict: needs follow-up because CI or manual QA coverage for the host-path test file is not shown.