Skip to content

[AMD] fix dsv4 indexer dtype dispatch on gfx950 #29479

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
HaiShaw merged 4 commits into from
Jul 9, 2026
Merged

[AMD] fix dsv4 indexer dtype dispatch on gfx950 #29479

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8e649be
[AMD] fix indexer dtype dispatch on gfx950
billishyahao Jun 27, 2026
72c7b01
cleanup
billishyahao Jun 29, 2026
a7fffbd
cleanup
billishyahao Jun 29, 2026
63e09b3
Merge branch 'main' into fix_fp8
HaiShaw Jul 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 3 additions & 6 deletions python/sglang/srt/layers/attention/dsv4/indexer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
PagedIndexerMetadata,
)
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
is_in_breakable_cuda_graph,
Expand All @@ -44,12 +45,8 @@
from sglang.srt.model_executor.forward_batch_info import ForwardBatch


if is_hip():
FP8_DTYPE = torch.float8_e4m3fnuz
FP8_MAX = torch.finfo(FP8_DTYPE).max
else:
FP8_DTYPE = torch.float8_e4m3fn
FP8_MAX = torch.finfo(FP8_DTYPE).max
FP8_DTYPE = torch.float8_e4m3fnuz if is_fp8_fnuz() else torch.float8_e4m3fn


IndexerQuery: TypeAlias = Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]

Expand Down
9 changes: 7 additions & 2 deletions python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,12 +58,15 @@
import triton.language as tl
from aiter.ops.triton.utils.device_info import get_num_sms

from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz

LOG2E = 1.4426950408889634 # log2(e); folded into qk_scale so softmax can use exp2.
_MAX_KV_SPLITS = 64 # Hard cap on kv_splits (see _kv_splits_heuristic).

# FP8 KV cache (1xGROUP_SIZE block-scale quantization).
#
# Storage: unified_kv[total_pages, D] in e4m3fnuz + kv_scales[total_pages,
# Storage: unified_kv[total_pages, D] in e4m3fnuz (gfx942 only) / e4m3fn +
# kv_scales[total_pages,
# D // GROUP_SIZE] in fp32. Per-slot, D is split into NUM_GROUPS chunks of
# GROUP_SIZE elements; each chunk shares one fp32 scale.
# Dequant in-kernel: kv_bf16 = kv_fp8.to(fp32) * scale[d // GROUP_SIZE], cast
Expand All @@ -73,7 +76,9 @@
# per slot, 4 bytes each → +6.25% storage on top of the fp8 pool (vs the
# halving from bf16→fp8 = 2× saving — net ~46% read bandwidth reduction).
_FP8_GROUP_SIZE = 64
_FP8_DTYPE = torch.float8_e4m3fnuz
# Match the KV-cache write format: fnuz only on AMD gfx942; Other platform e.g.
# gfx950 store e4m3fn. See quant_k_cache.py / indexer.py FP8_DTYPE.
_FP8_DTYPE = torch.float8_e4m3fnuz if is_fp8_fnuz() else torch.float8_e4m3fn


@functools.lru_cache(maxsize=1)
Expand Down
Loading