[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)
+    if USE_COMPACT_TILE_GRID:
-    cur_head = tl.program_id(1)
+        output_tile = tl.program_id(0)
-    cur_block_m = tl.program_id(2)
+        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:

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
     ):

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()

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()