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

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

Browse Source 
Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
billishyahao
2026-07-09 02:54:00 -07:00
committed by GitHub
co-authored by HAI
parent 078f06fbf4
commit 336b64ecce
2 changed files with 10 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/dsv4/indexer.py
+3 -6
View File
@@ -22,6 +22,7 @@ from sglang.srt.layers.attention.dsv4.metadata import (
PagedIndexerMetadata, PagedIndexerMetadata,
) )
from sglang.srt.layers.linear import ReplicatedLinear 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.forward_batch_info import ForwardMode
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
is_in_breakable_cuda_graph, is_in_breakable_cuda_graph,
@@ -44,12 +45,8 @@ if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_batch_info import ForwardBatch
if is_hip(): FP8_DTYPE = torch.float8_e4m3fnuz if is_fp8_fnuz() else torch.float8_e4m3fn
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
IndexerQuery: TypeAlias = Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] IndexerQuery: TypeAlias = Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
python/sglang/srt/layers/attention/dsv4/unified_kv_kernels/paged_decode.py
+7 -2
View File
@@ -58,12 +58,15 @@ import triton
import triton.language as tl import triton.language as tl
from aiter.ops.triton.utils.device_info import get_num_sms 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. 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). _MAX_KV_SPLITS = 64 # Hard cap on kv_splits (see _kv_splits_heuristic).
# FP8 KV cache (1xGROUP_SIZE block-scale quantization). # 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 # D // GROUP_SIZE] in fp32. Per-slot, D is split into NUM_GROUPS chunks of
# GROUP_SIZE elements; each chunk shares one fp32 scale. # GROUP_SIZE elements; each chunk shares one fp32 scale.
# Dequant in-kernel: kv_bf16 = kv_fp8.to(fp32) * scale[d // GROUP_SIZE], cast # Dequant in-kernel: kv_bf16 = kv_fp8.to(fp32) * scale[d // GROUP_SIZE], cast
@@ -73,7 +76,9 @@ _MAX_KV_SPLITS = 64 # Hard cap on kv_splits (see _kv_splits_heuristic).
# per slot, 4 bytes each → +6.25% storage on top of the fp8 pool (vs the # 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). # halving from bf16→fp8 = 2× saving — net ~46% read bandwidth reduction).
_FP8_GROUP_SIZE = 64 _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) @functools.lru_cache(maxsize=1)