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()