[AMD] Fix GLM-5 fp8 KV quant path dispatch on MI300 (#22314)

This commit is contained in:
Thomas Wang
2026-04-07 21:16:02 -07:00
committed by GitHub
parent 36f05810c9
commit 729b74d8dd
+3 -7
View File
@@ -45,7 +45,7 @@ from sglang.srt.layers.attention.nsa.quant_k_cache import (
quantize_k_cache,
quantize_k_cache_separate,
)
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.utils import (
get_mla_kv_buffer_triton,
@@ -1575,13 +1575,9 @@ class MLATokenToKVPool(KVCache):
):
layer_id = layer.layer_id
if self.nsa_kv_cache_store_fp8:
if _is_hip:
if _is_hip and self.use_nsa and self.dtype == fp8_dtype:
# HIP FP8 path uses raw MLA KV layout (nope + rope) without per-block scales.
# Fuse BF16/FP16 -> FP8 cast with paged KV write.
fp8_dtype = (
torch.float8_e4m3fnuz if _is_fp8_fnuz else torch.float8_e4m3fn
)
set_mla_kv_buffer_triton_fp8_quant(
self.kv_buffer[layer_id - self.start_layer],
loc,
@@ -1589,7 +1585,7 @@ class MLATokenToKVPool(KVCache):
cache_k_rope,
fp8_dtype,
)
else:
elif self.nsa_kv_cache_store_fp8:
# OPTIMIZATION: Quantize k_nope and k_rope separately to avoid concat overhead
# This also enables reuse of set_mla_kv_buffer_triton two-tensor write path
# quantize_k_cache_separate returns (nope_part, rope_part) as uint8 bytes