85 lines
2.5 KiB
Python
85 lines
2.5 KiB
Python
#!/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()
|