Skip to content

[None][feat] Add DeepSeek V4 Vanilla sparse attention #18005

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
yihwang-nv wants to merge 3 commits into
base: main
Choose a base branch
from
Open

[None][feat] Add DeepSeek V4 Vanilla sparse attention #18005

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f887350
[None][refactor] add Vanilla sparse attention primitives
yihwang-nv Aug 20, 2026
6982f23
[None][feat] add DeepSeek V4 Vanilla sparse attention
yihwang-nv Aug 21, 2026
318b448
Merge branch 'main' into vanilla-dpskv4-attention
yihwang-nv Aug 25, 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
3 changes: 3 additions & 0 deletions tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from .indexer import DeepseekV4Indexer
from .metadata import DeepseekV4TrtllmAttentionMetadata
from .params import DeepseekV4AttentionType, DeepSeekV4MetadataParams, DeepSeekV4Params
from .vanilla_backend import DeepseekV4VanillaAttention, DeepseekV4VanillaIndexer

__all__ = [
"DeepSeekV4MetadataParams",
Expand All @@ -27,4 +28,6 @@
"DeepseekV4Indexer",
"DeepseekV4TrtllmAttention",
"DeepseekV4TrtllmAttentionMetadata",
"DeepseekV4VanillaAttention",
"DeepseekV4VanillaIndexer",
]
2 changes: 2 additions & 0 deletions tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1157,6 +1157,7 @@ def _indexer_branch():
if self.apply_rotary_emb:
assert ctx_position_ids is not None
k_pe_ctx = self.apply_rope(q_ctx, k_pe_ctx, ctx_position_ids)
latent_cache_ctx = torch.cat([compressed_kv_ctx, k_pe_ctx], dim=-1)

context_o_lora_bmm_input = forward_context_sparse_attn(
self,
Expand Down Expand Up @@ -1185,6 +1186,7 @@ def _indexer_branch():
if self.apply_rotary_emb:
assert gen_position_ids is not None
k_pe_gen = self.apply_rope(q_gen, k_pe_gen, gen_position_ids)
latent_cache_gen = torch.cat([compressed_kv_gen, k_pe_gen], dim=-1)

generation_o_lora_bmm_input = forward_generation_sparse_attn(
self,
Expand Down
Loading
Loading