From 55486fb49dfd169495980d351b6651d3ab01e53d Mon Sep 17 00:00:00 2001 From: luzhenyu Date: Tue, 30 Jun 2026 14:57:05 +0800 Subject: [PATCH] Fix grid dim overflow in DSA backward convert kernel on SM100 The convert kernel grid was configured as [1, convert_grid_x, 1], placing the seq-block dimension on grid.y. CUDA caps grid.y/z at 65535, so large mKV.shape[0] / block_seq values trigger `invalid configuration argument`. grid.x supports up to 2^31-1, so move convert_grid_x to grid.x and update the corresponding block_idx() unpacking in the kernel accordingly. No behavior change for in-range sizes. --- .../sparse_attention_backward/dsa_bwd_sm100.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py index 0b4d3c320..7f23c5ef5 100644 --- a/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py +++ b/python/cudnn/deepseek_sparse_attention/sparse_attention_backward/dsa_bwd_sm100.py @@ -572,9 +572,9 @@ class SharedStorage: convert_grid_x = (mKV.shape[0] + self.block_seq - 1) // self.block_seq convert_grid = [ - 1, convert_grid_x, 1, + 1, ] convert_block = [self.num_threads_D_convert, self.num_threads_seq, 1] self.convert( @@ -616,8 +616,8 @@ def convert( ): tidx, tidy, _ = cute.arch.thread_idx() ( - _, seq_block_idx, + _, batch_idx, ) = cute.arch.block_idx()