[AMD][DSV4] DSV4 MTP graph + sparse triton attn optimizations (#26383)

Co-authored-by: wunhuang <wunhuang@amd.com>
Co-authored-by: Thomas Wang <1am9trash@gmail.com>
Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
Co-authored-by: amd-danli103 <danli103@amd.com>
Co-authored-by: Lin, Soga <soga.lin@amd.com>
Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com>
Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Co-authored-by: yichiche@amd.com <jacky.cheng>
Co-authored-by: yctseng0211 <yctseng@amd.com>
Co-authored-by: Bingxu Chen <bingxche@amd.com>

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)

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

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

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,

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

python/sglang/srt/models/deepseek_common/amd/__init__.py

@@ -0,0 +1 @@
+# AMD-specific DeepSeek common model helpers.

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

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):

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,
         )

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()