diff --git a/runs/frontier-qwen30-vllm020-profile-v1/frontier_vllm020_compat.py b/runs/frontier-qwen30-vllm020-profile-v1/frontier_vllm020_compat.py index 9730897..4366c9c 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/frontier_vllm020_compat.py +++ b/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 diff --git a/runs/frontier-qwen30-vllm020-profile-v1/jobs_frontier_linear_smoke.toml b/runs/frontier-qwen30-vllm020-profile-v1/jobs_frontier_linear_smoke.toml index 4baf066..2208aa7 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/jobs_frontier_linear_smoke.toml +++ b/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"