From 9852848ff38623ff032cf1a350ef3834c727ebfc Mon Sep 17 00:00:00 2001 From: terminator123 Date: Mon, 3 Aug 2026 15:24:41 +0800 Subject: [PATCH 1/2] add out parameter for dsa api --- python/cudnn/deepseek_sparse_attention/indexer_forward/api.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py index ccf4b722d..d0ad99450 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py @@ -224,6 +224,7 @@ def indexer_forward_wrapper( w: torch.Tensor, ratio: int = 4, qhead_per_kv_head: Optional[int] = None, + out: Optional[torch.Tensor] = None, m_block_size: int = 128, n_block_size: int = 128, q_stage: int = 2, @@ -266,6 +267,7 @@ def indexer_forward_wrapper( w, ratio=ratio, qhead_per_kv_head=qhead_per_kv_head, + out=out, sm_scale=sm_scale, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, From 53091b2f59f96a57c3aa6a2ebb92b22837e42e2f Mon Sep 17 00:00:00 2001 From: terminator123 Date: Mon, 3 Aug 2026 15:45:44 +0800 Subject: [PATCH 2/2] add out paramter for dsa indexer forward --- .../indexer_forward/_interface_sm90.py | 2 +- .../cudnn/deepseek_sparse_attention/indexer_forward/api.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.py index b96474668..c83d4b5ae 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/_interface_sm90.py @@ -47,8 +47,8 @@ def indexer_fwd( w: torch.Tensor, ratio: int = 4, qhead_per_kv_head: Optional[int] = None, - out: Optional[torch.Tensor] = None, sm_scale: float = 1.0, + out: Optional[torch.Tensor] = None, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_k: Optional[torch.Tensor] = None, max_seqlen_q: Optional[int] = None, diff --git a/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py b/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py index d0ad99450..bf3a0b1f9 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_forward/api.py @@ -224,12 +224,12 @@ def indexer_forward_wrapper( w: torch.Tensor, ratio: int = 4, qhead_per_kv_head: Optional[int] = None, - out: Optional[torch.Tensor] = None, m_block_size: int = 128, n_block_size: int = 128, q_stage: int = 2, kv_stage: int = 4, sm_scale: float = 1.0, + out: Optional[torch.Tensor] = None, stream: Optional[cuda.CUstream] = None, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_k: Optional[torch.Tensor] = None, @@ -267,8 +267,8 @@ def indexer_forward_wrapper( w, ratio=ratio, qhead_per_kv_head=qhead_per_kv_head, - out=out, sm_scale=sm_scale, + out=out, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, @@ -288,6 +288,7 @@ def indexer_forward_wrapper( w, ratio=ratio, qhead_per_kv_head=qhead_per_kv_head, + out=out, m_block_size=m_block_size, n_block_size=n_block_size, num_threads=384,