Adapt Frontier RMSNorm profiling to vLLM 0.20

runs/frontier-qwen30-vllm020-profile-v1/frontier_vllm020_compat.py

@@ -13,8 +13,14 @@ import os
 from typing import Any
 
 import torch
+from vllm import _custom_ops as vllm_ops
+from vllm.model_executor.layers.layernorm import (
+    GemmaRMSNorm as VllmGemmaRMSNorm,
+    fused_add_rms_norm as vllm_fused_add_rms_norm,
+)
 from vllm.model_executor.layers.rotary_embedding import get_rope as vllm_get_rope
 
+from frontier.profiling.common.layers import layernorm as frontier_layernorm
 from frontier.profiling.common.layers import rotary_embedding as frontier_rope
 
 
@@ -40,6 +46,16 @@ def _vllm020_get_rope_adapter(
     )
 
 
+def _vllm020_rms_norm_adapter(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    variance_epsilon: float,
+) -> torch.Tensor:
+    output = torch.empty_like(x)
+    vllm_ops.rms_norm(output, x, weight, variance_epsilon)
+    return output
+
+
 def main() -> None:
     parameters = inspect.signature(vllm_get_rope).parameters
     if "rope_parameters" not in parameters or "rotary_dim" in parameters:
@@ -50,6 +66,10 @@ def main() -> None:
 
     frontier_rope._VLLM_GET_ROPE = _vllm020_get_rope_adapter
     frontier_rope._VLLM_GET_ROPE_IMPORT_ERROR = None
+    frontier_layernorm.HAS_VLLM_RMSNORM = True
+    frontier_layernorm.VllmGemmaRMSNorm = VllmGemmaRMSNorm
+    frontier_layernorm.vllm_rms_norm = _vllm020_rms_norm_adapter
+    frontier_layernorm.vllm_fused_add_rms_norm = vllm_fused_add_rms_norm
 
     from frontier.profiling.linear_op.main import main as frontier_main
     from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config