From 6314e9e4f5bdadd975969e8781dd1edcb27bb4a7 Mon Sep 17 00:00:00 2001 From: "jacky.cheng" Date: Sun, 16 Aug 2026 12:02:54 +0800 Subject: [PATCH] [AMD][Fix] Qwen3.5: guard zero-grid launch in fused_qk_gemma_rmsnorm(_with_gate) (HIP invalid configuration on idle DP rank) (#31794) --- .../sglang/kernels/ops/attention/triton_gdn_fused_proj.py | 6 ++++++ python/sglang/srt/layers/attention/linear/gdn_backend.py | 8 ++++++++ python/sglang/srt/layers/sampler.py | 4 ++++ python/sglang/srt/models/qwen2_moe.py | 5 +++-- python/sglang/srt/models/utils.py | 6 ++++++ 5 files changed, 27 insertions(+), 2 deletions(-) diff --git a/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py b/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py index 8e39197ba..64d81b9d3 100644 --- a/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py +++ b/python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py @@ -132,6 +132,8 @@ def fused_qkvzba_split_reshape_cat( device=mixed_ba.device, ) a = torch.empty_like(b) + if _is_hip and batch * seq_len == 0: + return mixed_qkv, z, b, a grid = (batch * seq_len, num_heads_qk) fused_qkvzba_split_reshape_cat_kernel[grid]( mixed_qkv, @@ -297,6 +299,8 @@ def fused_qkvzba_split_reshape_cat_contiguous( device=mixed_ba.device, ) a = torch.empty_like(b) + if _is_hip and batch * seq_len == 0: + return mixed_qkv, z, b, a grid = (batch * seq_len, num_heads_qk) # Each program moves `v_per_group * head_v` elements for both v and z. For # the small head-group ratios (<= 512 elements) a single warp is the best @@ -402,6 +406,8 @@ def fused_qkv_split_gdn_prefill( ) qkv_dim = num_q_heads * head_q + num_k_heads * head_k + num_v_heads * head_v + if _is_hip and seq_len == 0: + return q, k, v fused_qkv_split_gdn_prefill_kernel[(seq_len,)]( q, k, diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 0bc8fbd56..8fdec861a 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -22,6 +22,8 @@ from sglang.srt.runtime_context import get_exec, get_memory, get_schedule from sglang.srt.utils import is_cpu, is_cuda, is_hip, is_npu, is_xpu from sglang.srt.utils.common import rank0_log +_is_hip = is_hip() + if not is_cpu(): from sglang.kernels.ops.attention.fla.chunk_delta_h import ( CHUNK_SIZE as FLA_CHUNK_SIZE, @@ -398,6 +400,9 @@ class GDNAttnBackend(MambaAttnBackendBase): b: torch.Tensor, **kwargs, ): + if _is_hip and isinstance(mixed_qkv, torch.Tensor) and mixed_qkv.shape[0] == 0: + return mixed_qkv.new_zeros((1, 0, layer.num_v_heads, layer.head_v_dim)) + layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id) conv_states = layer_cache.conv[0] ssm_states = layer_cache.temporal @@ -492,6 +497,9 @@ class GDNAttnBackend(MambaAttnBackendBase): assert isinstance(mixed_qkv, torch.Tensor) seq_len = mixed_qkv.shape[0] + if _is_hip and seq_len == 0: + return mixed_qkv.new_zeros((1, 0, layer.num_v_heads, layer.head_v_dim)) + is_target_verify = forward_batch.forward_mode.is_target_verify() forward_metadata = self.forward_metadata diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index d02db43fe..b491ed0bb 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -45,6 +45,7 @@ if is_musa(): top_p_renorm_prob, ) +_is_hip = is_hip() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip() if _use_aiter: from aiter import greedy_sample as _aiter_greedy_sample @@ -119,6 +120,9 @@ class Sampler(nn.Module): """ logits = logits_output.next_token_logits + if _is_hip and logits.shape[0] == 0: + return torch.empty((0,), dtype=torch.int64, device=logits.device) + # Preprocess logits (custom processors and NaN handling) logits = self._preprocess_logits(logits, sampling_info) return_sampling_mask = any(sampling_info.return_sampling_masks or []) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index dabfedbf1..d01ea904f 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -326,6 +326,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module): dict(tp_rank=0, tp_size=1) if ( get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_mori() or get_moe_a2a_backend().is_flashinfer() ) else {} @@ -344,7 +345,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module): else: self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) - if get_moe_a2a_backend().is_deepep(): + if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori(): # TODO: we will support tp < ep in the future self.ep_size = get_parallel().moe_ep_size self.num_experts = ( @@ -590,7 +591,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module): num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) - if get_moe_a2a_backend().is_deepep(): + if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori(): return self._forward_deepep(hidden_states, forward_batch) use_fused_gate = ( diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 3a11b5a6f..b9b18a6dd 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -593,6 +593,9 @@ def fused_qk_gemma_rmsnorm( q_out = torch.empty(q_rows, head_dim, dtype=q.dtype, device=q.device) k_out = torch.empty(k_rows, head_dim, dtype=k.dtype, device=k.device) + if _is_hip and q_rows == 0: + return q_out, k_out + BLOCK_HD = triton.next_power_of_2(head_dim) _fused_qk_gemma_rmsnorm_kernel[(q_rows,)]( @@ -709,6 +712,9 @@ def fused_qk_gemma_rmsnorm_with_gate( k_out = torch.empty(k_rows, head_dim, dtype=k.dtype, device=k.device) gate_out = torch.empty(q_rows, head_dim, dtype=q_gate.dtype, device=q_gate.device) + if _is_hip and q_rows == 0: + return q_out, k_out, gate_out + BLOCK_HD = triton.next_power_of_2(head_dim) _fused_qk_gemma_rmsnorm_gate_kernel[(q_rows,)](