[AMD] perf: compact Triton extend-attention for ragged prefill (AMD/HIP-only) (#29677)

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,

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

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)

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
 

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

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: