update qwen3_moe of layernorm and activation

vllm/attention/layer.py

@@ -29,7 +29,7 @@ try:
     tag_cudagraph_unsafe = (torch._C.Tag.cudagraph_unsafe, )
 except AttributeError:
     tag_cudagraph_unsafe = ()  # type: ignore[assignment]
-
+    
 
 class Attention(nn.Module):
     """Attention layer.
@@ -220,8 +220,8 @@ class Attention(nn.Module):
             output_shape = (output_shape
                             if output_shape is not None else query.shape)
             output = torch.zeros(output_shape,
-                                 dtype=query.dtype,
-                                 device=query.device)
+                                dtype=query.dtype,
+                                device=query.device)
             hidden_size = output_shape[-1]
             # We skip reshaping query, key and value tensors for the MLA
             # backend since these tensors have different semantics and are

vllm/envs.py

@@ -1124,7 +1124,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # vLLM will use lightop moe_align_block_size 
     "VLLM_USE_LIGHTOP_MOE_ALIGN":
         lambda: (os.environ.get("VLLM_USE_LIGHTOP_MOE_ALIGN", "True").lower() in
-                 ("true", "1")),     
+                 ("true", "1")),    
     # vLLM will use opt merge_aatn_states, not triton
     "VLLM_USE_MERGE_ATTN_STATES_OPT":
         lambda: (os.environ.get("VLLM_USE_MERGE_ATTN_STATES_OPT", "True").lower() in

vllm/model_executor/layers/activation.py

@@ -77,7 +77,7 @@ class SiluAndMul(CustomOp):
         """PyTorch-native implementation equivalent to forward()."""
         if not torch.compiler.is_compiling() and envs.VLLM_ENABLE_TBO:
             return self.forward_cuda(x)
-        elif envs.VLLM_USE_OPT_OP:
+        elif not torch.compiler.is_compiling() and envs.VLLM_USE_OPT_OP:
             return self.forward_cuda(x)
         else:
             d = x.shape[-1] // 2

vllm/model_executor/layers/layernorm.py

@@ -167,7 +167,7 @@ class RMSNorm(CustomOp):
     ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
         if not torch.compiler.is_compiling() and envs.VLLM_ENABLE_TBO:
             return self.forward_cuda(x, residual) 
-        elif envs.VLLM_USE_OPT_OP:
+        elif not torch.compiler.is_compiling() and envs.VLLM_USE_OPT_OP:
             return self.forward_cuda(x, residual)  
         else:
             orig_dtype = x.dtype

vllm/model_executor/models/qwen3_moe.py

@@ -234,7 +234,7 @@ class Qwen3MoeAttention(nn.Module):
         if envs.VLLM_USE_APEX_RN:
             q_by_head = self.q_norm.forward_apex(q_by_head)
         else:
-            q_by_head = self.q_norm(q_by_head)
+            q_by_head = self.q_norm.forward_cuda(q_by_head)
         q = q_by_head.view(q.shape)
 
         k_by_head = k.view(*k.shape[:-1], k.shape[-1] // self.head_dim,
@@ -242,7 +242,7 @@ class Qwen3MoeAttention(nn.Module):
         if envs.VLLM_USE_APEX_RN:
             k_by_head = self.k_norm.forward_apex(k_by_head)
         else:
-            k_by_head = self.k_norm(k_by_head)
+            k_by_head = self.k_norm.forward_cuda(k_by_head)
         k = k_by_head.view(k.shape)
         q, k = self.rotary_emb(positions, q, k)
         attn_output = self.attn(q, k, v)