kvcache-ai--ktransformers/kt-kernel/examples/test_fp4_moe_amx.py

import math
import os
import sys
from typing import Dict

sys.path.insert(0, os.path.dirname(__file__) + "/../build")

import torch
from kt_kernel import kt_kernel_ext

torch.manual_seed(42)

hidden_size = 7168
intermediate_size = 2048
max_len = 25600

expert_num = 16
num_experts_per_tok = 8

layer_num = 1
CPUInfer = kt_kernel_ext.CPUInfer(40)
validation_iter = 3
k_group_size = 32
debug_print_count = 16

# Forward dispatch in do_gate_up_gemm uses `qlen > 4 * expert_num / top_k`
# (= 8 with these constants), so qlen=1 hits mat-vec and qlen=32 hits the
# mat-mat 4×4 register tile (per-expert avg m = qlen*top_k/expert_num = 16).
QLEN_LIST = [1, 32]
DISPATCH_THRESHOLD = 4 * expert_num / num_experts_per_tok

physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous()

# E2M1 values: {0, ±0.5, ±1.0, ±1.5, ±2.0, ±3.0, ±4.0, ±6.0}
E2M1_VALUES = torch.tensor([
    0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
    -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0,
], dtype=torch.float32)

# Nibble encoding: index = 4-bit value
# 0b0000..0b0111 = positive, 0b1000..0b1111 = negative
E2M1_NIBBLE_MAP = torch.tensor([
    0,   # 0b0000 = 0.0
    1,   # 0b0001 = 0.5
    2,   # 0b0010 = 1.0
    3,   # 0b0011 = 1.5
    4,   # 0b0100 = 2.0
    5,   # 0b0101 = 3.0
    6,   # 0b0110 = 4.0
    7,   # 0b0111 = 6.0
    8,   # 0b1000 = -0.0
    9,   # 0b1001 = -0.5
    10,  # 0b1010 = -1.0
    11,  # 0b1011 = -1.5
    12,  # 0b1100 = -2.0
    13,  # 0b1101 = -3.0
    14,  # 0b1110 = -4.0
    15,  # 0b1111 = -6.0
], dtype=torch.int32)


def _pattern_uniform(groups: int) -> torch.Tensor:
    return torch.full((groups,), 0.02, dtype=torch.float32)


def _pattern_alternating(groups: int) -> torch.Tensor:
    vals = torch.full((groups,), 0.015, dtype=torch.float32)
    vals[1::2] = 0.03
    return vals


def _pattern_ramp(groups: int) -> torch.Tensor:
    return torch.linspace(0.005, 0.04, steps=groups, dtype=torch.float32)


WEIGHT_PATTERNS = {
    "uniform_scale": ("All k-groups share the same abs max / scale", _pattern_uniform),
    "alternating_scale": ("Alternate small / large abs max per k-group", _pattern_alternating),
    "ramp_scale": ("Linearly increasing abs max per k-group", _pattern_ramp),
    "random": ("Random bf16 weights (baseline)", None),
}


def act_fn(x):
    return x / (1.0 + torch.exp(-x))


def mlp_torch(input, gate_proj, up_proj, down_proj):
    gate_buf = torch.mm(input, gate_proj.t())
    up_buf = torch.mm(input, up_proj.t())
    intermediate = act_fn(gate_buf) * up_buf
    ret = torch.mm(intermediate, down_proj.t())
    return ret


def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj):
    cnts = expert_ids.new_zeros((expert_ids.shape[0], expert_num))
    cnts.scatter_(1, expert_ids, 1)
    tokens_per_expert = cnts.sum(dim=0)
    idxs = expert_ids.view(-1).argsort()
    sorted_tokens = input[idxs // expert_ids.shape[1]]

    outputs = []
    start_idx = 0

    for i, num_tokens in enumerate(tokens_per_expert):
        end_idx = start_idx + num_tokens
        if num_tokens == 0:
            continue
        tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
        expert_out = mlp_torch(tokens_for_this_expert, gate_proj[i], up_proj[i], down_proj[i])
        outputs.append(expert_out)
        start_idx = end_idx

    outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)

    new_x = torch.empty_like(outs)
    new_x[idxs] = outs
    t_output = (
        new_x.view(*expert_ids.shape, -1)
        .type(weights.dtype)
        .mul_(weights.unsqueeze(dim=-1))
        .sum(dim=1)
        .type(new_x.dtype)
    )
    return t_output


def quantize_mxfp4_tensor(weights: torch.Tensor, group_size: int):
    """
    MXFP4 E2M1 quantization per k-group.
    For each block of group_size (32) elements in K dimension:
      scale = max_abs / 6.0
      quantized = round(value / scale) to nearest E2M1 value
    Args:
        weights: [expert_num, rows (N), cols (K)] in bf16
    Returns:
        packed: int32 tensor storing 8 FP4 nibbles per int32, shape [expert_num, rows * (cols // 8)]
        scales: bfloat16 tensor with shape [expert_num, rows * (cols // group_size)]
    """
    weights_f32 = weights.to(torch.float32)
    e, rows, cols = weights_f32.shape
    if cols % group_size != 0 or cols % 2 != 0:
        raise ValueError(f"cols ({cols}) must be divisible by group_size ({group_size}) and 2")

    reshaped = weights_f32.view(e, rows, cols // group_size, group_size)
    max_abs = reshaped.abs().amax(dim=-1, keepdim=True)
    max_abs = torch.clamp(max_abs, min=1e-8)
    scales = (max_abs / 6.0).squeeze(-1)

    # Quantize: round(value / scale) to nearest E2M1 value
    normalized = reshaped / scales.unsqueeze(-1)

    # For each normalized value, find the closest E2M1 value
    # E2M1_VALUES shape: [16]
    e2m1_vals = E2M1_VALUES.view(1, 1, 1, 1, 16)  # broadcast over (e, rows, groups, group_size, 16)
    normalized_expanded = normalized.unsqueeze(-1)  # (e, rows, groups, group_size, 1)
    distances = torch.abs(normalized_expanded - e2m1_vals)
    closest_indices = distances.argmin(dim=-1)  # (e, rows, groups, group_size) — indices 0..15

    # Dequantized values for reference
    dequant = E2M1_VALUES[closest_indices].to(torch.float32) * scales.unsqueeze(-1)
    dequant = dequant.view(e, rows, cols)

    # Pack nibbles: each byte = (hi_nibble << 4) | lo_nibble
    # Column-major: consecutive K elements are consecutive nibbles
    # nibble at even K index goes to low nibble, odd K index goes to high nibble
    # But wait — looking at the kernel's mxfp4_to_bf16_32:
    #   lo = packed & 0x0F, hi = (packed >> 4) & 0x0F
    # And the interleaving is: [lo[0],hi[0], lo[1],hi[1], ...] = column order
    # So for column-major layout: each byte has lo nibble = element at col c, hi nibble = element at col c+1
    # But the weight buffer layout is: b_row[k_block] = 16 packed bytes covering 32 K elements
    # With column-major: K is the innermost dimension
    # For 32 elements per k_group: byte 0 = [K_0 | K_1], byte 1 = [K_2 | K_3], ..., byte 15 = [K_30 | K_31]
    # low nibble = even K index, high nibble = odd K index

    nibbles = closest_indices.to(torch.uint8)  # 0..15, each is already a 4-bit value
    nibbles = nibbles.view(e, rows, cols // 2, 2)
    lo = nibbles[..., 0]  # even K indices
    hi = nibbles[..., 1]  # odd K indices
    packed_bytes = (hi << 4) | lo  # low nibble first in memory (little-endian style)

    # Pack 4 bytes into int32
    bytes_view = packed_bytes.view(e, rows, cols // 8, 4)
    packed_int32 = (
        bytes_view[..., 0].to(torch.int32) |
        (bytes_view[..., 1].to(torch.int32) << 8) |
        (bytes_view[..., 2].to(torch.int32) << 16) |
        (bytes_view[..., 3].to(torch.int32) << 24)
    )
    packed_int32 = packed_int32.view(e, rows, cols // 8).contiguous()

    scales = scales.to(torch.bfloat16).contiguous().view(e, rows, cols // group_size).contiguous()

    return packed_int32, scales, dequant


def build_structured_tensor(shape: torch.Size, pattern: str) -> torch.Tensor:
    if pattern == "random":
        torch.manual_seed(42)
        return (torch.randn(shape, dtype=torch.bfloat16, device="cpu") / 100.0).contiguous()

    e, rows, cols = shape
    groups = cols // k_group_size
    group_builder = WEIGHT_PATTERNS[pattern][1]
    group_vals = group_builder(groups).to(torch.float32)
    block = group_vals.view(1, 1, groups, 1).expand(e, rows, groups, k_group_size).clone()
    row_signs = torch.where(
        (torch.arange(rows) % 2 == 0),
        torch.ones(rows, dtype=torch.float32),
        -torch.ones(rows, dtype=torch.float32),
    ).view(1, rows, 1, 1)
    col_offsets = torch.linspace(-0.0005, 0.0005, steps=k_group_size, dtype=torch.float32).view(1, 1, 1, k_group_size)
    block = block * row_signs + col_offsets
    return block.reshape(shape).to(torch.bfloat16).contiguous()


def prepare_mxfp4_quantized_weights(pattern: str) -> Dict[str, torch.Tensor]:
    if pattern not in WEIGHT_PATTERNS:
        raise ValueError(f"Unknown weight pattern: {pattern}")

    gate_proj = build_structured_tensor((expert_num, intermediate_size, hidden_size), pattern)
    up_proj = build_structured_tensor((expert_num, intermediate_size, hidden_size), pattern)
    down_proj = build_structured_tensor((expert_num, hidden_size, intermediate_size), pattern)

    gate_q, gate_scales, gate_dequant = quantize_mxfp4_tensor(gate_proj, k_group_size)
    up_q, up_scales, up_dequant = quantize_mxfp4_tensor(up_proj, k_group_size)
    down_q, down_scales, down_dequant = quantize_mxfp4_tensor(down_proj, k_group_size)

    return {
        "gate_qweight": gate_q.contiguous(),
        "up_qweight": up_q.contiguous(),
        "down_qweight": down_q.contiguous(),
        "gate_scales": gate_scales.contiguous(),
        "up_scales": up_scales.contiguous(),
        "down_scales": down_scales.contiguous(),
        "dequantized": {
            "gate_proj": gate_dequant.to(torch.bfloat16).contiguous(),
            "up_proj": up_dequant.to(torch.bfloat16).contiguous(),
            "down_proj": down_dequant.to(torch.bfloat16).contiguous(),
        },
    }


def build_moes_from_quantized_data(quant_data: Dict[str, torch.Tensor]):
    moes = []
    with torch.inference_mode(mode=True):
        for _ in range(layer_num):
            config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size, 0)
            config.max_len = max_len
            config.quant_config.bits = 4
            config.quant_config.group_size = k_group_size
            config.quant_config.zero_point = False

            config.gate_proj = quant_data["gate_qweight"].data_ptr()
            config.up_proj = quant_data["up_qweight"].data_ptr()
            config.down_proj = quant_data["down_qweight"].data_ptr()

            config.gate_scale = quant_data["gate_scales"].data_ptr()
            config.up_scale = quant_data["up_scales"].data_ptr()
            config.down_scale = quant_data["down_scales"].data_ptr()
            config.pool = CPUInfer.backend_

            moe = kt_kernel_ext.moe.AMXFP4_KGroup_MOE(config)
            CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
            CPUInfer.sync()
            moes.append(moe)
    return moes


def run_case(pattern: str, qlen: int) -> Dict[str, float]:
    print("\n" + "=" * 70)
    desc = WEIGHT_PATTERNS[pattern][0]
    path = "mat-vec" if qlen <= DISPATCH_THRESHOLD else "mat-mat"
    print(f"Running case: {pattern} -> {desc}  (qlen={qlen}, path={path})")
    print("=" * 70)

    quant_data = prepare_mxfp4_quantized_weights(pattern)
    moes = build_moes_from_quantized_data(quant_data)

    dequant_weights = quant_data["dequantized"]
    gate_bf16 = dequant_weights["gate_proj"]
    up_bf16 = dequant_weights["up_proj"]
    down_bf16 = dequant_weights["down_proj"]

    diffs = []
    with torch.inference_mode(mode=True):
        for i in range(validation_iter):
            torch.manual_seed(100 + i)
            bsz_tensor = torch.tensor([qlen], device="cpu")
            expert_ids = torch.stack(
                [torch.randperm(expert_num)[:num_experts_per_tok] for _ in range(qlen)]
            ).contiguous()
            weights = torch.randn((qlen, num_experts_per_tok), dtype=torch.float32).contiguous()
            input_tensor = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous() / 100
            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()

            moe = moes[i % layer_num]
            CPUInfer.submit(
                moe.forward_task(
                    bsz_tensor.data_ptr(),
                    num_experts_per_tok,
                    expert_ids.data_ptr(),
                    weights.data_ptr(),
                    input_tensor.data_ptr(),
                    output.data_ptr(),
                    False,
                )
            )
            CPUInfer.sync()

            # Torch reference: use dequantized weights
            input_tensor_bf16 = input_tensor.to(torch.bfloat16)
            t_output = moe_torch(input_tensor_bf16, expert_ids, weights, gate_bf16, up_bf16, down_bf16).to(
                torch.bfloat16
            )

            t_output = t_output.flatten()
            output = output.flatten()

            diff = torch.mean(torch.abs(output.float() - t_output.float())) / (
                torch.mean(torch.abs(t_output.float())) + 1e-12
            )
            diffs.append(diff.item())
            print(f"[{pattern}] Iteration {i}: relative L1 diff = {diff:.4f}")
            print(f"           output   {output[:debug_print_count]}")
            print(f"           t_output {t_output[:debug_print_count]}")

    mean_diff = float(sum(diffs) / len(diffs))
    max_diff = float(max(diffs))
    min_diff = float(min(diffs))
    return {"case": pattern, "description": desc, "mean": mean_diff, "max": max_diff, "min": min_diff}


def run_fp4_moe_test():
    summary_rows = []
    for qlen in QLEN_LIST:
        path = "mat-vec" if qlen <= DISPATCH_THRESHOLD else "mat-mat"
        print(f"\n##### qlen={qlen}  path={path} #####")
        for case_name in WEIGHT_PATTERNS.keys():
            results = run_case(case_name, qlen)
            results["qlen"] = qlen
            results["path"] = path
            summary_rows.append(results)

    print("\n=== Case vs. Relative Error Summary ===")
    print(f"{'Case':<20} {'qlen':>5} {'path':<8} {'Mean':>10} {'Max':>10} {'Min':>10}")
    for row in summary_rows:
        print(f"{row['case']:<20} {row['qlen']:>5} {row['path']:<8} "
              f"{row['mean']*100:9.2f}% {row['max']*100:9.2f}% {row['min']*100:9.2f}%")


if __name__ == "__main__":
    run_fp4_moe_test()