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

[AMD] Fix FP8 assertion failure in aiter MLA decode by falling back to self.k_scale (#19935)

Browse Source
This commit is contained in:
jacky.cheng
2026-03-12 16:48:51 +08:00
committed by GitHub
parent 067353f67b
commit 1e2983c98e
4 changed files with 613 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
python/sglang/srt/layers/attention/aiter_backend.py
View File

@@ -1974,8 +1974,12 @@ class AiterAttnBackend(AttentionBackend):
reduce_indptr=reduce_indptr,
reduce_final_map=reduce_final_map,
reduce_partial_map=reduce_partial_map,
q_scale=layer.k_scale,
kv_scale=layer.k_scale,
q_scale=(
layer.k_scale if layer.k_scale is not None else self.k_scale
),
kv_scale=(
layer.k_scale if layer.k_scale is not None else self.k_scale
),
intra_batch_mode=intra_batch_mode,
num_kv_splits=num_kv_splits,
)
@@ -2028,8 +2032,12 @@ class AiterAttnBackend(AttentionBackend):
reduce_indptr=reduce_indptr,
reduce_final_map=reduce_final_map,
reduce_partial_map=reduce_partial_map,
q_scale=layer.k_scale,
kv_scale=layer.k_scale,
q_scale=(
layer.k_scale if layer.k_scale is not None else self.k_scale
),
kv_scale=(
layer.k_scale if layer.k_scale is not None else self.k_scale
),
intra_batch_mode=intra_batch_mode,
num_kv_splits=num_kv_splits,
)
@@ -2059,8 +2067,12 @@ class AiterAttnBackend(AttentionBackend):
reduce_indptr=reduce_indptr,
reduce_final_map=reduce_final_map,
reduce_partial_map=reduce_partial_map,
q_scale=layer.k_scale,
kv_scale=layer.k_scale,
q_scale=(
layer.k_scale if layer.k_scale is not None else self.k_scale
),
kv_scale=(
layer.k_scale if layer.k_scale is not None else self.k_scale
),
intra_batch_mode=intra_batch_mode,
num_kv_splits=num_kv_splits,
)
@@ -2195,8 +2207,8 @@ class AiterAttnBackend(AttentionBackend):
reduce_indptr=reduce_indptr,
reduce_final_map=reduce_final_map,
reduce_partial_map=reduce_partial_map,
q_scale=layer.k_scale,
kv_scale=layer.k_scale,
q_scale=layer.k_scale if layer.k_scale is not None else self.k_scale,
kv_scale=layer.k_scale if layer.k_scale is not None else self.k_scale,
intra_batch_mode=intra_batch_mode,
num_kv_splits=num_kv_splits,
)