From 7b62afcc9cd33e07e3363f608fc3a0d9284872d1 Mon Sep 17 00:00:00 2001 From: David Orman Date: Sun, 6 Sep 2026 06:18:15 -0500 Subject: [PATCH 1/2] Reduce repeated attention setup during speculative decoding --- .../kernels/ops/attention/dsa_metadata.py | 342 ++++++++++++------ python/sglang/srt/environ.py | 2 + .../layers/attention/dsa/dsa_backend_kpool.py | 47 +++ .../dsa/dsa_backend_mtp_precompute.py | 6 +- .../srt/layers/attention/dsa_backend.py | 214 +++++++++-- .../ops/attention/test_dsa_metadata.py | 56 ++- .../test_kpool_metadata_fusion_contract.py | 69 ++++ 7 files changed, 588 insertions(+), 148 deletions(-) create mode 100644 test/registered/unit/layers/attention/test_kpool_metadata_fusion_contract.py diff --git a/python/sglang/kernels/ops/attention/dsa_metadata.py b/python/sglang/kernels/ops/attention/dsa_metadata.py index df430081a19a..8d029b61423b 100644 --- a/python/sglang/kernels/ops/attention/dsa_metadata.py +++ b/python/sglang/kernels/ops/attention/dsa_metadata.py @@ -4,12 +4,22 @@ import triton import triton.language as tl +# Bound the replay grid while retaining enough copy programs to saturate the GPU. +_TILE_PROGRAM_TARGET = 8192 + + +def _bounded_scan_num_splits(rows: int, num_col_blocks: int) -> int: + """Keep the grid capture-safe while bounding traversal by replay-time data.""" + assert rows > 0 + return max(1, min(num_col_blocks, _TILE_PROGRAM_TARGET // rows)) + @triton.jit( do_not_specialize=[ "page_table_stride_0", "real_page_table_stride_0", "max_len", + "num_splits", ] ) def _fused_dsa_decode_metadata_kernel( @@ -32,7 +42,9 @@ def _fused_dsa_decode_metadata_kernel( real_page_table_stride_1: tl.constexpr, bs: tl.constexpr, max_len, + num_splits, dsa_index_topk: tl.constexpr, + index_kpool: tl.constexpr, real_page_size: tl.constexpr, HAS_REAL_PAGE_TABLE: tl.constexpr, HAS_PAGE_TABLE_1: tl.constexpr, @@ -46,7 +58,14 @@ def _fused_dsa_decode_metadata_kernel( mask_b = offs_b < bs seq = tl.load(seq_lens + offs_b * seq_lens_stride, mask=mask_b, other=0) seq_i32 = seq.to(tl.int32) - dsa_seq = tl.minimum(seq_i32, dsa_index_topk) + if index_kpool <= 1: + dsa_seq = tl.minimum(seq_i32, dsa_index_topk) + else: + # Preserve the live partial pool after selecting pool-aligned history. + full_pool_tokens = (seq_i32 // index_kpool) * index_kpool + selected_history_tokens = tl.minimum(full_pool_tokens, dsa_index_topk) + tail_tokens = seq_i32 - full_pool_tokens + dsa_seq = selected_history_tokens + tail_tokens cu = tl.cumsum(seq_i32, 0) dsa_cu = tl.cumsum(dsa_seq, 0) @@ -59,54 +78,53 @@ def _fused_dsa_decode_metadata_kernel( tl.store(dsa_cu_seqlens_k + 1 + offs_b, dsa_cu, mask=mask_b) return - num_col_blocks = tl.cdiv(max_len, BLOCK_N) page_pid = pid - 1 - row = page_pid // num_col_blocks - col_block = page_pid - row * num_col_blocks - offs_n = col_block * BLOCK_N + tl.arange(0, BLOCK_N) - mask = (row < bs) & (offs_n < max_len) + row = page_pid // num_splits + split_id = page_pid - row * num_splits req_idx = tl.load( req_pool_indices + row * req_pool_indices_stride, mask=row < bs, other=0, ) - # Skip column blocks past the request's kv length: no consumer reads there - # (attention and the indexer both stay within cache_seqlens). Loaded after - # req_idx so the two scalar loads pipeline (no added latency when live). kv_len = tl.load( seq_lens + row * seq_lens_stride, mask=row < bs, other=0, ).to(tl.int32) - if col_block * BLOCK_N >= kv_len: - return - vals = tl.load( - req_to_token + req_idx * req_to_token_stride_0 + offs_n * req_to_token_stride_1, - mask=mask, - other=0, - ).to(tl.int32) - # Write the wide page_size=1 table only when the caller provides it; the - # fused decode CUDA graph drops it and consumes real_page_table alone. - if HAS_PAGE_TABLE_1: - tl.store( - page_table_1 - + row.to(tl.int64) * page_table_stride_0 - + offs_n * page_table_stride_1, - vals, + # Page-table row offsets can overflow int32 at 1M context. + row_i64 = row.to(tl.int64) + num_live_blocks = tl.minimum(tl.cdiv(kv_len, BLOCK_N), tl.cdiv(max_len, BLOCK_N)) + # Three stages hide latency across strided copy iterations. + for col_block in tl.range(split_id, num_live_blocks, num_splits, num_stages=3): + offs_n = col_block * BLOCK_N + tl.arange(0, BLOCK_N) + mask = (row < bs) & (offs_n < max_len) + vals = tl.load( + req_to_token + + req_idx * req_to_token_stride_0 + + offs_n * req_to_token_stride_1, mask=mask, - ) + other=0, + ).to(tl.int32) + if HAS_PAGE_TABLE_1: + tl.store( + page_table_1 + + row_i64 * page_table_stride_0 + + offs_n * page_table_stride_1, + vals, + mask=mask, + ) - if HAS_REAL_PAGE_TABLE: - real_mask = mask & ((offs_n % real_page_size) == 0) - real_cols = offs_n // real_page_size - tl.store( - real_page_table - + row.to(tl.int64) * real_page_table_stride_0 - + real_cols * real_page_table_stride_1, - vals // real_page_size, - mask=real_mask, - ) + if HAS_REAL_PAGE_TABLE: + real_mask = mask & ((offs_n % real_page_size) == 0) + real_cols = offs_n // real_page_size + tl.store( + real_page_table + + row_i64 * real_page_table_stride_0 + + real_cols * real_page_table_stride_1, + vals // real_page_size, + mask=real_mask, + ) def fused_dsa_decode_metadata( @@ -123,6 +141,7 @@ def fused_dsa_decode_metadata( max_len: int, dsa_index_topk: int, real_page_size: int, + index_kpool: int = 1, ) -> None: """Fill decode-graph DSA metadata (seqlens + page tables) from req_to_token. @@ -136,6 +155,11 @@ def fused_dsa_decode_metadata( Contract: each page-table row is written only over its live prefix ([:cache_seqlens]); the tail keeps stale values across CUDA-graph replays, so consumers must bound reads by cache_seqlens. + + The column scan is bounded inside the kernel by each row's own kv length + (read at run time), so the cost scales with the live sequence lengths and + not with ``max_len`` (the table width); the grid itself stays + data-independent. See :func:`_bounded_scan_num_splits`. """ assert seq_lens.is_cuda assert req_pool_indices.is_cuda @@ -149,6 +173,7 @@ def fused_dsa_decode_metadata( cu_seqlens_k[:1].zero_() dsa_cu_seqlens_k[:1].zero_() return + assert index_kpool > 0 has_real_page_table = real_page_size > 1 if has_real_page_table: @@ -171,7 +196,8 @@ def fused_dsa_decode_metadata( block_bs = triton.next_power_of_2(bs) block_n = 128 num_col_blocks = triton.cdiv(max_len, block_n) - grid = (1 + bs * num_col_blocks,) + num_splits = _bounded_scan_num_splits(bs, num_col_blocks) + grid = (1 + bs * num_splits,) _fused_dsa_decode_metadata_kernel[grid]( seq_lens, @@ -193,7 +219,9 @@ def fused_dsa_decode_metadata( real_page_table.stride(1) if has_real_page_table else 0, bs, max_len, + num_splits, dsa_index_topk, + index_kpool, real_page_size, has_real_page_table, has_page_table_1, @@ -207,6 +235,7 @@ def fused_dsa_decode_metadata( "page_table_stride_0", "real_page_table_stride_0", "max_seqlen_k", + "num_splits", ] ) def _fused_dsa_target_verify_metadata_kernel( @@ -233,7 +262,9 @@ def _fused_dsa_target_verify_metadata_kernel( paged_mqa_ctx_lens_stride_1: tl.constexpr, bs: tl.constexpr, max_seqlen_k, + num_splits, dsa_index_topk: tl.constexpr, + index_kpool: tl.constexpr, real_page_size: tl.constexpr, next_n: tl.constexpr, HAS_REAL_PAGE_TABLE: tl.constexpr, @@ -268,7 +299,14 @@ def _fused_dsa_target_verify_metadata_kernel( ).to(tl.int32) expanded_seq = base_seq + draft_off + 1 expanded_seq = tl.where(mask_e, expanded_seq, 0) - dsa_seq = tl.minimum(expanded_seq, dsa_index_topk) + if index_kpool <= 1: + dsa_seq = tl.minimum(expanded_seq, dsa_index_topk) + else: + # Preserve the live partial pool after selecting pool-aligned history. + full_pool_tokens = (expanded_seq // index_kpool) * index_kpool + selected_history_tokens = tl.minimum(full_pool_tokens, dsa_index_topk) + tail_tokens = expanded_seq - full_pool_tokens + dsa_seq = selected_history_tokens + tail_tokens dsa_cu = tl.cumsum(dsa_seq, 0) tl.store(seqlens_expanded + offs_e, expanded_seq, mask=mask_e) @@ -286,12 +324,9 @@ def _fused_dsa_target_verify_metadata_kernel( ) return - num_col_blocks = tl.cdiv(max_seqlen_k, BLOCK_N) page_pid = pid - 1 - out_row = page_pid // num_col_blocks - col_block = page_pid - out_row * num_col_blocks - offs_n = col_block * BLOCK_N + tl.arange(0, BLOCK_N) - mask = (out_row < expanded_size) & (offs_n < max_seqlen_k) + out_row = page_pid // num_splits + split_id = page_pid - out_row * num_splits req_row = out_row // next_n req_idx = tl.load( @@ -299,10 +334,6 @@ def _fused_dsa_target_verify_metadata_kernel( mask=out_row < expanded_size, other=0, ) - # Skip column blocks past the request's kv length (seq_len + next_n): no - # consumer reads there (attention and the indexer stay within cache_seqlens). - # Loaded after req_idx so the two scalar loads pipeline (no added latency - # when live). kv_len = ( tl.load( seq_lens + req_row * seq_lens_stride, @@ -311,37 +342,43 @@ def _fused_dsa_target_verify_metadata_kernel( ).to(tl.int32) + next_n ) - if col_block * BLOCK_N >= kv_len: - return - vals = tl.load( - req_to_token + req_idx * req_to_token_stride_0 + offs_n * req_to_token_stride_1, - mask=mask, - other=0, - ).to(tl.int32) - # Write the wide page_size=1 table only when the caller provides it (see - # fused_dsa_decode_metadata for the optional-page_table_1 contract). - if HAS_PAGE_TABLE_1: - tl.store( - page_table_1 - + out_row.to(tl.int64) * page_table_stride_0 - + offs_n * page_table_stride_1, - vals, + # Output-row offsets can overflow int32 at 1M context. + out_row_i64 = out_row.to(tl.int64) + num_live_blocks = tl.minimum( + tl.cdiv(kv_len, BLOCK_N), tl.cdiv(max_seqlen_k, BLOCK_N) + ) + for col_block in tl.range(split_id, num_live_blocks, num_splits, num_stages=3): + offs_n = col_block * BLOCK_N + tl.arange(0, BLOCK_N) + mask = (out_row < expanded_size) & (offs_n < max_seqlen_k) + vals = tl.load( + req_to_token + + req_idx * req_to_token_stride_0 + + offs_n * req_to_token_stride_1, mask=mask, - ) + other=0, + ).to(tl.int32) + if HAS_PAGE_TABLE_1: + tl.store( + page_table_1 + + out_row_i64 * page_table_stride_0 + + offs_n * page_table_stride_1, + vals, + mask=mask, + ) - if HAS_REAL_PAGE_TABLE: - real_mask = mask & ((offs_n % real_page_size) == 0) - real_cols = offs_n // real_page_size - tl.store( - real_page_table - + out_row.to(tl.int64) * real_page_table_stride_0 - + real_cols * real_page_table_stride_1, - vals // real_page_size, - mask=real_mask, - ) + if HAS_REAL_PAGE_TABLE: + real_mask = mask & ((offs_n % real_page_size) == 0) + real_cols = offs_n // real_page_size + tl.store( + real_page_table + + out_row_i64 * real_page_table_stride_0 + + real_cols * real_page_table_stride_1, + vals // real_page_size, + mask=real_mask, + ) -def fused_dsa_target_verify_metadata( +def _prep_fused_dsa_target_verify_metadata_launch( seq_lens: torch.Tensor, req_pool_indices: torch.Tensor, req_to_token: torch.Tensor, @@ -358,7 +395,8 @@ def fused_dsa_target_verify_metadata( real_page_size: int, next_n: int, paged_mqa_ctx_lens_2d: torch.Tensor = None, -) -> None: + index_kpool: int = 1, +): assert seq_lens.is_cuda assert req_pool_indices.is_cuda assert req_to_token.is_cuda @@ -368,11 +406,9 @@ def fused_dsa_target_verify_metadata( assert dsa_cache_seqlens.is_cuda assert dsa_cu_seqlens_k.is_cuda - if bs == 0: - cu_seqlens_k[:1].zero_() - dsa_cu_seqlens_k[:1].zero_() - return + assert bs > 0 assert next_n > 0 + assert index_kpool > 0 has_real_page_table = real_page_size > 1 if has_real_page_table: @@ -406,9 +442,10 @@ def fused_dsa_target_verify_metadata( block_expanded = triton.next_power_of_2(expanded_size) block_n = 128 num_col_blocks = triton.cdiv(max_seqlen_k, block_n) - grid = (1 + expanded_size * num_col_blocks,) + num_splits = _bounded_scan_num_splits(expanded_size, num_col_blocks) + grid = (1 + expanded_size * num_splits,) - _fused_dsa_target_verify_metadata_kernel[grid]( + args = ( seq_lens, req_pool_indices, req_to_token, @@ -432,16 +469,69 @@ def fused_dsa_target_verify_metadata( paged_mqa_ctx_lens_2d.stride(1) if has_paged_mqa_ctx_lens else 0, bs, max_seqlen_k, + num_splits, dsa_index_topk, + index_kpool, real_page_size, next_n, has_real_page_table, has_paged_mqa_ctx_lens, has_page_table_1, + ) + constexprs = dict( BLOCK_BS=block_bs, BLOCK_EXPANDED=block_expanded, BLOCK_N=block_n, ) + return grid, args, constexprs + + +def fused_dsa_target_verify_metadata( + seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + cache_seqlens: torch.Tensor, + cu_seqlens_k: torch.Tensor, + page_table_1: Optional[torch.Tensor], + seqlens_expanded: torch.Tensor, + dsa_cache_seqlens: torch.Tensor, + dsa_cu_seqlens_k: torch.Tensor, + real_page_table: torch.Tensor, + bs: int, + max_seqlen_k: int, + dsa_index_topk: int, + real_page_size: int, + next_n: int, + paged_mqa_ctx_lens_2d: torch.Tensor = None, + index_kpool: int = 1, +) -> None: + if bs == 0: + assert cu_seqlens_k.is_cuda + assert dsa_cu_seqlens_k.is_cuda + cu_seqlens_k[:1].zero_() + dsa_cu_seqlens_k[:1].zero_() + return + + grid, args, constexprs = _prep_fused_dsa_target_verify_metadata_launch( + seq_lens, + req_pool_indices, + req_to_token, + cache_seqlens, + cu_seqlens_k, + page_table_1, + seqlens_expanded, + dsa_cache_seqlens, + dsa_cu_seqlens_k, + real_page_table, + bs, + max_seqlen_k, + dsa_index_topk, + real_page_size, + next_n, + paged_mqa_ctx_lens_2d, + index_kpool, + ) + _fused_dsa_target_verify_metadata_kernel[grid](*args, **constexprs) @triton.jit( @@ -450,6 +540,7 @@ def fused_dsa_target_verify_metadata( "real_page_table_stride_0", "total_len", "max_seqlen_k", + "num_splits", ] ) def _fused_dsa_draft_extend_metadata_kernel( @@ -476,7 +567,9 @@ def _fused_dsa_draft_extend_metadata_kernel( bs: tl.constexpr, total_len, max_seqlen_k, + num_splits, dsa_index_topk: tl.constexpr, + index_kpool: tl.constexpr, real_page_size: tl.constexpr, HAS_REAL_PAGE_TABLE: tl.constexpr, HAS_PAGE_TABLE_1: tl.constexpr, @@ -536,7 +629,14 @@ def _fused_dsa_draft_extend_metadata_kernel( expanded_seq = base_seq - qo_len_for_row + local_off + 1 expanded_seq = tl.maximum(expanded_seq, 0) expanded_seq = tl.where(mask_e, expanded_seq, 0) - dsa_seq = tl.minimum(expanded_seq, dsa_index_topk) + if index_kpool <= 1: + dsa_seq = tl.minimum(expanded_seq, dsa_index_topk) + else: + # Preserve the live partial pool after selecting pool-aligned history. + full_pool_tokens = (expanded_seq // index_kpool) * index_kpool + selected_history_tokens = tl.minimum(full_pool_tokens, dsa_index_topk) + tail_tokens = expanded_seq - full_pool_tokens + dsa_seq = selected_history_tokens + tail_tokens dsa_cu = tl.cumsum(dsa_seq, 0) tl.store(seqlens_expanded + offs_e, expanded_seq, mask=mask_e) @@ -545,25 +645,25 @@ def _fused_dsa_draft_extend_metadata_kernel( tl.store(dsa_cu_seqlens_k + 1 + offs_e, dsa_cu, mask=mask_e) return - num_col_blocks = tl.cdiv(max_seqlen_k, BLOCK_N) page_pid = pid - 1 - req_row = page_pid // num_col_blocks - col_block = page_pid - req_row * num_col_blocks - offs_n = col_block * BLOCK_N + tl.arange(0, BLOCK_N) + req_row = page_pid // num_splits + split_id = page_pid - req_row * num_splits qo_len = tl.load( extend_seq_lens + req_row * extend_seq_lens_stride, mask=req_row < bs, other=0, ).to(tl.int32) - # Skip column blocks past the request's kv length: no consumer reads there - # (attention and the indexer both stay within cache_seqlens). kv_len = tl.load( seq_lens + req_row * seq_lens_stride, mask=req_row < bs, other=0, ).to(tl.int32) - if col_block * BLOCK_N >= kv_len: + # Bound the scan by the live replay-time kv length. + num_live_blocks = tl.minimum( + tl.cdiv(kv_len, BLOCK_N), tl.cdiv(max_seqlen_k, BLOCK_N) + ) + if split_id >= num_live_blocks: return if STATIC_EXTEND_LEN: prefix = req_row * qo_len @@ -577,41 +677,46 @@ def _fused_dsa_draft_extend_metadata_kernel( offs_r = tl.arange(0, BLOCK_ROWS) out_rows = prefix + offs_r row_mask = (req_row < bs) & (offs_r < qo_len) & (out_rows < total_len) - col_mask = offs_n < max_seqlen_k has_rows = (req_row < bs) & (qo_len > 0) - mask = row_mask[:, None] & col_mask[None, :] req_idx = tl.load( req_pool_indices + req_row * req_pool_indices_stride, mask=has_rows, other=0, ) - vals = tl.load( - req_to_token + req_idx * req_to_token_stride_0 + offs_n * req_to_token_stride_1, - mask=col_mask & has_rows, - other=0, - ).to(tl.int32) - # Write the wide page_size=1 table only when the caller provides it (see - # fused_dsa_decode_metadata for the optional-page_table_1 contract). - if HAS_PAGE_TABLE_1: - tl.store( - page_table_1 - + out_rows.to(tl.int64)[:, None] * page_table_stride_0 - + offs_n[None, :] * page_table_stride_1, - vals[None, :], - mask=mask, - ) + # Output-row offsets can overflow int32 at 1M context. + out_rows_i64 = out_rows.to(tl.int64) + for col_block in tl.range(split_id, num_live_blocks, num_splits, num_stages=3): + offs_n = col_block * BLOCK_N + tl.arange(0, BLOCK_N) + col_mask = offs_n < max_seqlen_k + mask = row_mask[:, None] & col_mask[None, :] + + vals = tl.load( + req_to_token + + req_idx * req_to_token_stride_0 + + offs_n * req_to_token_stride_1, + mask=col_mask & has_rows, + other=0, + ).to(tl.int32) + if HAS_PAGE_TABLE_1: + tl.store( + page_table_1 + + out_rows_i64[:, None] * page_table_stride_0 + + offs_n[None, :] * page_table_stride_1, + vals[None, :], + mask=mask, + ) - if HAS_REAL_PAGE_TABLE: - real_mask = mask & ((offs_n[None, :] % real_page_size) == 0) - real_cols = offs_n // real_page_size - tl.store( - real_page_table - + out_rows.to(tl.int64)[:, None] * real_page_table_stride_0 - + real_cols[None, :] * real_page_table_stride_1, - (vals // real_page_size)[None, :], - mask=real_mask, - ) + if HAS_REAL_PAGE_TABLE: + real_mask = mask & ((offs_n[None, :] % real_page_size) == 0) + real_cols = offs_n // real_page_size + tl.store( + real_page_table + + out_rows_i64[:, None] * real_page_table_stride_0 + + real_cols[None, :] * real_page_table_stride_1, + (vals // real_page_size)[None, :], + mask=real_mask, + ) def fused_dsa_draft_extend_metadata( @@ -634,6 +739,7 @@ def fused_dsa_draft_extend_metadata( max_extend_len: int, max_total_len: int, static_extend_len: bool = False, + index_kpool: int = 1, ) -> None: assert seq_lens.is_cuda assert extend_seq_lens.is_cuda @@ -662,6 +768,7 @@ def fused_dsa_draft_extend_metadata( # that would sync in the replay hot path. assert max_extend_len > 0 assert total_len <= bs * max_extend_len + assert index_kpool > 0 has_real_page_table = real_page_size > 1 if has_real_page_table: @@ -685,7 +792,8 @@ def fused_dsa_draft_extend_metadata( block_rows = triton.next_power_of_2(max_extend_len) block_n = 128 num_col_blocks = triton.cdiv(max_seqlen_k, block_n) - grid = (1 + bs * num_col_blocks,) + num_splits = _bounded_scan_num_splits(bs, num_col_blocks) + grid = (1 + bs * num_splits,) _fused_dsa_draft_extend_metadata_kernel[grid]( seq_lens, @@ -711,7 +819,9 @@ def fused_dsa_draft_extend_metadata( bs, total_len, max_seqlen_k, + num_splits, dsa_index_topk, + index_kpool, real_page_size, has_real_page_table, has_page_table_1, diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5767726654bc..7fe052d92e92 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1512,6 +1512,8 @@ class Envs: SGLANG_DSA_FUSE_TOPK = EnvBoolWithAlias( True, deprecated_name="SGLANG_NSA_FUSE_TOPK" ) + # Fused KPool replay metadata and derived-buffer reuse by later draft steps. + SGLANG_EXPERIMENTAL_DSA_KPOOL_METADATA_FUSION = EnvBool(False) SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC = EnvBool(False) SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK = EnvStr(None) SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD = EnvIntWithAlias( diff --git a/python/sglang/srt/layers/attention/dsa/dsa_backend_kpool.py b/python/sglang/srt/layers/attention/dsa/dsa_backend_kpool.py index acf04fcb7b86..5f1d1c21b3b1 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_backend_kpool.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_backend_kpool.py @@ -13,6 +13,7 @@ update_kpool_write_plan, update_pooled_paged_mqa_metadata, ) +from sglang.srt.utils import is_cuda if TYPE_CHECKING: from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import ( @@ -23,6 +24,17 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +def _is_kpool_metadata_fusion_supported( + index_kpool: int, page_size: int, index_topk: int +) -> bool: + return ( + index_kpool > 1 + and page_size == 64 + and page_size % index_kpool == 0 + and index_topk % index_kpool == 0 + ) + + @dataclass class _KPoolForwardInputs: full_real_page_table: Optional[torch.Tensor] = None @@ -302,3 +314,38 @@ def _update_kpool_metadata_from_precomputed( forward_mode=forward_mode, slots_per_page=slots_per_page, ) + + def _copy_kpool_metadata_from_sibling( + self, metadata: DSAMetadata, src_metadata: DSAMetadata + ) -> None: + """Copy KPool metadata derived from identical inputs from a sibling.""" + if self.dsa_index_kpool <= 1 or not is_cuda(): + return + + if metadata.pooled_cache_seqlens_int32 is not None: + metadata.pooled_cache_seqlens_int32.copy_( + src_metadata.pooled_cache_seqlens_int32 + ) + if metadata.pooled_real_page_table is not None: + metadata.pooled_real_page_table.copy_(src_metadata.pooled_real_page_table) + if metadata.pooled_paged_mqa_schedule_metadata is not None: + metadata.pooled_paged_mqa_schedule_metadata.copy_( + src_metadata.pooled_paged_mqa_schedule_metadata + ) + + dst_plan = metadata.kpool_write_plan + src_plan = src_metadata.kpool_write_plan + if dst_plan is None: + return + dst_plan.req.copy_(src_plan.req) + dst_plan.write_start.copy_(src_plan.write_start) + dst_plan.tail_logical_start.copy_(src_plan.tail_logical_start) + dst_plan.write_loc.copy_(src_plan.write_loc) + if dst_plan.pool_seqlens_per_q is not None: + dst_plan.pool_seqlens_per_q.copy_(src_plan.pool_seqlens_per_q) + if dst_plan.seqlens_per_q is not None: + dst_plan.seqlens_per_q.copy_(src_plan.seqlens_per_q) + if dst_plan.pool_schedule_metadata is not None: + dst_plan.pool_schedule_metadata.copy_(src_plan.pool_schedule_metadata) + if dst_plan.effective_n_per_batch is not None: + dst_plan.effective_n_per_batch.copy_(src_plan.effective_n_per_batch) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py index b6f1291a2ff3..aa9a73e82dfc 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py @@ -122,7 +122,10 @@ def _precompute_decode_mode( """Precompute metadata for normal decode mode.""" max_len = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1] - if (_is_cuda or _is_hip) and self.dsa_index_kpool <= 1: + if (_is_cuda or _is_hip) and ( + self.dsa_index_kpool <= 1 + or (not _is_hip and getattr(self, "experimental_kpool_metadata_fusion", False)) + ): from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_decode_metadata, ) @@ -160,6 +163,7 @@ def _precompute_decode_mode( max_len=max_len, dsa_index_topk=self.dsa_index_topk, real_page_size=self.real_page_size, + index_kpool=self.dsa_index_kpool, ) seqlens_expanded = cache_seqlens seqlens_expanded_size = bs diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 22c35d36604d..2d2b2b57fe74 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -57,6 +57,7 @@ from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.dsa.dsa_backend_kpool import ( DeepseekSparseAttnBackendKPoolMixin, + _is_kpool_metadata_fusion_supported, _KPoolForwardInputs, ) from sglang.srt.layers.attention.dsa.dsa_backend_mtp_precompute import ( @@ -346,6 +347,36 @@ def __init__( self.dsa_index_topk = get_dsa_index_topk(hf_config) self.dsa_index_kpool = get_dsa_index_kpool(hf_config) self.needs_cpu_seq_lens = self.dsa_index_kpool > 1 + # The env is global, so unsupported KPool geometry or platforms must fall back. + _kpool_fusion_requested = ( + envs.SGLANG_EXPERIMENTAL_DSA_KPOOL_METADATA_FUSION.get() + ) + _kpool_fusion_supported = _is_kpool_metadata_fusion_supported( + self.dsa_index_kpool, + self.real_page_size, + self.dsa_index_topk, + ) + _kpool_fusion_platform_ok = is_cuda() and not _is_hip + self.experimental_kpool_metadata_fusion = ( + _kpool_fusion_requested + and _kpool_fusion_supported + and _kpool_fusion_platform_ok + ) + if ( + _kpool_fusion_requested + and self.dsa_index_kpool > 1 + and not _kpool_fusion_supported + ): + logger.warning( + "SGLANG_EXPERIMENTAL_DSA_KPOOL_METADATA_FUSION is set but this " + "DSA backend's geometry is outside the validated envelope " + "(index_kpool=%s, page_size=%s, index_topk=%s; required: " + "page_size=64, page-aligned pool, pool-aligned topk) - keeping " + "the eager metadata path for this backend.", + self.dsa_index_kpool, + self.real_page_size, + self.dsa_index_topk, + ) self.max_context_len = model_runner.model_config.context_len self.num_q_heads = ( model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size @@ -1530,7 +1561,10 @@ def _apply_cuda_graph_metadata( # Normal Decode max_len = self._graph_page_table_width(metadata) - if (is_cuda() or _is_hip) and self.dsa_index_kpool <= 1: + if (is_cuda() or _is_hip) and ( + self.dsa_index_kpool <= 1 + or (not _is_hip and getattr(self, "experimental_kpool_metadata_fusion", False)) + ): fused_dsa_decode_metadata( seq_lens=seq_lens, req_pool_indices=req_pool_indices, @@ -1545,6 +1579,7 @@ def _apply_cuda_graph_metadata( max_len=max_len, dsa_index_topk=self.dsa_index_topk, real_page_size=self.real_page_size, + index_kpool=self.dsa_index_kpool, ) cache_seqlens = metadata.cache_seqlens_int32 dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32 @@ -1570,7 +1605,10 @@ def _apply_cuda_graph_metadata( elif forward_mode.is_target_verify(): max_seqlen_k = self._graph_page_table_width(metadata) - if (is_cuda() or _is_hip) and self.dsa_index_kpool <= 1: + if (is_cuda() or _is_hip) and ( + self.dsa_index_kpool <= 1 + or (not _is_hip and getattr(self, "experimental_kpool_metadata_fusion", False)) + ): paged_mqa_ctx_lens_2d = None if ( self.speculative_num_draft_tokens >= 2 @@ -1600,6 +1638,7 @@ def _apply_cuda_graph_metadata( real_page_size=self.real_page_size, next_n=self.speculative_num_draft_tokens, paged_mqa_ctx_lens_2d=paged_mqa_ctx_lens_2d, + index_kpool=self.dsa_index_kpool, ) target_verify_ctx_lens_written = paged_mqa_ctx_lens_2d is not None cache_seqlens = metadata.cache_seqlens_int32 @@ -1666,7 +1705,10 @@ def _apply_cuda_graph_metadata( device=self.device, ) - if (is_cuda() or _is_hip) and self.dsa_index_kpool <= 1: + if (is_cuda() or _is_hip) and ( + self.dsa_index_kpool <= 1 + or (not _is_hip and getattr(self, "experimental_kpool_metadata_fusion", False)) + ): fused_dsa_draft_extend_metadata( seq_lens=seq_lens, extend_seq_lens=extend_seq_lens, @@ -1687,6 +1729,7 @@ def _apply_cuda_graph_metadata( max_extend_len=self.speculative_num_draft_tokens, max_total_len=bs * self.speculative_num_draft_tokens, static_extend_len=True, + index_kpool=self.dsa_index_kpool, ) cache_seqlens = metadata.cache_seqlens_int32 seqlens_expanded = metadata.dsa_seqlens_expanded[:total_extend_len] @@ -1818,6 +1861,43 @@ def init_forward_metadata_replay_cuda_graph_from_precomputed( metadata = self.decode_cuda_graph_metadata[bs] + self._copy_base_replay_buffers(bs, metadata, precomputed, forward_mode) + + # Refresh the schedule because stale shape decomposition can deadlock + # DeepGEMM paged MQA. + if is_cuda(): + if forward_mode.is_decode_or_idle(): + seqlens_32_2d = _to_2d_context_lens(metadata.cache_seqlens_int32, bs) + else: + seqlens_32_2d = self._build_paged_mqa_schedule_2d_ctx_lens( + forward_mode, + metadata.cache_seqlens_int32, + metadata.dsa_seqlens_expanded, + bs, + ) + self._refresh_paged_mqa_schedule_metadata(metadata, seqlens_32_2d) + self._refresh_topk_v2_plan(metadata) + if metadata.paged_mqa_ctx_lens_2d is None: + object.__setattr__(metadata, "paged_mqa_ctx_lens_2d", seqlens_32_2d) + else: + metadata.paged_mqa_ctx_lens_2d.copy_(seqlens_32_2d) + + self._update_kpool_metadata_from_precomputed( + metadata, precomputed, forward_mode + ) + + self.forward_metadata = metadata + + def _copy_base_replay_buffers( + self, + bs: int, + metadata: DSAMetadata, + precomputed: PrecomputedMetadata, + forward_mode: ForwardMode, + ) -> None: + """Copy the captured base buffers of `metadata` from `precomputed` + (fused CUDA kernel with a per-tensor fallback). Shared between the + full replay path above and `_copy_replay_metadata_from_sibling`.""" # Track whether fused kernel succeeded fused_kernel_succeeded = False @@ -1938,28 +2018,86 @@ def init_forward_metadata_replay_cuda_graph_from_precomputed( flashmla_metadata = metadata.flashmla_metadata.slice(slice(0, size + 1)) flashmla_metadata.copy_(precomputed.flashmla_metadata) - # Refresh the schedule because stale shape decomposition can deadlock - # DeepGEMM paged MQA. + @staticmethod + def _sibling_replay_metadata_compatible(dst: DSAMetadata, src: DSAMetadata) -> bool: + """Check that both sides expose the same optional derived buffers.""" + + def _match(a, b) -> bool: + return (a is None) == (b is None) + + if not ( + _match(dst.paged_mqa_schedule_metadata, src.paged_mqa_schedule_metadata) + and _match(dst.topk_v2_plan, src.topk_v2_plan) + and _match(dst.pooled_cache_seqlens_int32, src.pooled_cache_seqlens_int32) + and _match(dst.pooled_real_page_table, src.pooled_real_page_table) + and _match( + dst.pooled_paged_mqa_schedule_metadata, + src.pooled_paged_mqa_schedule_metadata, + ) + and _match(dst.kpool_write_plan, src.kpool_write_plan) + ): + return False + dst_plan, src_plan = dst.kpool_write_plan, src.kpool_write_plan + if dst_plan is not None and not ( + _match(dst_plan.pool_seqlens_per_q, src_plan.pool_seqlens_per_q) + and _match(dst_plan.seqlens_per_q, src_plan.seqlens_per_q) + and _match(dst_plan.pool_schedule_metadata, src_plan.pool_schedule_metadata) + and _match(dst_plan.effective_n_per_batch, src_plan.effective_n_per_batch) + ): + return False + return True + + def _copy_replay_metadata_from_sibling( + self, + src_backend: DeepseekSparseAttnBackend, + bs: int, + precomputed: PrecomputedMetadata, + forward_mode: ForwardMode, + ) -> None: + """Copy replay metadata from a sibling using the same precomputed input.""" + metadata = self.decode_cuda_graph_metadata.get(bs) + src_metadata = src_backend.decode_cuda_graph_metadata.get(bs) + if ( + # The derived-copy body below is CUDA-only; any other platform + # must take the full recompute, not a partial copy that would + # leave the DeepGEMM schedule / top-k plan / kpool metadata + # stale. + not is_cuda() + or _is_hip + or not forward_mode.is_decode_or_idle() + or metadata is None + or src_metadata is None + # `src_backend` must have run the full recompute path for this bs + # in this replay, so its derived buffers are fresh. + or src_backend.forward_metadata is not src_metadata + or not self._sibling_replay_metadata_compatible(metadata, src_metadata) + ): + self.init_forward_metadata_replay_cuda_graph_from_precomputed( + bs=bs, precomputed=precomputed, forward_mode=forward_mode + ) + return + + self.set_dsa_prefill_impl(forward_batch=None) + self._copy_base_replay_buffers(bs, metadata, precomputed, forward_mode) + if is_cuda(): - if forward_mode.is_decode_or_idle(): - seqlens_32_2d = _to_2d_context_lens(metadata.cache_seqlens_int32, bs) - else: - seqlens_32_2d = self._build_paged_mqa_schedule_2d_ctx_lens( - forward_mode, - metadata.cache_seqlens_int32, - metadata.dsa_seqlens_expanded, - bs, + if metadata.paged_mqa_schedule_metadata is not None: + metadata.paged_mqa_schedule_metadata.copy_( + src_metadata.paged_mqa_schedule_metadata ) - self._refresh_paged_mqa_schedule_metadata(metadata, seqlens_32_2d) - self._refresh_topk_v2_plan(metadata) + if metadata.topk_v2_plan is not None: + metadata.topk_v2_plan.copy_(src_metadata.topk_v2_plan) + # Decode: the 2D ctx lens are a (bs, 1) view of this backend's own + # cache_seqlens_int32 (just refreshed by the base copy above); keep + # the exact refresh the recompute path performs -- it is a single + # small view/copy, not part of the duplicated derived work. + seqlens_32_2d = _to_2d_context_lens(metadata.cache_seqlens_int32, bs) if metadata.paged_mqa_ctx_lens_2d is None: object.__setattr__(metadata, "paged_mqa_ctx_lens_2d", seqlens_32_2d) else: metadata.paged_mqa_ctx_lens_2d.copy_(seqlens_32_2d) - self._update_kpool_metadata_from_precomputed( - metadata, precomputed, forward_mode - ) + self._copy_kpool_metadata_from_sibling(metadata, src_metadata) self.forward_metadata = metadata @@ -3744,6 +3882,12 @@ def init_forward_metadata(self, forward_batch: ForwardBatch): for i in range(self.speculative_num_steps - 1): self.attn_backends[i].init_forward_metadata(forward_batch) + def _multistep_dedup_enabled(self) -> bool: + """The effective fusion gate includes geometry and platform support.""" + return getattr( + self.attn_backends[0], "experimental_kpool_metadata_fusion", False + ) + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): for i in range(self.speculative_num_steps - 1): self.attn_backends[i].init_cuda_graph_state(max_bs, max_num_tokens) @@ -3915,15 +4059,31 @@ def init_forward_metadata_out_graph( forward_mode=ForwardMode.DECODE, ) else: - # Copy to each backend and refresh its derived metadata independently. - for i in range(self.speculative_num_steps - 1): - self.attn_backends[ - i - ].init_forward_metadata_replay_cuda_graph_from_precomputed( - bs=bs, - precomputed=precomputed, - forward_mode=ForwardMode.DECODE, - ) + # Backend 0 fully refreshes; later backends may copy its derived metadata. + dedup = self._multistep_dedup_enabled() + self.attn_backends[ + 0 + ].init_forward_metadata_replay_cuda_graph_from_precomputed( + bs=bs, + precomputed=precomputed, + forward_mode=ForwardMode.DECODE, + ) + for i in range(1, self.speculative_num_steps - 1): + if dedup: + self.attn_backends[i]._copy_replay_metadata_from_sibling( + src_backend=self.attn_backends[0], + bs=bs, + precomputed=precomputed, + forward_mode=ForwardMode.DECODE, + ) + else: + self.attn_backends[ + i + ].init_forward_metadata_replay_cuda_graph_from_precomputed( + bs=bs, + precomputed=precomputed, + forward_mode=ForwardMode.DECODE, + ) def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None: for i in range(self.speculative_num_steps - 1): diff --git a/test/registered/kernels/ops/attention/test_dsa_metadata.py b/test/registered/kernels/ops/attention/test_dsa_metadata.py index 971bf0df8eeb..54d9d1f498a7 100644 --- a/test/registered/kernels/ops/attention/test_dsa_metadata.py +++ b/test/registered/kernels/ops/attention/test_dsa_metadata.py @@ -21,7 +21,17 @@ def _cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor: return out -def _dsa_seqlens(seqlens: torch.Tensor, topk: int) -> torch.Tensor: +def _dsa_seqlens( + seqlens: torch.Tensor, topk: int, index_kpool: int = 1 +) -> torch.Tensor: + if index_kpool > 1: + full_pool_tokens = ( + torch.div(seqlens, index_kpool, rounding_mode="floor") * index_kpool + ) + return ( + torch.minimum(full_pool_tokens, torch.tensor(topk, device=seqlens.device)) + + seqlens % index_kpool + ) return torch.minimum( seqlens.to(torch.int32), torch.tensor(topk, device=seqlens.device) ) @@ -60,6 +70,7 @@ def _check_decode( max_len: int, dsa_index_topk: int, real_page_size: int, + index_kpool: int = 1, ): bs = len(seq_lens_values) pool_size = max(bs + 3, 8) @@ -96,11 +107,12 @@ def _check_decode( max_len=max_len, dsa_index_topk=dsa_index_topk, real_page_size=real_page_size, + index_kpool=index_kpool, ) expected_cache = seq_lens.to(torch.int32) expected_page_table = req_to_token[req_pool_indices, :max_len].contiguous() - expected_dsa = _dsa_seqlens(expected_cache, dsa_index_topk) + expected_dsa = _dsa_seqlens(expected_cache, dsa_index_topk, index_kpool) # Compare only the live prefix [:seq_len]: whole blocks starting past # the kv length are skipped (keep stale values), while the last @@ -142,6 +154,7 @@ def _check_target_verify( real_page_size: int, next_n: int, fill_ctx_lens: bool, + index_kpool: int = 1, ): bs = len(seq_lens_values) expanded_size = bs * next_n @@ -198,6 +211,7 @@ def _check_target_verify( dsa_index_topk=dsa_index_topk, real_page_size=real_page_size, next_n=next_n, + index_kpool=index_kpool, paged_mqa_ctx_lens_2d=paged_mqa_ctx_lens_2d, ) @@ -209,7 +223,7 @@ def _check_target_verify( draft_offsets = torch.arange(next_n, dtype=torch.int32, device=self.device) expected_expanded = seq_lens.to(torch.int32).view(-1, 1) + draft_offsets + 1 expected_expanded = expected_expanded.reshape(-1).contiguous() - expected_dsa = _dsa_seqlens(expected_expanded, dsa_index_topk) + expected_dsa = _dsa_seqlens(expected_expanded, dsa_index_topk, index_kpool) # Compare only the live prefix [:seq_len + next_n] per expanded row: # whole blocks starting past the kv length are skipped, the last @@ -260,6 +274,7 @@ def _check_draft_extend( max_extend_len: int, max_total_len: int, static_extend_len: bool, + index_kpool: int = 1, ): bs = len(seq_lens_values) total_len = sum(extend_seq_lens_values) @@ -315,6 +330,7 @@ def _check_draft_extend( max_extend_len=max_extend_len, max_total_len=max_total_len, static_extend_len=static_extend_len, + index_kpool=index_kpool, ) expected_cache = seq_lens.to(torch.int32) @@ -337,7 +353,7 @@ def _check_draft_extend( if expanded_parts else torch.empty(0, dtype=torch.int32, device=self.device) ) - expected_dsa = _dsa_seqlens(expected_expanded, dsa_index_topk) + expected_dsa = _dsa_seqlens(expected_expanded, dsa_index_topk, index_kpool) # Compare only the live prefix [:kv_len]: whole blocks starting past # kv_len are skipped, the last partially live block may still write @@ -441,6 +457,38 @@ def test_draft_extend_partial_fill(self): static_extend_len=False, ) + def test_kpool_live_tails_across_pool_page_and_topk_boundaries(self): + for index_kpool in (2, 4): + with self.subTest(index_kpool=index_kpool): + lengths = [1, 3, 4, 63, 64, 65, 255, 256, 257, 2048, 2049, 4099] + self._check_decode( + lengths, + max_len=4609, + dsa_index_topk=2048, + real_page_size=64, + index_kpool=index_kpool, + ) + self._check_target_verify( + lengths, + max_seqlen_k=4609, + dsa_index_topk=2048, + real_page_size=64, + next_n=4, + fill_ctx_lens=True, + index_kpool=index_kpool, + ) + self._check_draft_extend( + [66, 257, 2050, 4099], + [4, 4, 4, 4], + max_seqlen_k=4609, + dsa_index_topk=2048, + real_page_size=64, + max_extend_len=4, + max_total_len=16, + static_extend_len=True, + index_kpool=index_kpool, + ) + def test_empty_batch(self): self._check_decode( [], diff --git a/test/registered/unit/layers/attention/test_kpool_metadata_fusion_contract.py b/test/registered/unit/layers/attention/test_kpool_metadata_fusion_contract.py new file mode 100644 index 000000000000..95dabd625d49 --- /dev/null +++ b/test/registered/unit/layers/attention/test_kpool_metadata_fusion_contract.py @@ -0,0 +1,69 @@ +"""CPU checks for the opt-in metadata-fusion envelope and reuse guards.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from sglang.srt.layers.attention.dsa.dsa_backend_kpool import ( + _is_kpool_metadata_fusion_supported, +) +from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class TestFusionContract(unittest.TestCase): + def test_only_supported_pool_page_topk_geometry_is_enabled(self): + for pool, page, topk, expected in ( + (1, 64, 2048, False), + (2, 64, 2048, True), + (4, 64, 2048, True), + (3, 64, 2048, False), + (4, 128, 2048, False), + (4, 64, 2049, False), + ): + with self.subTest(pool=pool, page=page, topk=topk): + self.assertEqual( + _is_kpool_metadata_fusion_supported(pool, page, topk), expected + ) + + def test_optional_derived_buffer_mismatch_rejects_reuse(self): + fields = dict( + paged_mqa_schedule_metadata=None, + topk_v2_plan=None, + pooled_cache_seqlens_int32=None, + pooled_real_page_table=None, + pooled_paged_mqa_schedule_metadata=None, + kpool_write_plan=None, + ) + dst, src = SimpleNamespace(**fields), SimpleNamespace(**fields) + self.assertTrue( + DeepseekSparseAttnBackend._sibling_replay_metadata_compatible(dst, src) + ) + src.topk_v2_plan = object() + self.assertFalse( + DeepseekSparseAttnBackend._sibling_replay_metadata_compatible(dst, src) + ) + + def test_non_cuda_recomputes_instead_of_partial_copy(self): + dst = SimpleNamespace( + decode_cuda_graph_metadata={}, + init_forward_metadata_replay_cuda_graph_from_precomputed=MagicMock(), + ) + src = SimpleNamespace(decode_cuda_graph_metadata={}) + precomputed = object() + with patch( + "sglang.srt.layers.attention.dsa_backend.is_cuda", return_value=False + ): + DeepseekSparseAttnBackend._copy_replay_metadata_from_sibling( + dst, src, 2, precomputed, ForwardMode.DECODE + ) + dst.init_forward_metadata_replay_cuda_graph_from_precomputed.assert_called_once_with( + bs=2, precomputed=precomputed, forward_mode=ForwardMode.DECODE + ) + + +if __name__ == "__main__": + unittest.main() From 342fba81d1c7fcefd7a58fedab957a63cd719e91 Mon Sep 17 00:00:00 2001 From: David Orman Date: Mon, 7 Sep 2026 14:30:40 -0500 Subject: [PATCH 2/2] Adapt regression coverage and documentation to current CI --- .../attention/dsa/dsa_backend_mtp_precompute.py | 5 ++++- python/sglang/srt/layers/attention/dsa_backend.py | 15 ++++++++++++--- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py index aa9a73e82dfc..97434438bed4 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py @@ -124,7 +124,10 @@ def _precompute_decode_mode( if (_is_cuda or _is_hip) and ( self.dsa_index_kpool <= 1 - or (not _is_hip and getattr(self, "experimental_kpool_metadata_fusion", False)) + or ( + not _is_hip + and getattr(self, "experimental_kpool_metadata_fusion", False) + ) ): from sglang.kernels.ops.attention.dsa_metadata import ( fused_dsa_decode_metadata, diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 2d2b2b57fe74..6e4bce7ea501 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -1563,7 +1563,10 @@ def _apply_cuda_graph_metadata( if (is_cuda() or _is_hip) and ( self.dsa_index_kpool <= 1 - or (not _is_hip and getattr(self, "experimental_kpool_metadata_fusion", False)) + or ( + not _is_hip + and getattr(self, "experimental_kpool_metadata_fusion", False) + ) ): fused_dsa_decode_metadata( seq_lens=seq_lens, @@ -1607,7 +1610,10 @@ def _apply_cuda_graph_metadata( if (is_cuda() or _is_hip) and ( self.dsa_index_kpool <= 1 - or (not _is_hip and getattr(self, "experimental_kpool_metadata_fusion", False)) + or ( + not _is_hip + and getattr(self, "experimental_kpool_metadata_fusion", False) + ) ): paged_mqa_ctx_lens_2d = None if ( @@ -1707,7 +1713,10 @@ def _apply_cuda_graph_metadata( if (is_cuda() or _is_hip) and ( self.dsa_index_kpool <= 1 - or (not _is_hip and getattr(self, "experimental_kpool_metadata_fusion", False)) + or ( + not _is_hip + and getattr(self, "experimental_kpool_metadata_fusion", False) + ) ): fused_dsa_draft_extend_metadata( seq_lens=seq_lens,