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

runs/frontier-qwen30-vllm020-profile-v1/jobs_frontier_linear_smoke.toml

@@ -1,18 +1,18 @@
 version = 1
 
 [[jobs]]
-name = "qwen30-vllm020-frontier-linear-smoke-20260716-v3-model-context"
+name = "qwen30-vllm020-frontier-linear-smoke-20260716-v4-rmsnorm-api"
 gpus = 1
 gpu_model = "H20"
 hosts = ["dash0"]
 command = "cd /home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-v1/runs/frontier-qwen30-vllm020-profile-v1 && timeout --signal=TERM --kill-after=30s 1020 bash run_frontier_linear_smoke.sh"
-artifacts = ["artifacts/frontier-linear-smoke-v3"]
+artifacts = ["artifacts/frontier-linear-smoke-v4"]
 
 [jobs.env]
 HOME = "/tmp/wjh"
 XDG_CACHE_HOME = "/tmp/wjh/.cache"
 VLLM_CACHE_ROOT = "/tmp/wjh/.cache/vllm"
-OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/frontier-linear-smoke-v3"
+OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/frontier-linear-smoke-v4"
 VENV_ROOT = "/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1"
 FRONTIER_ROOT = "/home/admin/cpfs/wjh/frontier-qwen30-vllm020-profile-v1/Frontier"
 VLLM_SOURCE = "/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build"