[AMD] Fix FP4 indexer OOR (#37660)

Co-authored-by: Thomas Wang <1am9trash@gmail.com>

python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple, Union
 
 import torch
@@ -27,6 +28,15 @@ _PREFILL_BASE_CTA_TARGET = 1024
 # AITER varctx cta_info row: [batch_packed, chunk_start, chunk_count, ctx_len].
 _DECODE_CTA_INFO_WIDTH = 4
 
+# Budget for the pooled prefill logits block, in MiB. Rows are split to fit it
+# (see `logits_rows_per_chunk`), so this caps the indexer's transient footprint
+# independently of context length and chunked-prefill size; smaller budgets only
+# buy more row chunks. 2 GiB covers 4096 rows over ~512K tokens of context.
+_LOGITS_BUDGET_ELEMS = (
+    int(os.environ.get("SGLANG_DSV4_FP4_LOGITS_BUDGET_MB", "2048")) * 2**20 // 4
+)
+_LOGITS_POOL: dict = {}
+
 
 class FP4DecodeWorkspace(NamedTuple):
     guarded_page_table: torch.Tensor
@@ -112,11 +122,16 @@ def _decode_cta_count(num_queries: int, max_seq_len: int) -> int:
     return min(available_ctas, target_ctas)
 
 
+def _guarded_pages(logical_width: int) -> int:
+    """Page columns after padding for 256-token scheduling."""
+    return max(4, (logical_width + 3) // 4 * 4)
+
+
 def _guard_page_table(page_table: torch.Tensor, out: Optional[torch.Tensor] = None):
     """Pad page tables for 256-token scheduling and one-chunk lookahead."""
     page_table = page_table.to(dtype=torch.int32).contiguous()
     rows, logical_width = page_table.shape
-    padded_width = max(4, (logical_width + 3) // 4 * 4)
+    padded_width = _guarded_pages(logical_width)
     if out is None:
         out = page_table.new_zeros((rows, padded_width + 4))
     else:
@@ -125,6 +140,48 @@ def _guard_page_table(page_table: torch.Tensor, out: Optional[torch.Tensor] = No
     return out, padded_width * _KV_BLOCK_SIZE
 
 
+def logits_rows_per_chunk(page_table: torch.Tensor) -> int:
+    """Rows whose logits fit the pooled block, for callers that loop by row."""
+    width = _guarded_pages(page_table.shape[1]) * _KV_BLOCK_SIZE
+    return max(1, _LOGITS_BUDGET_ELEMS // width)
+
+
+def _alloc_logits(
+    num_tokens: int, max_seq_len: int, device: torch.device, is_decode: bool
+) -> torch.Tensor:
+    """Hand out the [num_tokens, max_seq_len] fp32 scratch the logits kernel fills.
+
+    Prefill rectangles are served from one fixed-size pooled block. A fresh
+    `torch.empty` per call would instead feed the caching allocator a
+    monotonically growing size sequence -- the width tracks context length, and
+    an agentic session's context only ever grows -- so every request is slightly
+    larger than any cached block, none can be reused, and each strands a whole
+    segment. `reserved` then climbs while `allocated` stays flat, and that
+    stranded memory is invisible to allocators that bypass torch: Triton kernel
+    scratch fails with HSA_STATUS_ERROR_OUT_OF_RESOURCES instead of surfacing as
+    a clean torch OOM. Serving every rectangle out of one block keeps the
+    request size constant, so the block is always reused and nothing strands.
+
+    Decode keeps the plain allocation: it is captured against the graph memory
+    pool (bounded, separate from the fragmenting general pool), and creating the
+    pooled block mid-capture would hand out graph-pool memory to later replays.
+    """
+    n = num_tokens * max_seq_len
+    if (
+        is_decode
+        or n > _LOGITS_BUDGET_ELEMS
+        or torch.cuda.is_current_stream_capturing()
+    ):
+        return torch.empty(
+            (num_tokens, max_seq_len), dtype=torch.float32, device=device
+        )
+    buf = _LOGITS_POOL.get(device)
+    if buf is None:
+        buf = torch.empty(_LOGITS_BUDGET_ELEMS, dtype=torch.float32, device=device)
+        _LOGITS_POOL[device] = buf
+    return buf[:n].view(num_tokens, max_seq_len)
+
+
 def prepare_fp4_decode_workspace(
     page_table: torch.Tensor,
     c4_seq_lens: torch.Tensor,
@@ -251,12 +308,11 @@ def aiter_fp4_paged_mqa_logits(
         page_table, max_seq_len = _guard_page_table(page_table)
     q_payload = q_fp4.view(torch.uint8)
     k_payload = k_payload.view(torch.uint8)
-    # Scored write-once and freed with this call. Recycling it through the
+    # Scored write-once and dead when the caller's top-k returns, so the pooled
-    # allocator costs nothing because a pinned cta_info makes the kernel skip
+    # block can be handed straight to the next call: a pinned cta_info makes the
-    # its -inf pre-fill and the length-aware top-k reads only [0, c4_seq_len).
+    # kernel skip its -inf pre-fill and the length-aware top-k reads only
-    logits = torch.empty(
+    # [0, c4_seq_len), so neither ever observes the previous chunk's leftovers.
-        (num_tokens, max_seq_len), dtype=torch.float32, device=q_fp4.device
+    logits = _alloc_logits(num_tokens, max_seq_len, q_fp4.device, is_decode)
-    )
     common = {
         "weight_scale": weight_scale,
         "block_k": 256,

python/sglang/srt/layers/attention/dsv4/indexer.py

@@ -19,12 +19,14 @@ import torch.nn.functional as F
 from sglang.kernels.ops.attention.dsv4 import (
     fused_q_indexer_rope_hadamard_fp4_quant,
     fused_q_indexer_rope_hadamard_quant,
+    plan_topk_v2,
     topk_transform_paged,
     topk_transform_paged_v2,
 )
 from sglang.kernels.ops.attention.dsv4.fp4_indexer_hip import (
     aiter_fp4_paged_mqa_logits,
     aiter_q_indexer_fp4,
+    logits_rows_per_chunk,
 )
 from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
@@ -873,7 +875,13 @@ class C4IndexerBackendMixin:
                     page_table[rows],
                     c4_sparse_page_indices[rows],
                     indexer_metadata.c4_page_size,
-                    indexer_metadata.topk_metadata,
+                    # The cached plan routes rows by their index in the full
+                    # range, so a chunk needs one built over its own rows.
+                    (
+                        indexer_metadata.topk_metadata
+                        if rows == all_rows or not is_hip()
+                        else plan_topk_v2(c4_seq_lens[rows])
+                    ),
                 )
             else:
                 topk_transform_paged(
@@ -897,24 +905,46 @@ class C4IndexerBackendMixin:
             run_topk_transform(all_rows, logits)
         elif use_aiter_fp4:
             q_fp4, q_scale = q
-            logits = aiter_fp4_paged_mqa_logits(
+            is_decode = forward_batch.forward_mode.is_decode()
-                q_fp4=q_fp4,
+            # Hoisted: these await this layer's KV transfer, which every chunk
-                q_scale=q_scale,
+            # would otherwise re-await.
-                k_payload=token_to_kv_pool.get_index_k_fp4_payload_buffer(
+            k_payload = token_to_kv_pool.get_index_k_fp4_payload_buffer(
-                    c4_indexer.layer_id
+                c4_indexer.layer_id
-                ),
-                k_scale=token_to_kv_pool.get_index_k_fp4_scale_buffer(
-                    c4_indexer.layer_id
-                ),
-                weights=weights,
-                page_table=page_table,
-                c4_seq_lens=c4_seq_lens,
-                weight_scale=c4_indexer.weight_scale,
-                is_decode=forward_batch.forward_mode.is_decode(),
-                decode_workspace=metadata.fp4_decode_workspace,
-                prefill_workspace=metadata.fp4_prefill_workspace,
             )
-            run_topk_transform(all_rows, logits)
+            k_scale = token_to_kv_pool.get_index_k_fp4_scale_buffer(c4_indexer.layer_id)
+
+            def run_fp4_indexer(rows: slice) -> None:
+                logits = aiter_fp4_paged_mqa_logits(
+                    q_fp4=q_fp4[rows],
+                    q_scale=q_scale[rows],
+                    k_payload=k_payload,
+                    k_scale=k_scale,
+                    weights=weights[rows],
+                    page_table=page_table[rows],
+                    c4_seq_lens=c4_seq_lens[rows],
+                    weight_scale=c4_indexer.weight_scale,
+                    is_decode=is_decode,
+                    decode_workspace=metadata.fp4_decode_workspace,
+                    prefill_workspace=metadata.fp4_prefill_workspace,
+                )
+                run_topk_transform(rows, logits)
+
+            # The scores are the layer's largest transient and their width tracks
+            # context length, so prefill splits the rows into whatever fits the
+            # pooled logits block and reduces each chunk before the next one
+            # reuses it. Rows are scored and reduced independently, so this
+            # matches a single pass. Decode's rectangle is bounded by its capture
+            # shapes, so it always stays whole.
+            rows_per_chunk = (
+                query_rows if is_decode else logits_rows_per_chunk(page_table)
+            )
+            if rows_per_chunk >= query_rows:
+                run_fp4_indexer(all_rows)
+            else:
+                for start in range(0, query_rows, max(1, rows_per_chunk)):
+                    run_fp4_indexer(
+                        slice(start, min(start + rows_per_chunk, query_rows))
+                    )
         else:
             c4_indexer_kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(
                 layer_id=c4_indexer.layer_id,

test/registered/kernels/ops/attention/test_fp4_indexer_hip.py

@@ -751,6 +751,9 @@ def test_pinned_schedule_matches_unpinned_logits(is_decode: bool) -> None:
             case["page_table"], case["c4_seq_lens"]
         )
         pinned = _run_logits(case, is_decode=False, prefill_ws=workspace)
+    # Prefill scores are views of one pooled block, so the second call would
+    # otherwise hand back the same memory and compare it against itself.
+    pinned = pinned.clone()
     unpinned = _run_logits(case, is_decode=is_decode)
 
     seq_len = case["seq_len"]
@@ -772,5 +775,55 @@ def test_stale_workspace_row_count_falls_back_to_inline_schedule() -> None:
     torch.testing.assert_close(with_stale[:, :seq_len], without[:, :seq_len])
 
 
+def test_prefill_logits_come_from_one_pooled_block() -> None:
+    """Prefill must score into one constant-size block, not a per-call rectangle.
+
+    The logits width tracks context length, so a fresh allocation per call feeds
+    the caching allocator a growing size sequence: each request outgrows every
+    cached block and strands a segment, until an allocator that bypasses torch
+    (Triton kernel scratch) is refused. One pooled block keeps the request size
+    constant, which is what makes the blocks reusable.
+    """
+    torch.manual_seed(14)
+    narrow = _run_logits(_build_logits_case(2, 256), is_decode=False)
+    wide = _run_logits(_build_logits_case(4, 512), is_decode=False)
+
+    assert narrow.shape != wide.shape
+    assert narrow.data_ptr() == wide.data_ptr()
+
+
+def test_row_chunks_reproduce_the_unsplit_batch() -> None:
+    """Rows are scored and reduced independently, which is what lets callers chunk.
+
+    ``forward_c4_indexer`` splits prefill rows to whatever fits the pooled block,
+    so a chunk must score its rows exactly as an unsplit call would.
+    """
+    torch.manual_seed(15)
+    batch, chunk_rows = 6, 2
+    case = _build_logits_case(batch, 512)
+    # Cloned: the chunk calls below score into the same pooled block.
+    full = _run_logits(case, is_decode=False).clone()
+
+    for start in range(0, batch, chunk_rows):
+        rows = slice(start, start + chunk_rows)
+        chunk = aiter_fp4_paged_mqa_logits(
+            q_fp4=case["q_fp4"][rows],
+            q_scale=case["q_scale"][rows],
+            k_payload=case["payload"],
+            k_scale=case["scale"],
+            weights=case["weights"][rows],
+            page_table=case["page_table"][rows],
+            c4_seq_lens=case["c4_seq_lens"][rows],
+            weight_scale=case["weight_scale"],
+            is_decode=False,
+        )
+        for row, ctx in enumerate(case["context"][rows].tolist()):
+            torch.testing.assert_close(
+                chunk[row, :ctx],
+                full[start + row, :ctx],
+                msg=f"row {start + row} (ctx={ctx})",
+            )
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v"]))