amd/tiny-qwen3-moe-w4a8

Model Overview

• Model Architecture: Qwen3MoeForCausalLM (tiny, randomly initialized)
  • Input: Text
  • Output: Text
• Supported Hardware Microarchitecture: AMD MI300 / MI350 / MI355 (gfx942 / gfx950)
• Inference Engine: vLLM
• Model Optimizer: AMD-Quark
  • Weight quantization: W4A8 — INT4 weights (per-channel, symmetric) produced via a progressive FP8→INT4 spec, following the amd/Kimi-K2.5-W4A8 recipe
  • Activation quantization: FP8 E4M3, per-tensor, dynamic
  • Quantized layers: routed MoE experts only (attention, router/gate, and lm_head are kept in the original precision)

This is a tiny, randomly-initialized Qwen3-MoE model quantized to W4A8, used purely as vLLM CI coverage for the Quark W4A8 fused-MoE path (QuarkW4A8Fp8MoEMethod), which dispatches through the ROCm AITER fused MoE kernel. It is not intended to produce meaningful text — it exists so CI can load a real W4A8 checkpoint and run a forward pass on GPU.

The dimensions (hidden 2048, MoE intermediate 1024, 8 experts, top-2) are multiples of 256 so the AITER W4A8 shuffle/GEMM tile constraints hold. The vocab_size matches the tokenizer so token ids stay within the embedding table.

Model Creation

Built and quantized with AMD-Quark, following the progressive FP8→INT4 weight spec from the amd/Kimi-K2.5-W4A8 model card.

Note: Quark quantizes nn.Linear modules. MoE experts are stored as individual nn.Linear layers in transformers ~4.57; quantize with that version so the routed experts are captured.

import argparse

import torch
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

from quark.torch import ModelQuantizer, export_safetensors
from quark.torch.quantization.config.config import (
    FP8E4M3PerTensorSpec,
    Int4PerChannelSpec,
    ProgressiveSpec,
    QConfig,
    QLayerConfig,
)


def get_config() -> QConfig:
    # Quantize the routed experts only.
    exclude_layers = ["*self_attn*", "*mlp.gate", "*lm_head"]
    input_spec = FP8E4M3PerTensorSpec(
        observer_method="min_max", scale_type="float", is_dynamic=True
    ).to_quantization_spec()
    # Progressive FP8 -> INT4 weight spec (Kimi-K2.5-W4A8 recipe).
    weight_spec = ProgressiveSpec(
        first_stage=FP8E4M3PerTensorSpec(
            observer_method="min_max", scale_type="float", is_dynamic=False
        ),
        second_stage=Int4PerChannelSpec(
            symmetric=True,
            scale_type="float",
            round_method="half_even",
            is_dynamic=False,
            ch_axis=0,
        ),
    ).to_quantization_spec()
    return QConfig(
        global_quant_config=QLayerConfig(input_tensors=input_spec, weight=weight_spec),
        exclude=exclude_layers,
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--export-path", required=True)
    parser.add_argument("--tokenizer", default="Qwen/Qwen1.5-MoE-A2.7B-Chat")
    parser.add_argument("--hidden", type=int, default=2048)
    parser.add_argument("--moe-intermediate", type=int, default=1024)
    parser.add_argument("--experts", type=int, default=8)
    parser.add_argument("--topk", type=int, default=2)
    parser.add_argument("--layers", type=int, default=2)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    torch.manual_seed(args.seed)
    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)

    # vocab_size MUST cover the tokenizer, else real prompts produce token ids
    # beyond the embedding table -> out-of-bounds embedding lookup (GPU fault).
    cfg = AutoConfig.for_model(
        "qwen3_moe",
        hidden_size=args.hidden,
        intermediate_size=args.hidden,
        moe_intermediate_size=args.moe_intermediate,
        num_hidden_layers=args.layers,
        num_attention_heads=16,
        num_key_value_heads=2,
        head_dim=128,
        num_experts=args.experts,
        num_experts_per_tok=args.topk,
        vocab_size=len(tokenizer),
        max_position_embeddings=2048,
    )
    model = AutoModelForCausalLM.from_config(cfg).to("cuda").eval().to(torch.bfloat16)

    ds = load_dataset("mit-han-lab/pile-val-backup", split="validation")
    samples = [
        tokenizer(ds[i]["text"], return_tensors="pt", truncation=True,
                  max_length=64).input_ids.to("cuda")
        for i in range(8)
    ]
    dataloader = DataLoader(samples, batch_size=1)

    quantizer = ModelQuantizer(get_config())
    with torch.no_grad():
        model = quantizer.quantize_model(model, dataloader)

    export_safetensors(
        model, args.export_path, custom_mode="quark",
        weight_format="real_quantized", pack_method="reorder",
    )
    tokenizer.save_pretrained(args.export_path)
    # Symmetric INT4 export emits all-zero `*_zero_point_2` tensors that vLLM's
    # W4A8 loader does not expect; drop them so the checkpoint loads directly.


if __name__ == "__main__":
    main()

Usage in vLLM

W4A8 dispatches through the ROCm AITER fused MoE kernel, so run on gfx942/gfx950 with AITER enabled:

VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 \
  vllm serve amd/tiny-qwen3-moe-w4a8 --enforce-eager

Because the weights are random, outputs are not meaningful — this model is a structural / smoke-test fixture only.

License

Apache-2.0. The tiny model is randomly initialized and derives no weights from any base model.

Modifications Copyright(c) 2026 Advanced Micro Devices, Inc. All rights reserved.