From c7ca332fb0e6f5d05986a57afbc6dbfaae62abe7 Mon Sep 17 00:00:00 2001 From: weireweire Date: Wed, 8 Jul 2026 11:35:33 +0800 Subject: [PATCH] Fix DSV4 prefill large Triton recompilation idle across context lengths (#30255) Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com> --- .../layers/attention/dsv4/metadata_kernel.py | 24 +++++++++---------- .../attention/dsv4/sparse_prefill_utils.py | 15 +++++++----- 2 files changed, 21 insertions(+), 18 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsv4/metadata_kernel.py b/python/sglang/srt/layers/attention/dsv4/metadata_kernel.py index 5e4b999f7..ce6b038e7 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata_kernel.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata_kernel.py @@ -5,7 +5,7 @@ import triton import triton.language as tl -@triton.jit +@triton.jit(do_not_specialize=["bs", "c128_cur_max_seq_len"]) def _init_compressed_attn_metadata_kernel( seq_lens_ptr, positions_ptr, @@ -22,8 +22,7 @@ def _init_compressed_attn_metadata_kernel( c128_page_indices_ptr, bs, max_pages, - page_size: tl.constexpr, - c128_max_seq_len: tl.constexpr, + c128_cur_max_seq_len, c128_page_size: tl.constexpr, BLOCK_SIZE: tl.constexpr, COMPUTE_PAGE_INDICES: tl.constexpr, @@ -59,10 +58,10 @@ def _init_compressed_attn_metadata_kernel( tl.store(c128_seq_lens_clamp1_ptr + batch_id, c128_seq_lens_clamp1) if COMPUTE_PAGE_INDICES: - page_indices_base = batch_id * c128_max_seq_len - for block_start in range(0, c128_max_seq_len, BLOCK_SIZE): + page_indices_base = batch_id * c128_cur_max_seq_len + for block_start in tl.range(0, c128_cur_max_seq_len, BLOCK_SIZE): offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < c128_max_seq_len + mask = offsets < c128_cur_max_seq_len page_idx = offsets // c128_page_size offset_in_page = offsets % c128_page_size @@ -121,18 +120,20 @@ def _init_compressed_attn_metadata_triton( assert ( page_table is not None ), "page_table required when compute_page_indices=True" - assert page_size > 0, "page_size required when compute_page_indices=True" + assert ( + page_size >= 128 and page_size % 128 == 0 + ), "page_size must be a multiple of 128 when compute_page_indices=True" max_pages = page_table.shape[1] c128_page_size = page_size // 128 - c128_max_seq_len = c128_page_size * max_pages + c128_cur_max_seq_len = c128_page_size * max_pages c128_page_indices = torch.empty( - bs, c128_max_seq_len, dtype=torch.int32, device=device + bs, c128_cur_max_seq_len, dtype=torch.int32, device=device ) BLOCK_SIZE = triton.next_power_of_2(max(c128_page_size, 64)) else: max_pages = 0 c128_page_size = 1 - c128_max_seq_len = 0 + c128_cur_max_seq_len = 0 c128_page_indices = None BLOCK_SIZE = 64 if page_table is None: @@ -159,8 +160,7 @@ def _init_compressed_attn_metadata_triton( ), bs, max_pages, - page_size if page_size > 0 else 128, - c128_max_seq_len, + c128_cur_max_seq_len, c128_page_size, BLOCK_SIZE, compute_page_indices, diff --git a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py index be5ef1915..351263cb0 100644 --- a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py +++ b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py @@ -135,7 +135,10 @@ def combine_topk_swa_indices( assert gather_lens.dtype == torch.int32 assert compressed_base.dtype == torch.int32 assert swa_base.dtype == torch.int32 - assert compress_ratio >= 1, "COMPRESS_RATIO must be >= 1 (use TOP_K=0 for SWA-only)" + assert compress_ratio >= 1, "compress_ratio must be >= 1 (use topk=0 for SWA-only)" + assert ( + topk_indices.shape[-1] >= topk + ), f"topk_indices width {topk_indices.shape[-1]} must be >= topk {topk}" num_tokens = topk_indices.shape[0] num_reqs = seq_lens.shape[0] @@ -172,7 +175,7 @@ def combine_topk_swa_indices( gather_lens, compressed_base, swa_base, - TOP_K=topk, + top_k=topk, COMPRESS_RATIO=compress_ratio, WINDOW_SIZE=window_size, PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]), @@ -279,7 +282,7 @@ def _build_swa_token_ids_kernel( tl.store(out_ptr + out_off + i, swa_id) -@triton.jit +@triton.jit(do_not_specialize=["top_k"]) def _combine_topk_swa_indices_kernel( combined_indices_ptr, combined_indices_stride, @@ -291,7 +294,7 @@ def _combine_topk_swa_indices_kernel( gather_lens_ptr, compressed_base_ptr, swa_base_ptr, - TOP_K: tl.constexpr, + top_k, COMPRESS_RATIO: tl.constexpr, WINDOW_SIZE: tl.constexpr, PADDED_TOP_K: tl.constexpr, @@ -321,8 +324,8 @@ def _combine_topk_swa_indices_kernel( pos = start_pos + token_idx_in_query # Both the C4 indexer and the C128 metadata builder emit # min((pos+1)//compress_ratio, topk_tokens) valid entries. Caller - # passes TOP_K=0 for SWA-only layers to zero this out. - topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) + # passes top_k=0 for SWA-only layers to zero this out. + topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, top_k) swa_len = tl.minimum(pos + 1, WINDOW_SIZE) combined_row = token_idx.to(tl.int64) * combined_indices_stride