minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Fix FP4 indexer OOR (#37660)

Browse Source 
Co-authored-by: Thomas Wang <1am9trash@gmail.com>
This commit is contained in:
Xinyi Song
2026-09-03 01:49:46 -07:00
committed by GitHub
co-authored by Thomas Wang
parent e59a576f03
commit 7ed29eba80
3 changed files with 164 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py
+63 -7
View File
@@ -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,
python/sglang/srt/layers/attention/dsv4/indexer.py
+48 -18
View File
@@ -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,