[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
-    # allocator costs nothing because a pinned cta_info makes the kernel skip
-    # its -inf pre-fill and the length-aware top-k reads only [0, c4_seq_len).
-    logits = torch.empty(
-        (num_tokens, max_seq_len), dtype=torch.float32, device=q_fp4.device
-    )
+    # Scored write-once and dead when the caller's top-k returns, so the pooled
+    # block can be handed straight to the next call: a pinned cta_info makes the
+    # kernel skip its -inf pre-fill and the length-aware top-k reads only
+    # [0, c4_seq_len), so neither ever observes the previous chunk's leftovers.
+    logits = _alloc_logits(num_tokens, max_seq_len, q_fp4.device, is_decode)
     common = {
         "weight_scale": weight_scale,
         "block_k": 256,