[AMD][Spec] Enable GDN ReplaySSM target-verify on ROCm (#38184)
Co-authored-by: chuyeh <298092489+chuyeh@users.noreply.github.com>
This commit is contained in:
@@ -43,6 +43,11 @@ import torch
|
|||||||
import triton
|
import triton
|
||||||
import triton.language as tl
|
import triton.language as tl
|
||||||
|
|
||||||
|
from sglang.kernels.ops.attention.fla.utils import is_tf32_supported
|
||||||
|
from sglang.srt.utils import is_gfx95_supported
|
||||||
|
|
||||||
|
_IS_GFX95 = is_gfx95_supported()
|
||||||
|
|
||||||
|
|
||||||
@triton.jit
|
@triton.jit
|
||||||
def gdn_replayssm_spec_circular_kernel(
|
def gdn_replayssm_spec_circular_kernel(
|
||||||
@@ -1269,7 +1274,8 @@ def gdn_replayssm_spec_decode(
|
|||||||
num_stages_flush: int = 2,
|
num_stages_flush: int = 2,
|
||||||
nk_flush: int = 2,
|
nk_flush: int = 2,
|
||||||
launch_mode: str = "both",
|
launch_mode: str = "both",
|
||||||
dot_precision: str = "tf32",
|
# Triton only accepts TF32 input precision on Ampere-or-newer NVIDIA GPUs.
|
||||||
|
dot_precision: str = "tf32" if is_tf32_supported else "ieee",
|
||||||
):
|
):
|
||||||
"""GDN cached speculative-decode on a CIRCULAR ring cache (split-qkv varlen).
|
"""GDN cached speculative-decode on a CIRCULAR ring cache (split-qkv varlen).
|
||||||
|
|
||||||
@@ -1286,6 +1292,18 @@ def gdn_replayssm_spec_decode(
|
|||||||
"""
|
"""
|
||||||
if scale is None:
|
if scale is None:
|
||||||
scale = checkpoint_state.shape[-1] ** -0.5
|
scale = checkpoint_state.shape[-1] ** -0.5
|
||||||
|
batch_size = query_start_loc.shape[0] - 1
|
||||||
|
# At the Qwen3.5 MTP shape, one wide V tile avoids duplicating q/k and gate
|
||||||
|
# work across two programs on gfx950.
|
||||||
|
if (
|
||||||
|
_IS_GFX95
|
||||||
|
and batch_size >= 32
|
||||||
|
and max_cache_len == 16
|
||||||
|
and max_spec_len == 4
|
||||||
|
and q.shape[-1] == 128
|
||||||
|
and v.shape[-1] == 128
|
||||||
|
):
|
||||||
|
block_v, num_stages = 128, 1
|
||||||
if is_flush.dtype != torch.int8:
|
if is_flush.dtype != torch.int8:
|
||||||
is_flush = is_flush.to(torch.int8)
|
is_flush = is_flush.to(torch.int8)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user