diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a91ab5b37..ef1c3869a 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -650,6 +650,7 @@ class Envs: SGLANG_OPT_DEEPGEMM_HC_PRENORM = EnvBool(True) SGLANG_OPT_USE_TILELANG_MHC_PRE = EnvBool(True) SGLANG_OPT_USE_TILELANG_MHC_POST = EnvBool(True) + SGLANG_OPT_USE_TRITON_FUSED_MHC = EnvBool(True) SGLANG_OPT_USE_TILELANG_INDEXER = EnvBool(False) SGLANG_OPT_USE_AITER_INDEXER = EnvBool(False) SGLANG_OPT_USE_JIT_INDEXER_METADATA = EnvBool(True) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index ec59114d9..88f27563a 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -51,6 +51,7 @@ from sglang.srt.layers.dp_attention import ( ) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc from sglang.srt.speculative.spec_info import SpecInput from sglang.srt.utils import ceil_align @@ -500,32 +501,21 @@ class DeepseekV4HipRadixBackend( req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, out_cache_loc: Optional[torch.Tensor] = None, + extend_seq_lens: Optional[torch.Tensor] = None, use_prefill_cuda_graph: bool = False, ) -> Union[DSV4Metadata, DSV4RawVerifyMetadata]: - if envs.SGLANG_PREP_IN_CUDA_GRAPH.get(): - assert out_cache_loc is not None - if not hasattr(self, "extend_seq_lens_buffer"): - self.extend_seq_lens_buffer = torch.tensor( - [self.speculative_num_draft_tokens] * 1025, device=self.device - ) - extend_seq_lens = self.extend_seq_lens_buffer[: len(seq_lens)] - - return DSV4RawVerifyMetadata( - req_pool_indices=req_pool_indices, - seq_lens=seq_lens, - out_cache_loc=out_cache_loc, - extend_seq_lens=extend_seq_lens, - ) - else: - seq_lens_cpu = seq_lens.tolist() - return self.init_forward_metadata_target_verify_old( - max_seq_len=max_seq_len, - req_pool_indices=req_pool_indices, - seq_lens=seq_lens, - seq_lens_cpu=seq_lens_cpu, - out_cache_loc=out_cache_loc, - use_prefill_cuda_graph=use_prefill_cuda_graph, - ) + # HIP path: build target-verify metadata eagerly even when + # SGLANG_PREP_IN_CUDA_GRAPH is enabled. The raw/lazy-upgrade route can + # hit planner invariants during graph capture for DSV4+EAGLE. + seq_lens_cpu = seq_lens.tolist() + return self.init_forward_metadata_target_verify_old( + max_seq_len=max_seq_len, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + out_cache_loc=out_cache_loc, + use_prefill_cuda_graph=use_prefill_cuda_graph, + ) def init_forward_metadata_target_verify_old( self, @@ -565,8 +555,15 @@ class DeepseekV4HipRadixBackend( out_cache_loc = raw_metadata.out_cache_loc bs, num_draft_tokens = len(seq_lens), self.speculative_num_draft_tokens - seq_lens = seq_lens + self.speculative_num_draft_tokens + seq_lens = seq_lens + num_draft_tokens extend_seq_lens = raw_metadata.extend_seq_lens + if extend_seq_lens is None or extend_seq_lens.numel() != bs: + extend_seq_lens = torch.full_like(seq_lens, num_draft_tokens) + else: + extend_seq_lens = extend_seq_lens.to( + device=seq_lens.device, dtype=seq_lens.dtype + ) + extend_seq_lens = torch.minimum(extend_seq_lens, seq_lens).clamp_min_(1) seq_lens_casual, req_pool_indices_repeated = ( self.expand_extend_with_same_length( @@ -678,11 +675,22 @@ class DeepseekV4HipRadixBackend( max_seq_len = int(seq_lens_cpu.max().item()) if forward_batch.forward_mode.is_decode_or_idle(): + # DSv4 bakes this step's KV write target (c4/c128) into metadata, + # so slice the shared multi-step out_cache_loc now rather than at + # forward time. + out_cache_loc = forward_batch.out_cache_loc + if self.topk > 0 and self.speculative_num_steps > 1: + out_cache_loc = per_step_draft_out_cache_loc( + out_cache_loc, + forward_batch.batch_size, + self.topk, + self.speculative_num_steps, + )[self.speculative_step_id] metadata = self.init_forward_metadata_decode( max_seq_len=max_seq_len, req_pool_indices=req_pool_indices, seq_lens=seq_lens, - out_cache_loc=forward_batch.out_cache_loc, + out_cache_loc=out_cache_loc, ) elif forward_batch.forward_mode.is_target_verify(): metadata = self.init_forward_metadata_target_verify( @@ -690,6 +698,7 @@ class DeepseekV4HipRadixBackend( req_pool_indices=req_pool_indices, seq_lens=seq_lens, out_cache_loc=forward_batch.out_cache_loc, + extend_seq_lens=forward_batch.extend_seq_lens, ) elif forward_batch.forward_mode.is_prefill(include_draft_extend_v2=True): extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu diff --git a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_fused.py b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_fused.py index 6167f58bf..41c29efb1 100644 --- a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_fused.py +++ b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_fused.py @@ -1208,6 +1208,9 @@ def _prune_splitk_configs(configs, named_args, **kwargs): triton.Config({"BLOCK_H": 64, "BLOCK_N": 128}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 128, "BLOCK_N": 64}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 128, "BLOCK_N": 128}, num_warps=4, num_stages=1), + # BLOCK_H=32: critical for cc=32 with h_q=128 (gives 256 blocks with split_k=2) + triton.Config({"BLOCK_H": 32, "BLOCK_N": 64}, num_warps=4, num_stages=1), + triton.Config({"BLOCK_H": 32, "BLOCK_N": 128}, num_warps=4, num_stages=1), ], key=["total_tokens_bucket", "h_q", "topk_per_split"], prune_configs_by={"early_config_prune": _prune_splitk_configs}, @@ -1590,6 +1593,7 @@ def fused_gather_attn_decode_dsv4_dual_scope( topk_length_extra: Optional[torch.Tensor] = None, attn_sink: Optional[torch.Tensor] = None, s_q: int = 1, + force_no_splitk: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Fused gather+dequant+attention for DSV4 with dual scope (main + extra). @@ -1608,6 +1612,9 @@ def fused_gather_attn_decode_dsv4_dual_scope( topk_length_extra: Optional per-batch topk length for extra [b] attn_sink: Optional attention sink values [h_q] s_q: Sequence length per batch + force_no_splitk: If True, skip split-K and use the non-splitk kernel + directly. Used by the dispatch layer for large batch prefill where + the non-splitk fused kernel avoids intermediate buffer allocation. Returns: output: Attention output [total_tokens, h_q, d_v] @@ -1647,6 +1654,10 @@ def fused_gather_attn_decode_dsv4_dual_scope( or kv_cache_size_extra > BUFFER_OPS_DISABLE_THRESHOLD ) + # When force_no_splitk is set, skip the split-K decision and fall + # through to the non-splitk kernel path below. + use_splitk = not force_no_splitk + # Use Split-K for dual scope in these cases: # 1. Small batch sizes with h_q=128 or large topk to increase GPU parallelism # 2. Large topk (>= 2048) with medium/large batch sizes @@ -1665,7 +1676,7 @@ def fused_gather_attn_decode_dsv4_dual_scope( # For h_q > 64 (e.g. h_q=128), the non-splitk grid has very few blocks # in the H dimension, leading to low GPU utilization at medium batch sizes. use_splitk_for_large_hq = h_q > 64 and total_tokens > 8 and total_topk >= 256 - if ( + if use_splitk and ( use_splitk_for_small_bs or use_splitk_for_h64_large_topk or use_splitk_for_large_topk diff --git a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_optimized.py b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_optimized.py index 02891b91b..7426dd839 100644 --- a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_optimized.py +++ b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_optimized.py @@ -28,6 +28,7 @@ from .triton_mla_kernels_decode_dsv4 import ( ) from .triton_mla_kernels_decode_fused import ( fused_gather_attn_decode_dsv4, + fused_gather_attn_decode_dsv4_dual_scope, fused_gather_attn_decode_dsv4_dual_scope_low_overhead, ) @@ -56,14 +57,8 @@ def triton_sparse_attn_decode( def _should_use_fused_dual_scope(total_tokens: int, h_q: int, total_topk: int) -> bool: """Determine whether to use fused kernel for dual-scope cases. - The fused kernel avoids allocating a large intermediate gathered_kv - buffer and eliminates a separate gather kernel launch. However, for - h_q > 64 with medium-to-large batch sizes and larger topk, the - non-splitk fused kernel suffers from low GPU utilization (the grid - has only cdiv(h_q, BLOCK_H) blocks in the H dimension). In those - cases the fallback (separate gather + attention) can be faster on - the GPU, though it incurs extra torch.empty() overhead in CUDA - graphs. + Returns True if the fused kernel (with splitk for small bs) should be used. + For large batch sizes (>= 256), use _should_use_fused_nosplitk instead. The thresholds below were determined empirically on MI355X (256 CUs). """ @@ -74,17 +69,52 @@ def _should_use_fused_dual_scope(total_tokens: int, h_q: int, total_topk: int) - if h_q <= 64 and total_topk >= 1024: return total_tokens <= 128 # h_q > 64 (e.g. h_q=128 when q is padded to full n_heads). - # For small topk (c128 layers, topk~192), fused always wins. - # For larger topk (c4 layers, topk~640), fused wins at small bs - # but the fallback catches up at bs>=16 due to better GPU utilization. - # However, the fallback has 4 extra torch.empty() calls that add - # ~30us CUDA-graph replay overhead, roughly cancelling the GPU gain. - # So we route to fused for all practical batch sizes. if h_q > 64: - return total_tokens <= 256 + if total_topk >= 400: + return total_tokens <= 32 + else: + return total_tokens <= 128 return True +def _should_use_fused_nosplitk(total_tokens: int, h_q: int, total_topk: int) -> bool: + """Determine whether to use the fused no-splitk kernel for large batches. + + Kernel-level benchmarking on MI355X shows that for large batch sizes + (total_tokens >= 256), the fused dual-scope kernel WITHOUT split-K + is ~10% faster than the separate gather+attention path: + + total_tokens=256: fused-noSK=169us vs separate=194us (14% faster) + total_tokens=512: fused-noSK=350us vs separate=408us (14% faster) + total_tokens=1024: fused-noSK=700us vs separate=777us (10% faster) + total_tokens=4096: fused-noSK=2761us vs separate=3063us (10% faster) + + The fused no-splitk kernel avoids: + 1. Materializing the large intermediate gathered_kv buffer + 2. The separate gather kernel launch + 3. The split-K combine overhead + + For total_tokens < 256, the separate path is faster because the + fused kernel has insufficient parallelism. + + For extend (total_tokens >= 1024), the fused kernel always wins + regardless of h_q or total_topk because: + - The grid already has thousands of blocks (good GPU utilization) + - It eliminates 1.5-5 GB gathered_kv buffer allocation + - It eliminates 2x gather_dequant kernel launches (~414 us) + - It avoids chunking that TP>1 configs require with the separate path + """ + if total_tokens >= 1024: + return True + if h_q <= 64: + return False # Not benchmarked for h_q <= 64 + if total_topk < 200: + return False # Small topk doesn't benefit + # For h_q > 64 and total_topk >= 200: + # Fused no-splitk wins for total_tokens >= 256 + return total_tokens >= 256 + + def _triton_sparse_attn_decode_dsv4( q: torch.Tensor, kv_scope, @@ -135,7 +165,39 @@ def _triton_sparse_attn_decode_dsv4( topk_extra = extra_kv_scope.indices_in_kvcache.shape[-1] total_topk = topk_main + topk_extra - # Check if chunking needed (fall back to original implementation) + # For large batch sizes, use fused no-splitk kernel (10% faster than separate). + # This check is BEFORE the chunking check because the fused kernel does NOT + # allocate the intermediate gathered_kv buffer, so buffer size limits don't apply. + if _should_use_fused_nosplitk(total_tokens, h_q, total_topk): + q_reshaped = q.reshape(total_tokens, h_q, d_qk).contiguous() + + indices_main = kv_scope.indices_in_kvcache.reshape( + total_tokens, topk_main + ).contiguous() + + block_size_extra = extra_kv_scope.blocked_k.shape[1] + indices_extra = extra_kv_scope.indices_in_kvcache.reshape( + total_tokens, topk_extra + ).contiguous() + + output, lse = fused_gather_attn_decode_dsv4_dual_scope( + q_reshaped, + kv_quantized_main, + indices_main, + block_size_main, + extra_kv_scope.blocked_k_quantized, + indices_extra, + block_size_extra, + sm_scale, + topk_length_main=kv_scope.topk_length, + topk_length_extra=extra_kv_scope.topk_length, + attn_sink=attn_sink, + s_q=s_q, + force_no_splitk=True, + ) + return output.view(b, s_q, h_q, d_v), lse.view(b, s_q, h_q).transpose(1, 2) + + # Check if chunking needed for separate path (fall back to original implementation) token_ranges = compute_token_ranges(total_tokens, total_topk, d_qk) if len(token_ranges) > 1: from .triton_mla_kernels_decode_dsv4 import triton_sparse_attn_decode_dsv4 @@ -146,20 +208,16 @@ def _triton_sparse_attn_decode_dsv4( # Use fused dual-scope kernel with low-overhead buffer pool if _should_use_fused_dual_scope(total_tokens, h_q, total_topk): - q_reshaped = q.reshape(total_tokens, h_q, d_qk) - if not q_reshaped.is_contiguous(): - q_reshaped = q_reshaped.contiguous() + q_reshaped = q.reshape(total_tokens, h_q, d_qk).contiguous() - indices_main = kv_scope.indices_in_kvcache.reshape(total_tokens, topk_main) - if not indices_main.is_contiguous(): - indices_main = indices_main.contiguous() + indices_main = kv_scope.indices_in_kvcache.reshape( + total_tokens, topk_main + ).contiguous() block_size_extra = extra_kv_scope.blocked_k.shape[1] indices_extra = extra_kv_scope.indices_in_kvcache.reshape( total_tokens, topk_extra - ) - if not indices_extra.is_contiguous(): - indices_extra = indices_extra.contiguous() + ).contiguous() output, lse = fused_gather_attn_decode_dsv4_dual_scope_low_overhead( q_reshaped, diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index 816702f3c..e3b8f4d79 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -19,6 +19,7 @@ from sglang.srt.server_args import get_global_server_args from sglang.srt.utils.common import ( get_bool_env_var, is_cuda, + is_hip, is_musa, is_npu, ) @@ -41,6 +42,9 @@ if is_musa(): top_p_renorm_prob, ) +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip() +if _use_aiter: + from aiter import greedy_sample as _aiter_greedy_sample if is_npu(): import torch_npu @@ -106,8 +110,13 @@ class Sampler(nn.Module): logits = self._preprocess_logits(logits, sampling_info) if sampling_info.is_all_greedy: - # Use torch.argmax if all requests use greedy sampling - batch_next_token_ids = torch.argmax(logits, -1) + if _use_aiter: + batch_next_token_ids = torch.empty( + logits.shape[0], device=logits.device, dtype=torch.int32 + ) + _aiter_greedy_sample(batch_next_token_ids, logits) + else: + batch_next_token_ids = torch.argmax(logits, -1) if return_logprob: original_logprobs = logprobs = torch.nn.functional.log_softmax( logits, dim=-1 diff --git a/python/sglang/srt/models/deepseek_common/amd/__init__.py b/python/sglang/srt/models/deepseek_common/amd/__init__.py new file mode 100644 index 000000000..f7fbd557d --- /dev/null +++ b/python/sglang/srt/models/deepseek_common/amd/__init__.py @@ -0,0 +1 @@ +# AMD-specific DeepSeek common model helpers. diff --git a/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_fused_mhc.py b/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_fused_mhc.py new file mode 100644 index 000000000..0020fb892 --- /dev/null +++ b/python/sglang/srt/models/deepseek_common/amd/deepseek_v4_fused_mhc.py @@ -0,0 +1,158 @@ +import logging +from typing import Optional, Tuple + +import torch +import triton + +from sglang.srt.environ import envs + +logger = logging.getLogger(__name__) + +_FUSED_HC_POST_PRE_M_THRESHOLD = 64 +_FUSED_HC_POST_PRE_CACHE: dict[tuple, dict[str, torch.Tensor]] = {} +_TRITON_MHC_POST_PRE_OPS = None +_TRITON_MHC_POST_PRE_RUNTIME_DISABLED = False + + +def _get_triton_mhc_post_pre_ops(): + global _TRITON_MHC_POST_PRE_OPS + + if _TRITON_MHC_POST_PRE_OPS is not None: + return _TRITON_MHC_POST_PRE_OPS + + try: + from aiter.ops.triton.fusions.mhc import mhc_post_pre + from aiter.ops.triton.utils.mhc_config_utils import get_mhc_config + except Exception as err: + logger.warning( + "Triton fused mHC (mhc_post_pre) is unavailable, falling back: %s", err + ) + return None + + _TRITON_MHC_POST_PRE_OPS = (mhc_post_pre, get_mhc_config) + return _TRITON_MHC_POST_PRE_OPS + + +def _get_fused_hc_post_pre_buffers( + num_tokens: int, + hidden_size: int, + hc_mult: int, + dtype: torch.dtype, + device: torch.device, +) -> Optional[dict[str, torch.Tensor]]: + ops = _get_triton_mhc_post_pre_ops() + if ops is None: + return None + _, get_mhc_config = ops + + key = (num_tokens, hidden_size, hc_mult, dtype, device.type, device.index) + bufs = _FUSED_HC_POST_PRE_CACHE.get(key) + if bufs is not None: + return bufs + + try: + cfg, _ = get_mhc_config("MHC_FUSED", num_tokens, hidden_size, mode="sinkhorn") + except Exception as err: + logger.warning("Failed to initialize fused mHC config, falling back: %s", err) + return None + + n_total = 2 * hc_mult + hc_mult * hc_mult + k_dim = hc_mult * hidden_size + block_k = cfg.get("BLOCK_K", min(512, triton.next_power_of_2(k_dim))) + block_k = min(block_k, triton.next_power_of_2(k_dim)) + block_c_split = max(block_k // hc_mult, 1) + num_ksplit = triton.cdiv(hidden_size, block_c_split) + + bufs = { + "residual_out": torch.empty( + num_tokens, hc_mult, hidden_size, dtype=dtype, device=device + ), + "layer_input_out": torch.empty( + num_tokens, hidden_size, dtype=dtype, device=device + ), + "h_post": torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=device), + "h_res": torch.empty( + num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=device + ), + "acc_partial": torch.empty( + num_ksplit, num_tokens, n_total, dtype=torch.float32, device=device + ), + "acc_sq_partial": torch.empty( + num_ksplit, num_tokens, dtype=torch.float32, device=device + ), + } + _FUSED_HC_POST_PRE_CACHE[key] = bufs + return bufs + + +def try_fused_hc_post_pre( + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + hc_fn_t: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_mult: int, + norm_eps: float, + hc_eps: float, + hc_post_mult: float, + sinkhorn_iters: int, + is_gfx95_supported: bool, +) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool]]: + global _TRITON_MHC_POST_PRE_RUNTIME_DISABLED + + if ( + _TRITON_MHC_POST_PRE_RUNTIME_DISABLED + or not envs.SGLANG_OPT_USE_TRITON_FUSED_MHC.get() + or not is_gfx95_supported + or x.shape[0] == 0 + or x.shape[0] > _FUSED_HC_POST_PRE_M_THRESHOLD + or x.dim() != 2 + or residual.dim() != 3 + ): + return None + + ops = _get_triton_mhc_post_pre_ops() + if ops is None: + return None + mhc_post_pre, _ = ops + + bufs = _get_fused_hc_post_pre_buffers( + x.shape[0], x.shape[1], hc_mult, residual.dtype, x.device + ) + if bufs is None: + return None + + try: + _, _, layer_input_out, new_residual = mhc_post_pre( + x, + residual, + post, + comb, + hc_fn_t, + hc_scale, + hc_base, + hc_mult, + norm_eps, + hc_eps, + hc_post_mult, + sinkhorn_iters, + # Match sglang's exp-domain asymmetric Sinkhorn used in hc_pre. + asymmetric_exp_domain=True, + hc_sinkhorn_eps=hc_eps, + residual_out=bufs["residual_out"], + h_post=bufs["h_post"], + h_res=bufs["h_res"], + layer_input_out=bufs["layer_input_out"], + acc_partial=bufs["acc_partial"], + acc_sq_partial=bufs["acc_sq_partial"], + ) + except Exception as err: + logger.warning( + "Triton fused mHC kernel failed, disabling fallback path: %s", err + ) + _TRITON_MHC_POST_PRE_RUNTIME_DISABLED = True + return None + + return new_residual, layer_input_out, bufs["h_post"], bufs["h_res"], False diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index e041e3655..fc82de69f 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -87,6 +87,9 @@ from sglang.srt.model_executor.forward_context import ( from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.dbrx import ReplicatedLinear +from sglang.srt.models.deepseek_common.amd.deepseek_v4_fused_mhc import ( + try_fused_hc_post_pre, +) from sglang.srt.models.deepseek_v2 import ParallelLMHead, _is_cuda, _is_hip, _is_npu if not _is_hip: @@ -133,6 +136,28 @@ def _fused_rmsnorm_fp8_quant(hidden_states, weight, eps): return x_quant, x_bf16 +_FREQS_CIS_TO_COS_SIN: dict[ + Tuple[int, torch.dtype, torch.device], Tuple[torch.Tensor, torch.Tensor] +] = {} + + +def _freqs_cis_to_cos_sin( + freqs_cis: torch.Tensor, dtype: torch.dtype, device: torch.device +) -> Tuple[torch.Tensor, torch.Tensor]: + """Derive (cos, sin) bf16 contiguous tables from a complex64 `freqs_cis`, + cached by `(id(freqs_cis), dtype, device)` so that all layers sharing the + same `freqs_cis` (via `precompute_freqs_cis`'s lru_cache) reuse one pair.""" + key = (id(freqs_cis), dtype, device) + cached = _FREQS_CIS_TO_COS_SIN.get(key) + if cached is not None: + return cached + fr = torch.view_as_real(freqs_cis) + cos = fr[..., 0].to(device=device, dtype=dtype).contiguous() + sin = fr[..., 1].to(device=device, dtype=dtype).contiguous() + _FREQS_CIS_TO_COS_SIN[key] = (cos, sin) + return cos, sin + + if TYPE_CHECKING: from sglang.srt.layers.attention.deepseek_v4_backend import ( DeepseekV4AttnBackend, @@ -319,7 +344,7 @@ class MQALayer(nn.Module): ) self.attn_sink = nn.Parameter(torch.empty(self.n_heads, dtype=torch.float32)) - self.fuse_wqa_wkv = not _is_hip and envs.SGLANG_OPT_FUSE_WQA_WKV.get() + self.fuse_wqa_wkv = envs.SGLANG_OPT_FUSE_WQA_WKV.get() if self.fuse_wqa_wkv: self.wqkv_a = ReplicatedLinear( self.hidden_size, @@ -641,6 +666,7 @@ class MQALayer(nn.Module): x=x, q_lora=q_lora, forward_batch=forward_batch, + attn_backend=attn_backend, skip_compressor=True, ) elif self.compressor is not None: @@ -795,6 +821,7 @@ class MQALayer(nn.Module): and get_is_capture_mode() and x.shape[0] <= self._multi_stream_bs_limit and not (self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)) + and not (_is_hip and self.compressor is None) ) tp_slice, q_padded, q_out = slice(None), None, None @@ -1187,15 +1214,33 @@ class DeepseekV4DecoderLayer(nn.Module): x_quant=x_quant, ) - hidden_states = self.hc_post(hidden_states, residual, post, comb) - residual = hidden_states - hidden_states, post, comb, norm_fused = self.hc_pre( + fused_mhc = try_fused_hc_post_pre( hidden_states, - self.hc_ffn_fn, + residual, + post, + comb, + self.hc_ffn_fn.T, self.hc_ffn_scale, self.hc_ffn_base, - norm=self.post_attention_layernorm, + self.hc_mult, + self.rms_norm_eps, + self.hc_eps, + 2.0, + self.hc_sinkhorn_iters, + _is_gfx95_supported, ) + if fused_mhc is not None: + residual, hidden_states, post, comb, norm_fused = fused_mhc + else: + hidden_states = self.hc_post(hidden_states, residual, post, comb) + residual = hidden_states # [n, hc, d] + hidden_states, post, comb, norm_fused = self.hc_pre( + hidden_states, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + norm=self.post_attention_layernorm, + ) # -> [n, d] if not norm_fused: hidden_states = self.post_attention_layernorm(hidden_states) @@ -1759,7 +1804,7 @@ class DeepseekV4ForCausalLM(nn.Module): cache_compressor_weight = {} COMPRESSOR_PART = ".compressor.w" - fuse_wqa_wkv = not _is_hip and envs.SGLANG_OPT_FUSE_WQA_WKV.get() + fuse_wqa_wkv = envs.SGLANG_OPT_FUSE_WQA_WKV.get() cache_wqkv_a_weight: dict[str, dict[str, torch.Tensor]] = {} def auto_weight_loader(module): diff --git a/python/sglang/srt/speculative/draft_utils.py b/python/sglang/srt/speculative/draft_utils.py index b09e0c6a0..b465612b6 100644 --- a/python/sglang/srt/speculative/draft_utils.py +++ b/python/sglang/srt/speculative/draft_utils.py @@ -1,7 +1,7 @@ import logging from sglang.srt.server_args import ServerArgs, get_global_server_args -from sglang.srt.utils.common import is_blackwell, is_musa +from sglang.srt.utils.common import is_blackwell, is_hip, is_musa logger = logging.getLogger(__name__) @@ -226,9 +226,14 @@ class DraftBackendFactory: ) def _create_dsv4_decode_backend(self): - from sglang.srt.layers.attention.deepseek_v4_backend import ( - DeepseekV4MultiStepBackend, - ) + if is_hip(): + from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import ( + DeepseekV4MultiStepBackend, + ) + else: + from sglang.srt.layers.attention.deepseek_v4_backend import ( + DeepseekV4MultiStepBackend, + ) return DeepseekV4MultiStepBackend( self.draft_model_runner, self.topk, self.speculative_num_steps @@ -318,6 +323,14 @@ class DraftBackendFactory: return None def _create_dsv4_prefill_backend(self): + if is_hip(): + from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import ( + DeepseekV4HipRadixBackend, + ) + + return DeepseekV4HipRadixBackend( + self.draft_model_runner, skip_prefill=False + ) from sglang.srt.layers.attention.deepseek_v4_backend import ( DeepseekV4AttnBackend, ) diff --git a/test/registered/ops/test_aiter_greedy_sample_amd.py b/test/registered/ops/test_aiter_greedy_sample_amd.py new file mode 100644 index 000000000..662296561 --- /dev/null +++ b/test/registered/ops/test_aiter_greedy_sample_amd.py @@ -0,0 +1,289 @@ +"""Unit tests for aiter greedy_sample kernel and Sampler integration. + +Validates that: +1. aiter.greedy_sample produces identical results to torch.argmax (kernel level) +2. Sampler.forward() correctly dispatches to aiter when _use_aiter=True +3. The fallback to torch.argmax works when _use_aiter=False +4. return_logprob path works with the aiter greedy branch + +The kernel is designed for production LLM inference (large vocab, bf16) and is +used when SGLANG_USE_AITER=1 on ROCm. +""" + +import unittest +from unittest import mock + +import torch + +from sglang.srt.utils.common import is_hip +from sglang.test.ci.ci_register import register_amd_ci + +register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd") + + +def _mock_global_server_args(backend="pytorch"): + from sglang.srt.layers import sampler as sampler_mod + from sglang.srt.server_args import ServerArgs + + sampler_mod.get_global_server_args = lambda: ServerArgs( + model_path="dummy", + sampling_backend=backend, + ) + + class _DummyTPGroup: + device_group = None + + sampler_mod.get_tp_group = lambda: _DummyTPGroup() + sampler_mod.is_dp_attention_enabled = lambda: False + + +def _make_sampling_info(batch_size, vocab_size, device="cuda"): + from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo + + return SamplingBatchInfo( + temperatures=torch.ones(batch_size, 1, device=device, dtype=torch.float), + top_ps=torch.ones(batch_size, device=device), + top_ks=torch.zeros(batch_size, device=device, dtype=torch.int32), + min_ps=torch.zeros(batch_size, device=device), + is_all_greedy=True, + need_top_p_sampling=False, + need_top_k_sampling=False, + need_min_p_sampling=False, + vocab_size=vocab_size, + device=device, + ) + + +@unittest.skipUnless(is_hip(), "aiter greedy_sample requires ROCm") +class TestAiterGreedySample(unittest.TestCase): + """Kernel-level correctness: aiter.greedy_sample vs torch.argmax.""" + + @classmethod + def setUpClass(cls): + try: + from aiter import greedy_sample + + cls.greedy_sample = staticmethod(greedy_sample) + except ImportError: + raise unittest.SkipTest("aiter not installed") + cls.device = "cuda" + + def setUp(self): + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) + + def _run_and_compare(self, batch_size, vocab_size): + logits = torch.randn( + batch_size, vocab_size, device=self.device, dtype=torch.bfloat16 + ) + + expected = torch.argmax(logits, dim=-1) + + actual = torch.empty(logits.shape[0], device=logits.device, dtype=torch.int32) + self.greedy_sample(actual, logits) + + self.assertTrue( + torch.equal(actual.to(expected.dtype), expected), + f"Mismatch for shape ({batch_size}, {vocab_size}): " + f"expected={expected[:8].tolist()}, got={actual[:8].tolist()}", + ) + + def test_single_request(self): + self._run_and_compare(1, 32000) + + def test_small_batch(self): + self._run_and_compare(4, 32000) + + def test_medium_batch(self): + self._run_and_compare(32, 32000) + + def test_large_batch(self): + self._run_and_compare(128, 32000) + + def test_realistic_vocab_deepseek(self): + self._run_and_compare(64, 129280) + + def test_realistic_vocab_llama3(self): + self._run_and_compare(64, 128256) + + def test_realistic_vocab_qwen(self): + self._run_and_compare(64, 151936) + + def test_various_batch_sizes(self): + configs = [ + (1, 128256), + (2, 128256), + (8, 128256), + (16, 128256), + (32, 129280), + (64, 129280), + (128, 129280), + (256, 129280), + ] + for batch_size, vocab_size in configs: + with self.subTest(batch_size=batch_size, vocab_size=vocab_size): + self._run_and_compare(batch_size, vocab_size) + + def test_tied_values(self): + vocab_size = 32000 + logits = torch.zeros(8, vocab_size, device=self.device, dtype=torch.bfloat16) + logits[:, 0] = 1.0 + + expected = torch.argmax(logits, dim=-1) + actual = torch.empty(8, device=self.device, dtype=torch.int32) + self.greedy_sample(actual, logits) + + self.assertTrue(torch.equal(actual.to(expected.dtype), expected)) + + def test_negative_logits(self): + vocab_size = 32000 + logits = ( + torch.randn(16, vocab_size, device=self.device, dtype=torch.bfloat16) - 5.0 + ) + + expected = torch.argmax(logits, dim=-1) + actual = torch.empty(16, device=self.device, dtype=torch.int32) + self.greedy_sample(actual, logits) + + self.assertTrue(torch.equal(actual.to(expected.dtype), expected)) + + def test_extreme_values(self): + vocab_size = 32000 + logits = torch.randn(16, vocab_size, device=self.device, dtype=torch.bfloat16) + logits[0, 42] = 1e4 + logits[1, 100] = -1e4 + + expected = torch.argmax(logits, dim=-1) + actual = torch.empty(16, device=self.device, dtype=torch.int32) + self.greedy_sample(actual, logits) + + self.assertTrue(torch.equal(actual.to(expected.dtype), expected)) + + +@unittest.skipUnless(is_hip(), "aiter greedy_sample requires ROCm") +class TestAiterGreedyIntegration(unittest.TestCase): + """Integration: Sampler.forward() with _use_aiter on/off.""" + + @classmethod + def setUpClass(cls): + try: + from aiter import greedy_sample + + cls._greedy_sample_fn = staticmethod(greedy_sample) + except ImportError: + raise unittest.SkipTest("aiter not installed") + cls.device = "cuda" + + def setUp(self): + torch.manual_seed(42) + torch.cuda.manual_seed_all(42) + + def _run_sampler(self, use_aiter, logits, sampling_info, return_logprob=False): + from sglang.srt.layers import sampler as sampler_mod + from sglang.srt.layers.logits_processor import LogitsProcessorOutput + + _mock_global_server_args() + + patches = {"_use_aiter": use_aiter} + if use_aiter: + patches["_aiter_greedy_sample"] = self._greedy_sample_fn + + with mock.patch.multiple(sampler_mod, **patches): + sampler = sampler_mod.Sampler() + batch_size = logits.shape[0] + positions = torch.arange(batch_size, device=self.device, dtype=torch.int32) + + return sampler.forward( + logits_output=LogitsProcessorOutput(next_token_logits=logits.clone()), + sampling_info=sampling_info, + return_logprob=return_logprob, + top_logprobs_nums=[0] * batch_size, + token_ids_logprobs=[None] * batch_size, + positions=positions, + ) + + def test_aiter_matches_argmax_through_sampler(self): + batch_size, vocab_size = 64, 129280 + logits = torch.randn( + batch_size, vocab_size, device=self.device, dtype=torch.bfloat16 + ) + sampling_info = _make_sampling_info(batch_size, vocab_size, self.device) + + out_aiter = self._run_sampler(True, logits, sampling_info) + out_argmax = self._run_sampler(False, logits, sampling_info) + + self.assertTrue( + torch.equal(out_aiter.cpu(), out_argmax.cpu()), + f"Sampler mismatch: aiter={out_aiter[:8].tolist()}, " + f"argmax={out_argmax[:8].tolist()}", + ) + + def test_fallback_to_argmax_when_disabled(self): + batch_size, vocab_size = 32, 32000 + logits = torch.randn( + batch_size, vocab_size, device=self.device, dtype=torch.bfloat16 + ) + sampling_info = _make_sampling_info(batch_size, vocab_size, self.device) + + out = self._run_sampler(False, logits, sampling_info) + expected = torch.argmax(logits, dim=-1) + + self.assertTrue( + torch.equal(out.cpu(), expected.cpu()), + "Fallback path should produce torch.argmax results", + ) + + def test_aiter_greedy_with_return_logprob(self): + batch_size, vocab_size = 16, 32000 + logits = torch.randn( + batch_size, vocab_size, device=self.device, dtype=torch.bfloat16 + ) + sampling_info = _make_sampling_info(batch_size, vocab_size, self.device) + + out_aiter = self._run_sampler(True, logits, sampling_info, return_logprob=True) + out_argmax = self._run_sampler( + False, logits, sampling_info, return_logprob=True + ) + + self.assertTrue( + torch.equal(out_aiter.cpu(), out_argmax.cpu()), + "Token IDs should match with return_logprob=True", + ) + + def test_aiter_output_dtype(self): + """Document that the aiter path returns int32 (vs int64 from argmax).""" + batch_size, vocab_size = 16, 32000 + logits = torch.randn( + batch_size, vocab_size, device=self.device, dtype=torch.bfloat16 + ) + sampling_info = _make_sampling_info(batch_size, vocab_size, self.device) + + out_aiter = self._run_sampler(True, logits, sampling_info) + out_argmax = self._run_sampler(False, logits, sampling_info) + + self.assertEqual(out_aiter.dtype, torch.int32) + self.assertEqual(out_argmax.dtype, torch.int64) + + def test_various_batch_sizes_through_sampler(self): + vocab_size = 129280 + for batch_size in [1, 4, 16, 64, 128]: + with self.subTest(batch_size=batch_size): + logits = torch.randn( + batch_size, + vocab_size, + device=self.device, + dtype=torch.bfloat16, + ) + sampling_info = _make_sampling_info(batch_size, vocab_size, self.device) + + out_aiter = self._run_sampler(True, logits, sampling_info) + out_argmax = self._run_sampler(False, logits, sampling_info) + + self.assertTrue( + torch.equal(out_aiter.cpu(), out_argmax.cpu()), + f"Mismatch at batch_size={batch_size}", + ) + + +if __name__ == "__main__": + unittest.main()