[torch.compile] support all attention backends (#10558)

Signed-off-by: youkaichao <youkaichao@gmail.com>

vllm/model_executor/models/phimoe.py

@@ -294,6 +294,7 @@ class PhiMoEAttention(nn.Module):
         cache_config: Optional[CacheConfig] = None,
         quant_config: Optional[QuantizationConfig] = None,
         rope_scaling: Optional[dict] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
@@ -347,6 +348,7 @@ class PhiMoEAttention(nn.Module):
             num_kv_heads=self.num_kv_heads,
             cache_config=cache_config,
             quant_config=quant_config,
+            prefix=f"{prefix}.attn",
         )
 
     def forward(
@@ -371,6 +373,7 @@ class PhiMoEDecoderLayer(nn.Module):
         config: PhiMoEConfig,
         cache_config: Optional[CacheConfig] = None,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -385,6 +388,7 @@ class PhiMoEDecoderLayer(nn.Module):
             cache_config=cache_config,
             quant_config=quant_config,
             rope_scaling=config.rope_scaling,
+            prefix=f"{prefix}.self_attn",
         )
         self.block_sparse_moe = PhiMoE(
             num_experts=config.num_local_experts,
@@ -454,8 +458,8 @@ class PhiMoEModel(nn.Module):
         )
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
-            lambda prefix: PhiMoEDecoderLayer(config, cache_config,
-                                              quant_config),
+            lambda prefix: PhiMoEDecoderLayer(
+                config, cache_config, quant_config, prefix=prefix),
             prefix=f"{prefix}.layers")
         self.norm = nn.LayerNorm(config.hidden_size,
                                  eps=config.rms_norm_eps,