Skip to content

[None][feat] DSv4 prep: attention op plumbing #15384

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

Merged
lfr-0531 merged 5 commits into from
Jun 17, 2026
Merged

[None][feat] DSv4 prep: attention op plumbing #15384

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9e86d65
[None][feat] DSv4 prep: attention op plumbing
lfr-0531 Jun 15, 2026
7e8d1ff
[None][feat] address attention op review comments
lfr-0531 Jun 16, 2026
bfacdce
[None][feat] keep spec decoding draft tokens behavior
lfr-0531 Jun 16, 2026
1a5e6d6
Revert "[None][feat] keep spec decoding draft tokens behavior"
lfr-0531 Jun 16, 2026
edcd002
[None][feat] group spec decoding draft token kwarg
lfr-0531 Jun 16, 2026
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
4 changes: 4 additions & 0 deletions tensorrt_llm/_torch/attention_backend/interface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -733,6 +733,10 @@ class AttentionSparseArgs:
sparse_attn_indices: Optional[torch.Tensor] = None
sparse_attn_offsets: Optional[torch.Tensor] = None
sparse_attn_indices_block_size: int = 0
# DeepSeek-V4 sparse-MLA only: per-token compressed top-k lengths and the
# base pointer of the compressed KV cache pool (compress_ratio > 1).
sparse_mla_topk_lens: Optional[torch.Tensor] = None
compressed_kv_cache_pool_ptr: Optional[int] = None


@dataclass(kw_only=True, slots=True)
Expand Down
40 changes: 18 additions & 22 deletions tensorrt_llm/_torch/attention_backend/trtllm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,9 @@
get_model_extra_attrs)
from .interface import (AttentionBackend, AttentionForwardArgs,
AttentionInputType, AttentionMask, AttentionMetadata,
AttentionSparseArgs, KVCacheParams, MLAParams,
PositionalEmbeddingParams, PredefinedAttentionMask,
RopeParams, merge_attention_forward_args)
KVCacheParams, MLAParams, PositionalEmbeddingParams,
PredefinedAttentionMask, RopeParams,
merge_attention_forward_args)

# Enable TRTLLM-Gen attention backend by default. Set
# TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION=0 to force the thop.attention path.
Expand All @@ -44,10 +44,7 @@
# ``thop.attention`` kwargs hard-wired to a literal at the call site (no
# rich object owns them). Sync test enforces both the kwarg name and the
# literal value.
_THOP_LITERALS: dict = {
"sparse_mla_topk_lens": None,
"compressed_kv_cache_pool_ptr": None,
}
_THOP_LITERALS: dict = {}


@functools.cache
Expand Down Expand Up @@ -114,6 +111,7 @@ class TrtllmAttentionMetadata(AttentionMetadata):
is_spec_dec_dynamic_tree: bool = False

# parameters required for spec-dec mode
max_total_draft_tokens: Optional[int] = None
spec_decoding_position_offsets: Optional[torch.Tensor] = None
# C++ attention op requires a 2-D position_offsets tensor and reads
# sizes()[1] as the generation length / packed-mask row stride.
Expand Down Expand Up @@ -1631,6 +1629,8 @@ def _run(
spec_decoding_bl_tree_mask_offset=metadata.
spec_decoding_bl_tree_mask_offset,
spec_decoding_bl_tree_mask=metadata.spec_decoding_bl_tree_mask,
spec_decoding_target_max_draft_tokens=metadata.
max_total_draft_tokens,
spec_bl_tree_first_sparse_mask_offset_kv=metadata.
spec_bl_tree_first_sparse_mask_offset_kv,
num_sparse_topk=metadata.num_sparse_topk,
Expand Down Expand Up @@ -1719,14 +1719,11 @@ def _run(
sparse_attn_offsets=forward_args.sparse.sparse_attn_offsets,
sparse_attn_indices_block_size=forward_args.sparse.
sparse_attn_indices_block_size,
sparse_mla_topk_lens=forward_args.sparse.sparse_mla_topk_lens,
compressed_kv_cache_pool_ptr=forward_args.sparse.
compressed_kv_cache_pool_ptr,

# --- Literals intentionally None (see _THOP_LITERALS) ---
Comment thread
yuxianq marked this conversation as resolved.
# ``sparse_mla_topk_lens`` and ``compressed_kv_cache_pool_ptr``
# stay as literal ``None`` until DeepSeek V4 sparse-MLA lands.
sparse_mla_topk_lens=None,
compressed_kv_cache_pool_ptr=None,
spec_decoding_target_max_draft_tokens=getattr(
metadata, 'max_total_draft_tokens', None),
# --- Literals intentionally in _THOP_LITERALS ---
)

if self.print_skip_softmax_stat:
Expand Down Expand Up @@ -1811,14 +1808,13 @@ def forward(
forward_args)
at_idx, at_off = self.sparse_attn_predict(q, k, metadata,
forward_args)
forward_args.sparse = AttentionSparseArgs(
sparse_kv_indices=kv_idx,
sparse_kv_offsets=kv_off,
sparse_attn_indices=at_idx,
sparse_attn_offsets=at_off,
sparse_attn_indices_block_size=self.sparse_attention_config.
get_indices_block_size(),
)
sparse_args = forward_args.sparse
sparse_args.sparse_kv_indices = kv_idx
sparse_args.sparse_kv_offsets = kv_off
sparse_args.sparse_attn_indices = at_idx
sparse_args.sparse_attn_offsets = at_off
sparse_args.sparse_attn_indices_block_size = (
self.sparse_attention_config.get_indices_block_size())

# Compute FlashMLA tile-scheduler metadata once per forward pass.
# The flag is reset in prepare_flash_mla() and update_for_spec_dec() to trigger
Expand Down
Loading