From 7ed29eba80b5dd04b81b55bbc59a1dbf807b3c97 Mon Sep 17 00:00:00 2001 From: Xinyi Song Date: Thu, 3 Sep 2026 01:49:46 -0700 Subject: [PATCH] [AMD] Fix FP4 indexer OOR (#37660) Co-authored-by: Thomas Wang <1am9trash@gmail.com> --- .../ops/attention/dsv4/fp4_indexer_hip.py | 70 +++++++++++++++++-- .../srt/layers/attention/dsv4/indexer.py | 66 ++++++++++++----- .../ops/attention/test_fp4_indexer_hip.py | 53 ++++++++++++++ 3 files changed, 164 insertions(+), 25 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py index 5380f104e..641ec673d 100644 --- a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 6f21371e7..1be191040 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/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( - q_fp4=q_fp4, - q_scale=q_scale, - k_payload=token_to_kv_pool.get_index_k_fp4_payload_buffer( - 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, + is_decode = forward_batch.forward_mode.is_decode() + # Hoisted: these await this layer's KV transfer, which every chunk + # would otherwise re-await. + k_payload = token_to_kv_pool.get_index_k_fp4_payload_buffer( + c4_indexer.layer_id ) - 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, diff --git a/test/registered/kernels/ops/attention/test_fp4_indexer_hip.py b/test/registered/kernels/ops/attention/test_fp4_indexer_hip.py index 11b877b97..cf6aacb6a 100644 --- a/test/registered/kernels/ops/attention/test_fp4_indexer_hip.py +++ b/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"]))