[AMD][Fix] Qwen3.5: guard zero-grid launch in fused_qk_gemma_rmsnorm(_with_gate) (HIP invalid configuration on idle DP rank)

python/sglang/kernels/ops/attention/triton_gdn_fused_proj.py

@@ -127,6 +127,8 @@ def fused_qkvzba_split_reshape_cat(
         device=mixed_ba.device,
     )
     a = torch.empty_like(b)
+    if 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,
@@ -292,6 +294,8 @@ def fused_qkvzba_split_reshape_cat_contiguous(
         device=mixed_ba.device,
     )
     a = torch.empty_like(b)
+    if batch * seq_len == 0:
+        return mixed_qkv, z, b, a
     grid = (batch * seq_len, num_heads_qk)
     fused_qkvzba_split_reshape_cat_contiguous_kernel[grid](
         mixed_qkv,
@@ -387,6 +391,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 seq_len == 0:
+        return q, k, v
     fused_qkv_split_gdn_prefill_kernel[(seq_len,)](
         q,
         k,

python/sglang/srt/layers/attention/linear/gdn_backend.py

@@ -380,6 +380,9 @@ def forward_decode(
         b: torch.Tensor,
         **kwargs,
     ):
+        if 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
@@ -474,6 +477,9 @@ def forward_extend(
         assert isinstance(mixed_qkv, torch.Tensor)
         seq_len = mixed_qkv.shape[0]
 
+        if 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
 

python/sglang/srt/layers/sampler.py

@@ -119,6 +119,9 @@ def forward(
         """
         logits = logits_output.next_token_logits
 
+        if 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 [])

python/sglang/srt/models/qwen2_moe.py

@@ -326,6 +326,7 @@ def __init__(
                     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 @@ def __init__(
         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 = (
@@ -560,7 +561,7 @@ def forward(
         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 = (

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