chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[AMD] Support ds3.2 on gfx942 platform (#17504)

Browse Source 
Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
This commit is contained in:
wufann
2026-01-23 05:57:08 +08:00
committed by GitHub
parent 15b511771d
commit a921029b97
4 changed files with 135 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
python/sglang/srt/layers/attention/nsa/index_buf_accessor.py
View File

@@ -4,8 +4,12 @@ import torch
import triton
import triton.language as tl
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.utils import is_hip
_is_hip = is_hip()
_is_fp8_fnuz = is_fp8_fnuz()
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
@@ -349,21 +353,24 @@ def _set_k_and_s_triton(
raise ValueError(
f"index_k_scale must be 1D or 2D, got shape {index_k_scale.shape}"
)
if is_hip():
if _is_hip:
assert buf_numel_per_page == 1 * (128 + 4)
else:
assert buf_numel_per_page == 64 * (128 + 4)
assert num_tokens_to_write == num_tokens_to_write_ == num_tokens_to_write__
assert index_head_dim == 128
assert scale_dim == 1
if is_hip():
if _is_hip:
assert page_size == 1
else:
assert page_size == 64
assert buf.dtype == torch.uint8
assert loc.dtype == torch.int64, f"{loc.dtype=}" # can be int32
assert index_k.dtype == torch.float8_e4m3fn
if _is_fp8_fnuz:
assert index_k.dtype == torch.float8_e4m3fnuz
else:
assert index_k.dtype == torch.float8_e4m3fn
assert index_k_scale.dtype == torch.float32
assert buf.is_contiguous()
@@ -371,7 +378,10 @@ def _set_k_and_s_triton(
assert index_k.is_contiguous()
assert index_k_scale.is_contiguous()
buf_fp8 = buf.view(torch.float8_e4m3fn)
if _is_fp8_fnuz:
buf_fp8 = buf.view(torch.float8_e4m3fnuz)
else:
buf_fp8 = buf.view(torch.float8_e4m3fn)
buf_fp32 = buf.view(torch.float32)
_set_k_and_s_triton_kernel[(num_tokens_to_write,)](

17
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -8,6 +8,7 @@ import torch
from einops import rearrange
from sglang.srt.layers.layernorm import LayerNorm
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.layers.utils import MultiPlatformOp
from sglang.srt.utils import add_prefix, ceil_align, is_cuda, is_hip, is_npu
@@ -15,6 +16,7 @@ global _use_multi_stream
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_npu = is_npu()
_is_fp8_fnuz = is_fp8_fnuz()
if _is_cuda:
try:
import deep_gemm
@@ -504,8 +506,10 @@ class Indexer(MultiPlatformOp):
)
k_fp8_list.append(k_fp8)
k_scale_list.append(k_scale)
k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn)
if _is_fp8_fnuz:
k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fnuz)
else:
k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn)
k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1)
kv_fp8 = (k_fp8, k_scale)
ks, ke = metadata.get_indexer_kvcache_range()
@@ -569,9 +573,9 @@ class Indexer(MultiPlatformOp):
from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits
kv, scale = kv_fp8
logits = fp8_mqa_logits(
logits_chunk = fp8_mqa_logits(
q_fp8[start:end],
kv_fp8,
kv,
scale,
weights[start:end],
ks[start:end],
@@ -1001,11 +1005,12 @@ class Indexer(MultiPlatformOp):
.view(m, ng, group)
.mul_(x_s.to(torch.float32).unsqueeze(-1))
.view(m, n)
.to(torch.bfloat16)
)
else:
x_for_gate = x_q.to(torch.float32)
x_for_gate = x_q.to(torch.bfloat16)
else:
x_for_gate = x_q.to(torch.float32)
x_for_gate = x_q.to(torch.bfloat16)
else:
x_for_gate = x

46
python/sglang/srt/layers/attention/nsa/tilelang_kernel.py
View File

@@ -4,21 +4,28 @@ import tilelang
import tilelang.language as T
import torch
from sglang.srt.utils import is_hip
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.utils import is_gfx95_supported, is_hip
tilelang.set_log_level("WARNING")
pass_configs = {
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
tilelang.PassConfigKey.TL_DISABLE_FAST_MATH: True,
}
BF16 = "bfloat16"
FP8 = "float8_e4m3"
FP32 = "float32"
# TL_DISABLE_FAST_MATH has deprecated in v0.1.7.post1 tilelang
if hasattr(tilelang.PassConfigKey, "TL_DISABLE_FAST_MATH"):
pass_configs[tilelang.PassConfigKey.TL_DISABLE_FAST_MATH] = True
elif hasattr(tilelang.PassConfigKey, "TL_ENABLE_FAST_MATH"):
pass_configs[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] = False
_is_hip = is_hip()
_is_gfx95_supported = is_gfx95_supported()
_is_fp8_fnuz = is_fp8_fnuz()
BF16 = "bfloat16"
FP8 = "float8_e4m3fnuz" if _is_fp8_fnuz else "float8_e4m3"
FP32 = "float32"
def fast_log2_ceil(x):
@@ -42,8 +49,8 @@ def act_quant_kernel(
N, in_dtype=BF16, out_dtype=FP8, scale_dtype=FP32, round_scale=False
):
M = T.symbolic("M")
fp8_min = -448.0
fp8_max = 448.0
fp8_min = -224.0 if _is_fp8_fnuz else -448.0
fp8_max = 224.0 if _is_fp8_fnuz else 448.0
fp8_max_inv = 1 / fp8_max
num_stages = 0 if round_scale else 2
blk_m = 32
@@ -108,7 +115,10 @@ def act_quant(
x.size(-1) % block_size == 0
), f"Last dimension size must be divisible by block_size (block_size={block_size})"
N = x.size(-1)
y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
if _is_fp8_fnuz:
y = torch.empty_like(x, dtype=torch.float8_e4m3fnuz)
else:
y = torch.empty_like(x, dtype=torch.float8_e4m3fn)
s = x.new_empty(*x.size()[:-1], N // block_size, dtype=torch.float32)
kernel = act_quant_kernel(N, round_scale=scale_fmt is not None)
kernel(x.view(-1, N), y.view(-1, N), s.view(-1, N // block_size))
@@ -777,9 +787,21 @@ def tilelang_sparse_fwd(
topk = indices.shape[-1]
assert topk == 2048
if _is_hip:
kernel = sparse_attention_fwd_kernel_v1(
num_heads, d_v, tail_dim, topk, sm_scale=sm_scale, num_stages=1
)
if _is_gfx95_supported:
kernel = sparse_attention_fwd_kernel_v1(
num_heads, d_v, tail_dim, topk, sm_scale=sm_scale, num_stages=1
)
else: # reduce LDS usage on gfx942 target
kernel = sparse_attention_fwd_kernel_v1(
num_heads,
d_v,
tail_dim,
topk,
sm_scale=sm_scale,
block_I=32,
num_stages=1,
threads=128,
)
else:
kernel = sparse_attention_fwd_kernel_v2(
num_heads, d_v, tail_dim, topk, sm_scale=sm_scale