Add Qwen235 Frontier MoE portability smoke
This commit is contained in:
@@ -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()
|
||||||
Reference in New Issue
Block a user