From 18e6c61c21ad39725522c008190d2b540dd6228d Mon Sep 17 00:00:00 2001 From: valechen <115046356+valechen@users.noreply.github.com> Date: Thu, 6 Aug 2026 14:46:10 -0700 Subject: [PATCH] [AMD] perf: compact Triton extend-attention for ragged prefill (AMD/HIP-only) (#29677) --- .../kernels/ops/attention/extend_attention.py | 104 +++++++++++++- .../kernels/ops/attention/verify_splitkv.py | 6 + python/sglang/srt/environ.py | 11 ++ .../srt/layers/attention/triton_backend.py | 14 ++ python/sglang/srt/managers/schedule_policy.py | 104 +++++++++++++- python/sglang/srt/managers/scheduler.py | 8 ++ .../test_triton_attention_kernels.py | 127 ++++++++++++++++++ .../attention/test_verify_splitkv.py | 14 ++ .../unit/managers/test_prefill_adder.py | 67 ++++++++- 9 files changed, 449 insertions(+), 6 deletions(-) diff --git a/python/sglang/kernels/ops/attention/extend_attention.py b/python/sglang/kernels/ops/attention/extend_attention.py index c41da1d56..82fa160e4 100644 --- a/python/sglang/kernels/ops/attention/extend_attention.py +++ b/python/sglang/kernels/ops/attention/extend_attention.py @@ -25,6 +25,7 @@ from sglang.kernels.ops.attention.prefill_attention import ( context_attention_fwd, ) from sglang.kernels.ops.attention.score_mod import unpack_aux_tensors +from sglang.srt.environ import envs from sglang.srt.utils import is_cuda, is_gfx95_supported, is_hip _is_cuda = is_cuda() @@ -126,6 +127,53 @@ def _get_block_sizes_for_extend_attention(Lq: int, Lv: int): return BLOCK_DMODEL, BLOCK_DPE, BLOCK_DV, BLOCK_M, BLOCK_N, num_warps +def _compact_extend_q_tiles_per_head( + *, + batch_size: int, + max_len_extend: int, + total_extend_tokens: int, + block_m: int, + extend_seq_lens_cpu=None, +) -> int | None: + """Return compact query tiles per head when it reduces launch work. + + The legacy extend grid is rectangular -- ``batch_size * cdiv(max_len_extend, + BLOCK_M)`` -- so in a ragged mixed-prefill batch every short row pays tile + work sized by the longest row. This computes the *compact* tile count + (``sum_i cdiv(extend_len_i, BLOCK_M)``), i.e. work proportional to the real + per-request lengths. That is the same ragged-aware launch the flash-attn + varlen kernels (used by the aiter backend via ``flash_attn_varlen_func`` / + ``mha_batch_prefill_func``) already get from their cu_seqlens scheduler -- + this closes that triton-vs-flash-attn gap rather than inventing a new + technique. Returns ``None`` (keep the legacy grid) when compacting would not + reduce launch work, e.g. a uniform batch. + """ + if batch_size <= 1 or max_len_extend <= 0: + return None + + legacy_tiles = batch_size * triton.cdiv(max_len_extend, block_m) + if legacy_tiles <= 0: + return None + + if extend_seq_lens_cpu is not None: + if isinstance(extend_seq_lens_cpu, torch.Tensor): + extend_seq_lens_cpu = extend_seq_lens_cpu.tolist() + if len(extend_seq_lens_cpu) < batch_size: + return None + compact_tiles = sum( + triton.cdiv(max(0, int(extend_seq_lens_cpu[i])), block_m) + for i in range(batch_size) + ) + else: + if total_extend_tokens == batch_size * max_len_extend: + return None + compact_tiles = (total_extend_tokens + batch_size * (block_m - 1)) // block_m + + if compact_tiles <= 0 or compact_tiles >= legacy_tiles: + return None + return int(compact_tiles) + + @triton.jit def tanh(x): # Tanh is just a scaled sigmoid @@ -277,6 +325,7 @@ def _fwd_kernel( stride_buf_ktok, stride_buf_vpage, stride_buf_vtok, + compact_batch_size, SLIDING_WINDOW_SIZE: tl.constexpr, logit_cap: tl.constexpr, xai_temperature_len: tl.constexpr, @@ -295,6 +344,7 @@ def _fwd_kernel( SKIP_EXTEND: tl.constexpr, STORE_TRANSPOSE: tl.constexpr, HAS_SINK: tl.constexpr, + USE_COMPACT_TILE_GRID: tl.constexpr, PAGE_SIZE: tl.constexpr = 1, SCORE_MOD: tl.constexpr = None, Aux0=None, @@ -302,9 +352,33 @@ def _fwd_kernel( aux0_stride_h=0, aux0_len=0, ): - cur_seq = tl.program_id(0) - cur_head = tl.program_id(1) - cur_block_m = tl.program_id(2) + if USE_COMPACT_TILE_GRID: + output_tile = tl.program_id(0) + cur_head = tl.program_id(1) + + cur_seq = tl.full((), 0, tl.int64) + cum_tiles = tl.full((), 0, tl.int64) + found = tl.full((), 0, tl.int32) + while (cur_seq < compact_batch_size) & (found == 0): + seq_q_start = tl.load(qo_indptr + cur_seq) + seq_q_end = tl.load(qo_indptr + cur_seq + 1) + seq_q_len = seq_q_end - seq_q_start + seq_tiles = (seq_q_len + BLOCK_M - 1) // BLOCK_M + next_cum_tiles = cum_tiles + seq_tiles + if next_cum_tiles > output_tile: + found = 1 + else: + cum_tiles = next_cum_tiles + cur_seq = cur_seq + 1 + + if found == 0: + return + + cur_block_m = output_tile - cum_tiles + else: + cur_seq = tl.program_id(0) + cur_head = tl.program_id(1) + cur_block_m = tl.program_id(2) cur_kv_head = cur_head // kv_group_num cur_seq_extend_start_idx = tl.load(qo_indptr + cur_seq) @@ -690,6 +764,7 @@ def extend_attention_fwd( page_size: int = 1, score_mod=None, aux_tensors=None, + extend_seq_lens_cpu=None, ): """ q_extend, k_extend, v_extend, o_extend: contiguous tensors @@ -727,7 +802,26 @@ def extend_attention_fwd( stride_lse_bs = lse_extend.stride(0) if STORE_LSE else 0 stride_lse_h = lse_extend.stride(1) if STORE_LSE else 0 - grid = (batch_size, head_num, triton.cdiv(max_len_extend, BLOCK_M)) + # Compact grid: AMD/HIP-only optimization (parity with flash-attn's ragged-aware + # launch). Explicitly check _is_hip and allow env var override. + use_compact_tile_grid = ( + _is_hip and envs.SGLANG_TRITON_COMPACT_EXTEND_ATTENTION.get() + ) + compact_q_tiles = None + if use_compact_tile_grid: + compact_q_tiles = _compact_extend_q_tiles_per_head( + batch_size=batch_size, + max_len_extend=max_len_extend, + total_extend_tokens=q_extend.shape[0], + block_m=BLOCK_M, + extend_seq_lens_cpu=extend_seq_lens_cpu, + ) + + use_compact_tile_grid = compact_q_tiles is not None + if use_compact_tile_grid: + grid = (compact_q_tiles, head_num) + else: + grid = (batch_size, head_num, triton.cdiv(max_len_extend, BLOCK_M)) num_stages = 1 extra_kargs = {} @@ -782,6 +876,7 @@ def extend_attention_fwd( k_tok_stride, v_page_stride, v_tok_stride, + batch_size, SLIDING_WINDOW_SIZE=sliding_window_size, logit_cap=logit_cap, xai_temperature_len=xai_temperature_len, @@ -800,6 +895,7 @@ def extend_attention_fwd( SKIP_EXTEND=skip_extend, HAS_SINK=HAS_SINK, STORE_TRANSPOSE=_is_hip, + USE_COMPACT_TILE_GRID=use_compact_tile_grid, PAGE_SIZE=page_size, SCORE_MOD=score_mod, Aux0=aux0, diff --git a/python/sglang/kernels/ops/attention/verify_splitkv.py b/python/sglang/kernels/ops/attention/verify_splitkv.py index e262124c6..94bb8526f 100644 --- a/python/sglang/kernels/ops/attention/verify_splitkv.py +++ b/python/sglang/kernels/ops/attention/verify_splitkv.py @@ -707,6 +707,12 @@ def can_handle( return False if v_extend.shape[2] != v_buffer.shape[2]: return False + # MLA (head_dim != v_head_dim, e.g. DeepSeek 576 vs 512) uses a shared + # latent KV cache and an absorbed-attention layout the split-KV verify + # kernel is not built for; it GPU-faults on that shape. Fall back to + # extend_attention_fwd, which handles MLA correctly. + if q_extend.shape[2] != v_extend.shape[2]: + return False # NOTE: must NOT read any tensor *values* here (no .item()/.cpu()): the # target-verify step runs inside a captured CUDA/HIP graph, where a # device->host sync raises hipErrorStreamCaptureUnsupported. We therefore diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 33cd636ba..9d5987f25 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -469,6 +469,13 @@ class Envs: # an accurate TTFT for benchmarking; the upstream default of 50 trades # off some TTFT-metric accuracy for less IPC overhead. SGLANG_FORCE_STREAM_INTERVAL = EnvInt(50) + # Compact extend-attention scheduler tile-budget admission (AMD/HIP-only). + # Budget <= 0 disables; >0 sets the max prefix-extend tiles per batch. + SGLANG_PREFILL_TILE_BUDGET = EnvInt(0) + # Tile-budget mode: "compact" (default, counts actual per-request tiles) or + # "legacy" (rectangular grid, max_extend_len-shaped). + # Internal/testing only - users should not need to change this. + SGLANG_PREFILL_TILE_BUDGET_MODE = EnvStr("compact") # Test: pd-disaggregation SGLANG_TEST_PD_DISAGG_BACKEND = EnvStr("mooncake") @@ -710,6 +717,10 @@ class Envs: # Triton SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS = EnvBool(False) SGLANG_USE_CUSTOM_TRITON_KERNEL_CACHE = EnvBool(False) + # Compact extend-attention query-tile grid: AMD/HIP-only optimization + # (parity with flash-attn's ragged-aware launch). The feature checks _is_hip + # explicitly in code; this env var allows override (0=force off, 1=force on). + SGLANG_TRITON_COMPACT_EXTEND_ATTENTION = EnvBool(True) # Torch Compile SGLANG_ENABLE_TORCH_COMPILE = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index c378c80ab..c1f921ce1 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -301,6 +301,19 @@ class TritonAttnBackend(AttentionBackend): # Tree-mask scratch is fetched from the target backend only. self.is_draft_runner = model_runner.is_draft_worker + # Auto-detect BLOCK_M that extend_attention kernel will use for this model. + # This is used by the scheduler's tile-budget admission logic to match + # the kernel's actual tile size. + head_dim = model_runner.model_config.head_dim + from sglang.kernels.ops.attention.extend_attention import ( + _get_block_sizes_for_extend_attention, + ) + + _, _, _, block_m, _, _ = _get_block_sizes_for_extend_attention( + Lq=head_dim, Lv=head_dim + ) + self.extend_attention_block_m = block_m + def get_num_kv_splits( self, num_kv_splits: torch.Tensor, @@ -1433,6 +1446,7 @@ class TritonAttnBackend(AttentionBackend): page_size=self.page_size, score_mod=score_mod, aux_tensors=aux_tensors, + extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu, ) return o diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 4d79fc88e..797a8dfcf 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -6,7 +6,7 @@ from array import array from sglang.srt.environ import envs from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor from sglang.srt.runtime_context import get_disagg -from sglang.srt.utils import get_bool_env_var +from sglang.srt.utils import get_bool_env_var, is_hip _ROUTING_KEY_POLICY_DEBUG_LOG = get_bool_env_var("SGLANG_ROUTING_KEY_POLICY_DEBUG_LOG") logger = logging.getLogger(__name__) @@ -91,6 +91,48 @@ IN_BATCH_PREFIX_CACHING_DEPRIORITIZE_THRESHOLD = int( IGNORE_EOS_RESERVE_TOKENS = 1 +# AMD/HIP-only: the prefill tile-budget admission control is part of the AMD +# compact extend-attention work and is gated on HIP (see _check_prefill_tile_budget), +# so non-AMD vendors keep the exact legacy scheduler behavior. +_IS_HIP = is_hip() +PREFILL_TILE_BUDGET = envs.SGLANG_PREFILL_TILE_BUDGET.get() +PREFILL_TILE_BUDGET_MODE = envs.SGLANG_PREFILL_TILE_BUDGET_MODE.get().strip().lower() +if PREFILL_TILE_BUDGET_MODE not in {"legacy", "compact"}: + logger.warning( + "Unsupported SGLANG_PREFILL_TILE_BUDGET_MODE=%s. Falling back to compact.", + PREFILL_TILE_BUDGET_MODE, + ) + PREFILL_TILE_BUDGET_MODE = "compact" + + +def _ceil_div(value: int, divisor: int) -> int: + return -(-value // divisor) + + +def estimate_prefill_extend_tile_metrics( + extend_lens: List[int], block_m: int +) -> Dict[str, Union[int, float, List[int], None]]: + """Estimate extend-attention query tiles per head for a prefill batch.""" + normalized_lens = [max(0, int(length)) for length in extend_lens] + q_tiles = [ + _ceil_div(length, block_m) if length > 0 else 0 for length in normalized_lens + ] + legacy_tiles = len(q_tiles) * max(q_tiles) if q_tiles else 0 + compact_tiles = sum(q_tiles) + saved_tiles = legacy_tiles - compact_tiles + saved_ratio = saved_tiles / legacy_tiles if legacy_tiles else None + return { + "block_m": int(block_m), + "request_count": len(normalized_lens), + "extend_lens": normalized_lens, + "q_tiles_per_request": q_tiles, + "max_extend_len": max(normalized_lens) if normalized_lens else 0, + "sum_extend_len": sum(normalized_lens), + "legacy_q_tiles_per_head": legacy_tiles, + "compact_q_tiles_per_head": compact_tiles, + "saved_q_tiles_per_head": saved_tiles, + "saved_q_tile_ratio": saved_ratio, + } def match_prefix_for_req( @@ -463,8 +505,10 @@ class PrefillAdder: prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None, dllm_config: Optional[DllmConfig] = None, waiting_queue_len: int = 0, + prefill_tile_block_m: int = 64, ): self.page_size = page_size + self.prefill_tile_block_m = prefill_tile_block_m self.tree_cache = tree_cache self.token_to_kv_pool_allocator = token_to_kv_pool_allocator self.running_batch = running_batch @@ -557,6 +601,36 @@ class PrefillAdder: # prefill pass. Used by PrefillDelayer's queue-based trigger. self.waiting_queue_len = waiting_queue_len + def _admitted_extend_lens(self) -> List[int]: + return [int(getattr(req, "extend_input_len", 0)) for req in self.can_run_list] + + def _tile_admission_metric_key(self) -> str: + return f"{PREFILL_TILE_BUDGET_MODE}_q_tiles_per_head" + + def _candidate_tile_metrics(self, candidate_extend_len: int) -> Dict[str, object]: + return estimate_prefill_extend_tile_metrics( + [*self._admitted_extend_lens(), int(candidate_extend_len)], + block_m=self.prefill_tile_block_m, + ) + + def _check_prefill_tile_budget( + self, candidate_extend_len: int + ) -> Optional[AddReqResult]: + # AMD-only: leave non-AMD scheduler admission unchanged even if the env + # budget is set. + if not _IS_HIP or PREFILL_TILE_BUDGET <= 0: + return None + + if not self.can_run_list: + return None + + metrics = self._candidate_tile_metrics(candidate_extend_len) + candidate_metric = int(metrics.get(self._tile_admission_metric_key()) or 0) + if candidate_metric <= PREFILL_TILE_BUDGET: + return None + + return AddReqResult.OTHER + def _init_dllm_meta(self, dllm_config: DllmConfig): self.dllm_block_size = dllm_config.block_size max_running_reqs = dllm_config.max_running_requests @@ -1057,11 +1131,21 @@ class PrefillAdder: if self.rem_dllm_tokens <= 0: return AddReqResult.OTHER + if ( + tile_stop := self._check_prefill_tile_budget(cand_extend_input_len) + ) is not None: + return tile_stop + self._add_dllm_req(req, 0) elif ( self.rem_chunk_tokens is None # chunked prefill is disabled or cand_extend_input_len <= self.rem_chunk_tokens # it is the last chunk ): + if ( + tile_stop := self._check_prefill_tile_budget(cand_extend_input_len) + ) is not None: + return tile_stop + # Non-chunked prefill — the whole sequence is committed this iter. req.set_extend_range( len(req.prefix_indices), len(req.full_untruncated_fill_ids) @@ -1081,6 +1165,9 @@ class PrefillAdder: # Chunked prefill trunc_len = self.rem_chunk_tokens + if (tile_stop := self._check_prefill_tile_budget(trunc_len)) is not None: + return tile_stop + assert len(req.prefix_indices) == 0 req.set_extend_range( len(req.prefix_indices), len(req.prefix_indices) + trunc_len @@ -1243,9 +1330,19 @@ class PrefillAdder: truncation_align_size is None ), "truncation_align_size is not supported for dllm prefill" + if ( + tile_stop := self._check_prefill_tile_budget(input_tokens) + ) is not None: + return tile_stop + self._add_dllm_req(req, prefix_len) self._req_inc_lock_ref(req) elif chunk_tokens_limit is None or input_tokens <= chunk_tokens_limit: + if ( + tile_stop := self._check_prefill_tile_budget(input_tokens) + ) is not None: + return tile_stop + # Non-chunked prefill — the whole sequence is committed this iter. req.set_extend_range( len(req.prefix_indices), len(req.full_untruncated_fill_ids) @@ -1290,6 +1387,11 @@ class PrefillAdder: if trunc_len <= 0: return AddReqResult.OTHER + if ( + tile_stop := self._check_prefill_tile_budget(trunc_len) + ) is not None: + return tile_stop + # Chunked prefill req.set_extend_range( len(req.prefix_indices), len(req.prefix_indices) + trunc_len diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 33643749b..c8bdceb57 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3157,6 +3157,13 @@ class Scheduler( chunked_prefill_size = dynamic_size # Prefill policy + # Get BLOCK_M from the backend for tile-budget admission logic + attn_backend = self.tp_worker.model_runner.attn_backend + if hasattr(attn_backend, "extend_attention_block_m"): + prefill_tile_block_m = attn_backend.extend_attention_block_m + else: + prefill_tile_block_m = 64 # Fallback for non-Triton backends + adder = PrefillAdder( self.page_size, self.tree_cache, @@ -3173,6 +3180,7 @@ class Scheduler( prefill_delayer_single_pass=prefill_delayer_single_pass, dllm_config=self.dllm_config, waiting_queue_len=len(self.waiting_queue), + prefill_tile_block_m=prefill_tile_block_m, ) if self.chunked_req is not None: diff --git a/test/registered/attention/test_triton_attention_kernels.py b/test/registered/attention/test_triton_attention_kernels.py index 10d2e9717..bf8897476 100644 --- a/test/registered/attention/test_triton_attention_kernels.py +++ b/test/registered/attention/test_triton_attention_kernels.py @@ -10,6 +10,7 @@ from sglang.kernels.ops.attention.decode_attention import ( decode_attention_fwd_normal, ) from sglang.kernels.ops.attention.extend_attention import ( + _compact_extend_q_tiles_per_head, build_unified_kv_indices, extend_attention_fwd, extend_attention_fwd_unified, @@ -19,6 +20,7 @@ from sglang.kernels.ops.attention.prefill_attention import ( context_attention_fwd, ) from sglang.srt.utils import get_device +from sglang.srt.utils.common import temp_set_env from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase, is_in_amd_ci @@ -338,6 +340,131 @@ class TestTritonAttention(CustomTestCase): ea._get_block_sizes_for_extend_attention(576, 576)[3:], (64, 64, 4) ) + def test_compact_extend_attention_tile_count(self): + self.assertEqual( + _compact_extend_q_tiles_per_head( + batch_size=16, + max_len_extend=1000, + total_extend_tokens=1015, + block_m=64, + extend_seq_lens_cpu=[1] * 15 + [1000], + ), + 31, + ) + self.assertEqual( + _compact_extend_q_tiles_per_head( + batch_size=2, + max_len_extend=4224, + total_extend_tokens=5376, + block_m=64, + extend_seq_lens_cpu=[1152, 4224], + ), + 84, + ) + self.assertIsNone( + _compact_extend_q_tiles_per_head( + batch_size=4, + max_len_extend=64, + total_extend_tokens=256, + block_m=64, + extend_seq_lens_cpu=[64, 64, 64, 64], + ) + ) + + def test_extend_attention_compact_grid(self): + dtype = torch.bfloat16 + device = get_device() + B, H_Q, H_KV, D = 4, 8, 2, 64 + + b_seq_len_prefix = torch.tensor( + [8, 16, 32, 64], dtype=torch.int32, device=device + ) + b_seq_len_extend = torch.tensor( + [1, 7, 13, 129], dtype=torch.int32, device=device + ) + b_seq_len = b_seq_len_prefix + b_seq_len_extend + b_start_loc = torch.zeros((B,), dtype=torch.int32, device=device) + b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0) + b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device=device) + b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0) + + kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) + kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix, dim=0) + kv_indices = torch.empty( + (int(b_seq_len_prefix.sum().item()),), dtype=torch.int32, device=device + ) + for i in range(B): + kv_indices[int(kv_indptr[i]) : int(kv_indptr[i + 1])] = torch.arange( + int(b_start_loc[i].item()), + int(b_start_loc[i].item()) + int(b_seq_len_prefix[i].item()), + device=device, + ) + + total_token_num = int(b_seq_len.sum().item()) + extend_token_num = int(b_seq_len_extend.sum().item()) + k_buffer = torch.empty( + (total_token_num, H_KV, D), dtype=dtype, device=device + ).normal_(mean=0.1, std=0.2) + v_buffer = torch.empty( + (total_token_num, H_KV, D), dtype=dtype, device=device + ).normal_(mean=0.1, std=0.2) + k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device) + v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device) + q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device) + for i in range(B): + extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i] + extend_end_in_buffer = b_start_loc[i] + b_seq_len[i] + extend_start = b_start_loc_extend[i] + extend_end = b_start_loc_extend[i] + b_seq_len_extend[i] + k_extend[extend_start:extend_end] = k_buffer[ + extend_start_in_buffer:extend_end_in_buffer + ] + v_extend[extend_start:extend_end] = v_buffer[ + extend_start_in_buffer:extend_end_in_buffer + ] + q_extend[extend_start:extend_end] = torch.empty( + (int(b_seq_len_extend[i].item()), H_Q, D), + dtype=dtype, + device=device, + ).normal_(mean=0.1, std=0.2) + + max_len_extend = int(b_seq_len_extend.max().item()) + qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device) + qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend, dim=0) + extend_seq_lens_cpu = b_seq_len_extend.cpu().tolist() + + o_legacy = torch.empty_like(q_extend) + o_compact = torch.empty_like(q_extend) + for output, use_compact in ((o_legacy, False), (o_compact, True)): + with temp_set_env( + allow_sglang=True, + SGLANG_TRITON_COMPACT_EXTEND_ATTENTION=str(use_compact), + ): + extend_attention_fwd( + q_extend, + k_extend, + v_extend, + output, + k_buffer, + v_buffer, + qo_indptr, + kv_indptr, + kv_indices, + custom_mask=None, + is_causal=True, + mask_indptr=None, + max_len_extend=max_len_extend, + k_scale=1.0, + v_scale=1.0, + extend_seq_lens_cpu=extend_seq_lens_cpu, + ) + + self.assertTrue( + torch.allclose(o_legacy, o_compact, rtol=1e-2, atol=1e-3), + f"compact grid output differs from legacy grid. " + f"Max diff: {(o_legacy - o_compact).abs().max()}", + ) + def _test_extend_attention_sliding_window_once( self, B, N_CTX, H_Q, H_KV, D, WINDOW_SIZE ): diff --git a/test/registered/attention/test_verify_splitkv.py b/test/registered/attention/test_verify_splitkv.py index 72123bda2..fd15aba0c 100644 --- a/test/registered/attention/test_verify_splitkv.py +++ b/test/registered/attention/test_verify_splitkv.py @@ -234,6 +234,20 @@ class TestVerifySplitKV(CustomTestCase): can_handle(q, k, v, kb, vb, qo, kvp, kvi, None, True, None, mle + 1) ) + def test_fallback_mla_head_dim_mismatch(self): + # MLA (DeepSeek) has head_dim != v_head_dim (576 vs 512): the shared + # latent KV / absorbed layout is not something the split-KV verify + # kernel is built for -- it GPU-faults on that shape. can_handle() must + # reject it so the backend falls back to extend_attention_fwd. + q, k, v, kb, vb, qo, kvp, kvi, mle = _build_verify_inputs( + [512, 512], 4, 16, 1, 576, 512, torch.bfloat16, "cuda" + ) + self.assertEqual(q.shape[2], 576) + self.assertEqual(v.shape[2], 512) + self.assertFalse( + can_handle(q, k, v, kb, vb, qo, kvp, kvi, None, True, None, mle) + ) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 33943b4a2..8c63e45e9 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -2,8 +2,13 @@ import unittest from types import SimpleNamespace from unittest.mock import MagicMock, patch +import sglang.srt.managers.schedule_policy as schedule_policy from sglang.srt.managers.schedule_batch import Req -from sglang.srt.managers.schedule_policy import AddReqResult, PrefillAdder +from sglang.srt.managers.schedule_policy import ( + AddReqResult, + PrefillAdder, + estimate_prefill_extend_tile_metrics, +) from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefResult, IncLockRefResult, @@ -775,6 +780,66 @@ class TestPrefillAdder(CustomTestCase): req.set_extend_range.assert_called_once_with(0, 200) self.assertIn(req, adder.can_run_list) + def _adder_with_extend_lens(self, extend_lens): + adder = PrefillAdder.__new__(PrefillAdder) + adder.can_run_list = [ + SimpleNamespace(extend_input_len=length) for length in extend_lens + ] + # BLOCK_M is auto-detected from the attention backend in production; the + # __new__ helper bypasses __init__, so set it explicitly. 64 matches the + # block_m the tile-count assertions below are computed against. + adder.prefill_tile_block_m = 64 + return adder + + def test_estimate_prefill_extend_tile_metrics(self): + metrics = estimate_prefill_extend_tile_metrics([1, 7, 13, 129], block_m=64) + + self.assertEqual(metrics["q_tiles_per_request"], [1, 1, 1, 3]) + self.assertEqual(metrics["legacy_q_tiles_per_head"], 12) + self.assertEqual(metrics["compact_q_tiles_per_head"], 6) + self.assertEqual(metrics["saved_q_tiles_per_head"], 6) + self.assertEqual(metrics["saved_q_tile_ratio"], 0.5) + + def test_compact_prefill_tile_budget_admits_more_than_legacy(self): + adder = self._adder_with_extend_lens([1, 7, 13]) + + # The tile-budget admission is gated on HIP in production; force the gate + # on so this vendor-neutral admission-math check runs on any CI runner. + with ( + patch.object(schedule_policy, "_IS_HIP", True), + patch.object(schedule_policy, "PREFILL_TILE_BUDGET", 6), + patch.object(schedule_policy, "PREFILL_TILE_BUDGET_MODE", "compact"), + ): + self.assertIsNone(adder._check_prefill_tile_budget(129)) + + with ( + patch.object(schedule_policy, "_IS_HIP", True), + patch.object(schedule_policy, "PREFILL_TILE_BUDGET", 6), + patch.object(schedule_policy, "PREFILL_TILE_BUDGET_MODE", "legacy"), + ): + self.assertEqual(adder._check_prefill_tile_budget(129), AddReqResult.OTHER) + + def test_prefill_tile_budget_always_allows_first_request(self): + adder = self._adder_with_extend_lens([]) + + with ( + patch.object(schedule_policy, "_IS_HIP", True), + patch.object(schedule_policy, "PREFILL_TILE_BUDGET", 1), + ): + self.assertIsNone(adder._check_prefill_tile_budget(4096)) + + def test_prefill_tile_budget_disabled_on_non_hip(self): + # AMD-only: on non-HIP vendors the tile-budget admission must be a no-op + # even when the budget env is set, so scheduler behavior is unchanged. + adder = self._adder_with_extend_lens([1, 7, 13]) + + with ( + patch.object(schedule_policy, "_IS_HIP", False), + patch.object(schedule_policy, "PREFILL_TILE_BUDGET", 6), + patch.object(schedule_policy, "PREFILL_TILE_BUDGET_MODE", "legacy"), + ): + self.assertIsNone(adder._check_prefill_tile_budget(129)) + if __name__ == "__main__": unittest.main()