diff --git a/runs/frontier-fidelity-envelope-v1/smoke_qwen235_v020_frontier_moe.py b/runs/frontier-fidelity-envelope-v1/smoke_qwen235_v020_frontier_moe.py new file mode 100644 index 0000000..5157422 --- /dev/null +++ b/runs/frontier-fidelity-envelope-v1/smoke_qwen235_v020_frontier_moe.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""One-cell gate for the two Qwen235 vLLM 0.20 MoE runtime backends.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import torch + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--frontier-source", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def routing(tokens: int, experts: int = 128, topk: int = 8): + ids = torch.arange(tokens * topk, device="cuda", dtype=torch.int64) + ids = (ids % experts).view(tokens, topk) + weights = torch.full((tokens, topk), 1.0 / topk, device="cuda") + return weights, ids + + +def main() -> None: + args = parse_args() + sys.path.insert(0, str(args.frontier_source.resolve())) + from frontier.profiling.moe.moe_vllm_kernel import profile_fused_moe_kernel + + weights, ids = routing(8) + cells = [] + cells.append( + { + "name": "tp4_ep1_triton", + "stats": profile_fused_moe_kernel( + num_tokens=8, + num_experts=128, + hidden_dim=4096, + expert_hidden_dim=1536, + top_k=8, + topk_weights=weights, + topk_ids=ids, + tensor_parallel_size=4, + use_fp8=True, + block_shape=[128, 128], + warmup_steps=1, + active_steps=2, + ), + } + ) + expert_map = torch.full((128,), -1, device="cuda", dtype=torch.int32) + expert_map[:16] = torch.arange(16, device="cuda", dtype=torch.int32) + cells.append( + { + "name": "tp1_ep8_flashinfer_cutlass", + "stats": profile_fused_moe_kernel( + num_tokens=8, + num_experts=16, + hidden_dim=4096, + expert_hidden_dim=1536, + top_k=8, + topk_weights=weights, + topk_ids=ids, + tensor_parallel_size=1, + use_fp8=True, + block_shape=[128, 128], + warmup_steps=1, + active_steps=2, + global_num_experts=128, + expert_map=expert_map, + ), + } + ) + payload = {"schema": "qwen235-v020-frontier-moe-smoke-v1", "cells": cells} + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + print(json.dumps(payload, sort_keys=True)) + + +if __name__ == "__main__": + main()