chore: import upstream snapshot with attribution

kt-kernel/examples/.gitignore

@@ -0,0 +1 @@
+debug

kt-kernel/examples/bench_moe_amx_int8.py

@@ -0,0 +1,357 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""
+AMX INT8 MoE Benchmark Script
+
+Benchmarks performance of AMX-accelerated INT8 MOE operations with configurable parameters.
+Supports uniform workload distribution across experts and optional CUDA stream mode.
+
+Usage:
+    python bench_moe_amx_int8.py [options]
+
+Examples:
+    # Default parameters
+    python bench_moe_amx_int8.py
+
+    # Custom parameters
+    python bench_moe_amx_int8.py --layer_num 4 --expert_num 256 --workload 8 --use_cuda_stream
+
+    # Full configuration
+    python bench_moe_amx_int8.py --layer_num 2 --expert_num 128 --num_experts_per_tok 8 \
+        --workload 4 --hidden_size 7168 --intermediate_size 2048 \
+        --warmup_iter 100 --test_iter 1000 --use_cuda_stream
+"""
+
+import os
+import sys
+import time
+import argparse
+
+# Add build path for development
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+
+import torch
+
+try:
+    from kt_kernel import kt_kernel_ext
+
+    HAS_KT_KERNEL = True
+except ImportError as e:
+    HAS_KT_KERNEL = False
+    import_error = str(e)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="AMX INT8 MoE Benchmark", formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    # Model parameters
+    parser.add_argument("--layer_num", type=int, default=2, help="Number of MoE layers")
+    parser.add_argument("--expert_num", type=int, default=256, help="Number of experts per layer")
+    parser.add_argument(
+        "--num_experts_per_tok", type=int, default=8, help="Number of experts selected per token (top-k)"
+    )
+    parser.add_argument("--hidden_size", type=int, default=7168, help="Hidden dimension size")
+    parser.add_argument("--intermediate_size", type=int, default=2048, help="Intermediate dimension size")
+
+    # Workload parameters
+    parser.add_argument("--workload", type=int, default=1, help="Workload (qlen, number of tokens)")
+    parser.add_argument("--max_len", type=int, default=25600, help="Maximum sequence length for buffer allocation")
+
+    # Benchmark parameters
+    parser.add_argument("--warmup_iter", type=int, default=100, help="Number of warmup iterations")
+    parser.add_argument("--test_iter", type=int, default=1000, help="Number of test iterations")
+
+    # Execution mode
+    parser.add_argument("--use_cuda_stream", action="store_true", help="Use CUDA stream mode (submit_with_cuda_stream)")
+    parser.add_argument("--profile", action="store_true", help="Enable PyTorch profiler and export trace.json")
+    parser.add_argument("--profile_path", type=str, default="./trace.json", help="Path to save profile trace")
+
+    # Worker configuration
+    parser.add_argument("--cpuinfer_threads", type=int, default=60, help="Total CPU inference threads")
+    parser.add_argument("--numa_count", type=int, default=2, help="Number of NUMA nodes")
+    parser.add_argument(
+        "--num_gpu_experts", type=int, default=0, help="Number of experts to place on GPU (first N experts)"
+    )
+
+    return parser.parse_args()
+
+
+def generate_uniform_workload(expert_num, num_experts_per_tok, workload):
+    """
+    Generate expert_ids and weights with uniform workload distribution.
+
+    workload = qlen (number of tokens)
+    Each token selects num_experts_per_tok experts.
+    Total expert calls = workload * num_experts_per_tok
+    """
+    qlen = workload
+
+    # Randomly select num_experts_per_tok experts (uniform, no duplicates)
+    # All tokens will use the same expert combination
+    selected_experts = torch.randperm(expert_num)[:num_experts_per_tok].tolist()
+
+    # Create expert_ids: all tokens use the same expert combination
+    expert_ids = [selected_experts for _ in range(qlen)]
+
+    # Create on GPU then copy to CPU (faster)
+    expert_ids = torch.tensor(expert_ids, dtype=torch.long, device="cuda").to("cpu").contiguous()
+    print(f"Selected experts (all tokens use same): {selected_experts}")
+    print(f"Expert IDs shape: {expert_ids.shape}")
+
+    # Uniform weights (normalized) - create on GPU then copy
+    weights = torch.ones((qlen, num_experts_per_tok), dtype=torch.float32, device="cuda") / num_experts_per_tok
+    weights = weights.to("cpu").contiguous()
+
+    return expert_ids, weights, qlen
+
+
+def run_benchmark(args):
+    """Run the AMX INT8 MoE benchmark."""
+
+    print("=" * 60)
+    print("AMX INT8 MoE Benchmark")
+    print("=" * 60)
+    print(f"\nConfiguration:")
+    print(f"  Layers:              {args.layer_num}")
+    print(f"  Experts per layer:   {args.expert_num}")
+    print(f"  Experts per token:   {args.num_experts_per_tok}")
+    print(f"  Hidden size:         {args.hidden_size}")
+    print(f"  Intermediate size:   {args.intermediate_size}")
+    print(f"  Workload (qlen):     {args.workload}")
+    print(f"  Use CUDA stream:     {args.use_cuda_stream}")
+    print(f"  Warmup iterations:   {args.warmup_iter}")
+    print(f"  Test iterations:     {args.test_iter}")
+    print(f"  CPU threads:         {args.cpuinfer_threads}")
+    print(f"  NUMA nodes:          {args.numa_count}")
+
+    # Generate uniform workload
+    expert_ids, weights, qlen = generate_uniform_workload(args.expert_num, args.num_experts_per_tok, args.workload)
+    print(f"\nActual qlen:           {qlen}")
+    print(f"Total expert calls:    {qlen * args.num_experts_per_tok}")
+
+    with torch.inference_mode():
+        # Initialize CPUInfer
+        if args.numa_count > 1:
+            worker_config = kt_kernel_ext.WorkerPoolConfig()
+            worker_config.subpool_count = args.numa_count
+            worker_config.subpool_numa_map = list(range(args.numa_count))
+            threads_per_numa = args.cpuinfer_threads // args.numa_count
+            worker_config.subpool_thread_count = [threads_per_numa] * args.numa_count
+            cpu_infer = kt_kernel_ext.CPUInfer(worker_config)
+        else:
+            cpu_infer = kt_kernel_ext.CPUInfer(args.cpuinfer_threads)
+
+        # Physical to logical mapping (identity)
+        physical_to_logical_map = torch.arange(args.expert_num, dtype=torch.int64, device="cpu").contiguous()
+
+        # GPU experts mask - set first num_gpu_experts to True if specified
+        gpu_experts_mask = torch.zeros(args.expert_num, dtype=torch.bool, device="cpu")
+        if args.num_gpu_experts > 0:
+            num_gpu = min(args.num_gpu_experts, args.expert_num)
+            gpu_experts_mask[:num_gpu] = True
+            print(f"  GPU experts: {num_gpu} (experts 0-{num_gpu-1})")
+
+        # Initialize MoE layers
+        print("\nInitializing MoE layers...")
+        moes = []
+        for layer_idx in range(args.layer_num):
+            # Create random weights on GPU then copy to CPU (faster)
+            gate_proj = (
+                torch.randn(
+                    (args.expert_num, args.intermediate_size, args.hidden_size), dtype=torch.bfloat16, device="cuda"
+                )
+                .to("cpu")
+                .contiguous()
+            )
+            up_proj = (
+                torch.randn(
+                    (args.expert_num, args.intermediate_size, args.hidden_size), dtype=torch.bfloat16, device="cuda"
+                )
+                .to("cpu")
+                .contiguous()
+            )
+            down_proj = (
+                torch.randn(
+                    (args.expert_num, args.hidden_size, args.intermediate_size), dtype=torch.bfloat16, device="cuda"
+                )
+                .to("cpu")
+                .contiguous()
+            )
+
+            # Configure MoE
+            config = kt_kernel_ext.moe.MOEConfig(
+                args.expert_num,
+                args.num_experts_per_tok,
+                args.hidden_size,
+                args.intermediate_size,
+                gpu_experts_mask.data_ptr(),
+            )
+            config.max_len = args.max_len
+            config.gate_proj = gate_proj.data_ptr()
+            config.up_proj = up_proj.data_ptr()
+            config.down_proj = down_proj.data_ptr()
+            config.pool = cpu_infer.backend_
+
+            moe = kt_kernel_ext.moe.AMXInt8_MOE(config)
+            cpu_infer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+            cpu_infer.sync()
+
+            moes.append(moe)
+            print(f"  Layer {layer_idx} initialized")
+
+        # Prepare input/output tensors
+        input_tensor = torch.randn((qlen, args.hidden_size), dtype=torch.bfloat16, device="cpu").contiguous()
+        output_tensor = torch.zeros((qlen, args.hidden_size), dtype=torch.bfloat16, device="cpu").contiguous()
+        bsz_tensor = torch.tensor([qlen], dtype=torch.int32, device="cpu")
+
+        # CUDA stream setup (if enabled)
+        cuda_stream = None
+        if args.use_cuda_stream:
+            if not torch.cuda.is_available():
+                print("\nWarning: CUDA not available, falling back to non-stream mode")
+                args.use_cuda_stream = False
+            else:
+                cuda_stream = torch.cuda.current_stream().cuda_stream
+                print(f"\nUsing CUDA stream: {cuda_stream}")
+
+        # Warmup
+        print(f"\nWarmup ({args.warmup_iter} iterations)...")
+        for i in range(args.warmup_iter):
+            moe = moes[i % args.layer_num]
+            task = moe.forward_task(
+                bsz_tensor.data_ptr(),
+                args.num_experts_per_tok,
+                expert_ids.data_ptr(),
+                weights.data_ptr(),
+                input_tensor.data_ptr(),
+                output_tensor.data_ptr(),
+                False,  # incremental
+            )
+
+            if args.use_cuda_stream:
+                cpu_infer.submit_with_cuda_stream(cuda_stream, task)
+                cpu_infer.sync_with_cuda_stream(cuda_stream)
+            else:
+                cpu_infer.submit(task)
+                cpu_infer.sync()
+
+        # Benchmark
+        print(f"Benchmarking ({args.test_iter} iterations)...")
+
+        if args.use_cuda_stream:
+            torch.cuda.synchronize()
+
+        # Setup profiler if enabled
+        profiler = None
+        if args.profile:
+            profiler = torch.profiler.profile(
+                activities=[
+                    torch.profiler.ProfilerActivity.CPU,
+                    torch.profiler.ProfilerActivity.CUDA,
+                ],
+                record_shapes=False,
+                with_stack=False,
+            )
+            profiler.__enter__()
+
+        start_time = time.perf_counter()
+
+        for i in range(args.test_iter):
+            moe = moes[i % args.layer_num]
+
+            if args.profile:
+                torch.cuda.nvtx.range_push(f"iter_{i}")
+
+            task = moe.forward_task(
+                bsz_tensor.data_ptr(),
+                args.num_experts_per_tok,
+                expert_ids.data_ptr(),
+                weights.data_ptr(),
+                input_tensor.data_ptr(),
+                output_tensor.data_ptr(),
+                False,
+            )
+
+            if args.use_cuda_stream:
+                if args.profile:
+                    torch.cuda.nvtx.range_push("submit")
+                cpu_infer.submit_with_cuda_stream(cuda_stream, task)
+                if args.profile:
+                    torch.cuda.nvtx.range_pop()
+                    torch.cuda.nvtx.range_push("sync")
+                cpu_infer.sync_with_cuda_stream(cuda_stream)
+                if args.profile:
+                    torch.cuda.nvtx.range_pop()
+            else:
+                cpu_infer.submit(task)
+                cpu_infer.sync()
+
+            if args.profile:
+                torch.cuda.nvtx.range_pop()
+
+        if args.use_cuda_stream:
+            torch.cuda.synchronize()
+
+        end_time = time.perf_counter()
+        total_time = end_time - start_time
+
+        # Export profiler trace
+        if profiler:
+            profiler.__exit__(None, None, None)
+            profiler.export_chrome_trace(args.profile_path)
+            print(f"\nProfile trace saved to: {args.profile_path}")
+
+        # Calculate metrics
+        # Note: each iteration processes ONE layer (round-robin: moe = moes[i % layer_num])
+        time_per_iter_us = total_time / args.test_iter * 1e6
+
+        # Bandwidth calculation
+        # Weight size per expert: 3 * hidden_size * intermediate_size * bytes_per_elem
+        bytes_per_elem = 1.0  # INT8
+        weight_bytes_per_expert = 3 * args.hidden_size * args.intermediate_size * bytes_per_elem
+
+        # Total weight bytes accessed per iteration (one layer per iteration)
+        # Each token activates num_experts_per_tok experts
+        total_experts_activated = qlen * args.num_experts_per_tok
+        weight_bytes_per_iter = total_experts_activated * weight_bytes_per_expert
+
+        bandwidth_gbs = weight_bytes_per_iter * args.test_iter / total_time / 1e9
+
+        # FLOPS calculation
+        # Per expert: 3 * hidden * intermediate * 2 (multiply-add)
+        flops_per_expert = 3 * args.hidden_size * args.intermediate_size * 2
+        total_flops = total_experts_activated * flops_per_expert * args.test_iter
+        tflops = total_flops / total_time / 1e12
+
+        # Results
+        print("\n" + "=" * 60)
+        print("Results")
+        print("=" * 60)
+        print(f"  Total time:           {total_time:.3f} s")
+        print(f"  Time per iteration:   {time_per_iter_us:.2f} us  (= time per layer)")
+        print(f"  Memory bandwidth:     {bandwidth_gbs:.2f} GB/s")
+        print(f"  Compute throughput:   {tflops:.3f} TFLOPS")
+        print("=" * 60)
+
+        return {
+            "total_time_s": total_time,
+            "time_per_iter_us": time_per_iter_us,
+            "bandwidth_gbs": bandwidth_gbs,
+            "tflops": tflops,
+        }
+
+
+def main():
+    args = parse_args()
+
+    if not HAS_KT_KERNEL:
+        print(f"Error: kt_kernel not available: {import_error}")
+        sys.exit(1)
+
+    run_benchmark(args)
+
+
+if __name__ == "__main__":
+    main()

kt-kernel/examples/configuration_deepseek_v3.py

@@ -0,0 +1,193 @@
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+DEEPSEEK_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+class DeepseekV3Config(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`DeepseekV3Model`]. It is used to instantiate an DeepSeek
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to that of the DeepSeek-V3.
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+    Args:
+        vocab_size (`int`, *optional*, defaults to 129280):
+            Vocabulary size of the Deep model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`DeepseekV3Model`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 11008):
+            Dimension of the MLP representations.
+        moe_intermediate_size (`int`, *optional*, defaults to 1407):
+            Dimension of the MoE representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer decoder.
+        num_nextn_predict_layers (`int`, *optional*, defaults to 1):
+            Number of nextn predict layers in the DeepSeekV3 Model.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        n_shared_experts (`int`, *optional*, defaults to None):
+            Number of shared experts, None means dense model.
+        n_routed_experts (`int`, *optional*, defaults to None):
+            Number of routed experts, None means dense model.
+        routed_scaling_factor (`float`, *optional*, defaults to 1.0):
+            Scaling factor or routed experts.
+        topk_method (`str`, *optional*, defaults to `gready`):
+            Topk method used in routed gate.
+        n_group (`int`, *optional*, defaults to None):
+            Number of groups for routed experts.
+        topk_group (`int`, *optional*, defaults to None):
+            Number of selected groups for each token(for each token, ensuring the selected experts is only within `topk_group` groups).
+        num_experts_per_tok (`int`, *optional*, defaults to None):
+            Number of selected experts, None means dense model.
+        moe_layer_freq (`int`, *optional*, defaults to 1):
+            The frequency of the MoE layer: one expert layer for every `moe_layer_freq - 1` dense layers.
+        first_k_dense_replace (`int`, *optional*, defaults to 0):
+            Number of dense layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head).
+                                                            \--k dense layers--/
+        norm_topk_prob (`bool`, *optional*, defaults to False):
+            Whether to normalize the weights of the routed experts.
+        scoring_func (`str`, *optional*, defaults to 'softmax'):
+            Method of computing expert weights.
+        aux_loss_alpha (`float`, *optional*, defaults to 0.001):
+            Auxiliary loss weight coefficient.
+        seq_aux = (`bool`, *optional*, defaults to True):
+            Whether to compute the auxiliary loss for each individual sample.
+        num_key_value_heads (`int`, *optional*):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+            `num_attention_heads`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        pad_token_id (`int`, *optional*):
+            Padding token id.
+        bos_token_id (`int`, *optional*, defaults to 1):
+            Beginning of stream token id.
+        eos_token_id (`int`, *optional*, defaults to 2):
+            End of stream token id.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether to tie weight embeddings
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
+            strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
+            `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
+            `max_position_embeddings` to the expected new maximum.
+        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+    ```python
+    >>> from transformers import DeepseekV3Model, DeepseekV3Config
+    >>> # Initializing a Deepseek-V3 style configuration
+    >>> configuration = DeepseekV3Config()
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "deepseek_v3"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=129280,
+        hidden_size=7168,
+        intermediate_size=18432,
+        moe_intermediate_size = 2048,
+        num_hidden_layers=61,
+        num_nextn_predict_layers=1,
+        num_attention_heads=128,
+        num_key_value_heads=128,
+        n_shared_experts = 1,
+        n_routed_experts = 256,
+        ep_size = 1,
+        routed_scaling_factor = 2.5,
+        kv_lora_rank = 512,
+        q_lora_rank = 1536,
+        qk_rope_head_dim = 64,
+        v_head_dim = 128,
+        qk_nope_head_dim = 128,
+        topk_method = 'noaux_tc',
+        n_group = 8,
+        topk_group = 4,
+        num_experts_per_tok = 8,
+        moe_layer_freq = 1,
+        first_k_dense_replace = 3,
+        norm_topk_prob = True,
+        scoring_func = 'sigmoid',
+        hidden_act="silu",
+        max_position_embeddings=4096,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=0,
+        eos_token_id=1,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.moe_intermediate_size = moe_intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_nextn_predict_layers = num_nextn_predict_layers
+        self.num_attention_heads = num_attention_heads
+        self.n_shared_experts = n_shared_experts
+        self.n_routed_experts = n_routed_experts
+        self.ep_size = ep_size
+        self.routed_scaling_factor = routed_scaling_factor
+        self.kv_lora_rank = kv_lora_rank
+        self.q_lora_rank = q_lora_rank
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.topk_method = topk_method
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.num_experts_per_tok = num_experts_per_tok
+        self.moe_layer_freq = moe_layer_freq
+        self.first_k_dense_replace = first_k_dense_replace
+        self.norm_topk_prob = norm_topk_prob
+        self.scoring_func = scoring_func
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

kt-kernel/examples/repro_llamafile_re.py

@@ -0,0 +1,168 @@
+#!/usr/bin/env python3
+"""
+Minimal LLAMAFILE repro harness to catch intermittent RuntimeError/RE.
+
+Requirements:
+- kt_kernel_ext built with LLAMAFILE (and CUDA stream integration)
+- Valid GGUF weights directory (WEIGHT_PATH)
+
+Usage:
+  WEIGHT_PATH=/path/to/gguf python examples/repro_llamafile_re.py
+
+Optional env:
+  DEVICE=cuda|cpu           # default: auto (cuda if available)
+  N_ITERS=1000              # iterations
+  BATCH=4                   # batch size
+  H=2048                    # hidden size
+  EXPERTS=128               # total experts
+  TOPK=8                    # experts per token
+  INTER=768                 # intermediate size (must be divisible by 256)
+  GPU_EXPERTS=100           # num experts on GPU side
+  TP=2                      # threadpool_count
+  CPU_THREADS=32            # cpuinfer_threads
+  MAX_DEFER=2               # max_deferred_experts_per_token
+  MODE=split|forward        # split=submit+sync, forward=wrapper.forward
+  SEED=1                    # random seed
+
+Debug tips:
+  - Set CUDA_LAUNCH_BLOCKING=1 to catch async errors deterministically.
+  - Try varying N_ITERS, BATCH, TOPK, MAX_DEFER.
+  - Capture stdout/stderr for failure iteration index.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import faulthandler
+import torch
+
+from kt_kernel import KTMoEWrapper
+
+
+def getenv_int(name: str, default: int) -> int:
+    try:
+        return int(os.environ.get(name, default))
+    except Exception:
+        return default
+
+
+def get_stream_for(device: torch.device | str):
+    device = torch.device(device)
+    if device.type == "cuda" and torch.cuda.is_available():
+        return torch.cuda.current_stream(device).cuda_stream
+    return 0
+
+
+def main() -> int:
+    faulthandler.enable()
+
+    weight_path = (os.environ.get("WEIGHT_PATH") or "").strip()
+    if not weight_path:
+        print("ERROR: WEIGHT_PATH env is required.")
+        return 2
+    if not os.path.exists(weight_path):
+        print(f"ERROR: WEIGHT_PATH does not exist: {weight_path}")
+        return 2
+
+    device_str = os.environ.get("DEVICE") or ("cuda" if torch.cuda.is_available() else "cpu")
+    device = torch.device(device_str)
+
+    n_iters = getenv_int("N_ITERS", 1000)
+    batch = getenv_int("BATCH", 4)
+    hidden = getenv_int("H", 2048)
+    experts = getenv_int("EXPERTS", 128)
+    topk = getenv_int("TOPK", 8)
+    inter = getenv_int("INTER", 768)
+    gpu_experts = getenv_int("GPU_EXPERTS", 100)
+    tp = getenv_int("TP", 2)
+    cpu_threads = getenv_int("CPU_THREADS", 32)
+    max_defer = getenv_int("MAX_DEFER", 2)
+    seed = getenv_int("SEED", 1)
+    mode = (os.environ.get("MODE") or "split").lower()
+
+    if inter % 256 != 0:
+        print(f"ERROR: INTER must be divisible by 256 for LLAMAFILE (got {inter}).")
+        return 2
+
+    print(
+        f"LLAMAFILE Repro: device={device}, iters={n_iters}, batch={batch}, H={hidden}, topk={topk}, E={experts}, inter={inter}, TP={tp}, CPU_THREADS={cpu_threads}, mode={mode}"
+    )
+    print(f"Weights: {weight_path}")
+
+    torch.manual_seed(seed)
+
+    # Create wrapper and load weights once
+    wrapper = KTMoEWrapper(
+        layer_idx=0,
+        num_experts=experts,
+        num_experts_per_tok=topk,
+        hidden_size=hidden,
+        moe_intermediate_size=inter,
+        num_gpu_experts=gpu_experts,
+        cpuinfer_threads=cpu_threads,
+        threadpool_count=tp,
+        weight_path=weight_path,
+        chunked_prefill_size=512,
+        method="LLAMAFILE",
+        max_deferred_experts_per_token=max_defer,
+    )
+    wrapper.load_weights()
+
+    # Optional capture of small batch sizes
+    KTMoEWrapper.set_capture_batch_sizes([1, 2, 4, 8, 16])
+
+    stream = get_stream_for(device)
+
+    # Allocate once and reuse to reduce allocator noise
+    hidden_states = torch.empty(batch, hidden, dtype=torch.bfloat16, device=device)
+    topk_ids = torch.empty(batch, topk, dtype=torch.long, device=device)
+    topk_weights = torch.empty(batch, topk, dtype=torch.float32, device=device)
+
+    def fill_random():
+        hidden_states.normal_(mean=0.0, std=1.0)
+        topk_ids.random_(0, experts)
+        topk_weights.uniform_()
+        topk_weights.div_(topk_weights.sum(dim=-1, keepdim=True) + 1e-6)
+
+    # Warmup
+    fill_random()
+    _ = wrapper.forward(hidden_states, topk_ids, topk_weights, stream)
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+
+    # Main loop
+    for i in range(n_iters):
+        try:
+            fill_random()
+            if mode == "forward":
+                _ = wrapper.forward(hidden_states, topk_ids, topk_weights, stream)
+            else:
+                wrapper.submit_forward(hidden_states, topk_ids, topk_weights, stream)
+                # Optional small GPU op to put work on the same stream
+                if device.type == "cuda":
+                    hidden_states.add_(0)  # no-op but enqueued on current stream
+                _ = wrapper.sync_forward(hidden_states, stream)
+
+            if (i + 1) % 50 == 0:
+                print(f"ok: iter {i + 1}/{n_iters}")
+                if device.type == "cuda":
+                    torch.cuda.synchronize(device)
+
+        except Exception as e:
+            print(f"FAIL at iter {i}: {repr(e)}")
+            # Flush GPU work for better diagnostics
+            if device.type == "cuda":
+                try:
+                    torch.cuda.synchronize(device)
+                except Exception as _:
+                    pass
+            return 1
+
+    print("All iterations completed without error.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
+

kt-kernel/examples/test-debug.py

@@ -0,0 +1,39 @@
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+import torch
+import ctypes
+from kt_kernel import kt_kernel_ext
+from kt_kernel_ext.moe import MOEConfig, MOE, AMXBF16_MOE, AMXInt8_MOE, AMXInt4_MOE, AMXInt4_1_MOE
+
+intermediate_size_full = 2048
+moe_intermediate_size = 3072
+hidden_size = 7168
+experts_num = 256
+num_experts_per_tok = 8
+cpu_infer = kt_kernel_ext.CPUInfer(97)
+
+up = torch.empty(experts_num, intermediate_size_full, hidden_size, dtype=torch.bfloat16, device="cpu")
+
+gate = torch.empty(experts_num, intermediate_size_full, hidden_size, dtype=torch.bfloat16, device="cpu")
+
+down = torch.empty(experts_num, hidden_size, intermediate_size_full, dtype=torch.bfloat16, device="cpu")
+
+gate_ptr = ctypes.addressof(ctypes.cast(gate.data_ptr(), ctypes.POINTER(ctypes.c_uint64)).contents)
+up_ptr = ctypes.addressof(ctypes.cast(up.data_ptr(), ctypes.POINTER(ctypes.c_uint64)).contents)
+down_ptr = ctypes.addressof(ctypes.cast(down.data_ptr(), ctypes.POINTER(ctypes.c_uint64)).contents)
+moe_config = MOEConfig(
+    experts_num,
+    num_experts_per_tok,
+    hidden_size,
+    moe_intermediate_size,
+)
+moe_config.layer_idx = 45
+moe_config.pool = cpu_infer.backend_
+moe_config.max_len = 1024  # TODO(zbx): multi cuda graph
+moe_config.gate_proj = gate_ptr
+moe_config.up_proj = up_ptr
+moe_config.down_proj = down_ptr
+moe_config.path = ""
+moe = AMXInt4_MOE(moe_config)

kt-kernel/examples/test_apply_rope.py

@@ -0,0 +1,52 @@
+import torch
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(q, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    b, h, s, d = q.shape
+    q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    return q_embed
+
+def my_apply(q,cos,sin):
+    
+    qa = q[:,:,range(0,64,2)]
+    qb = q[:,:,range(1,65,2)]
+    q1 = (qa * cos - qb * sin)
+    q2 = (qb*cos + qa*sin)
+    return torch.cat((q1,q2),-1)
+
+
+num_heads = 128
+seq_len = 1024
+rope_size = 64
+
+# theta = torch.randn(, dtype=torch.float32)
+
+

kt-kernel/examples/test_attention.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""
+Description  :
+Author       : Jianwei Dong
+Date         : 2024-08-28 10:32:05
+Version      : 1.0.0
+LastEditors  : chenht2022
+LastEditTime : 2024-08-28 10:32:05
+Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+"""
+import os, sys
+import time
+
+sys.path.append(os.path.dirname(__file__) + "/../build")
+from kt_kernel import kt_kernel_ext
+from flash_attn import flash_attn_with_kvcache
+import torch
+
+layer_num = 10
+kv_head_num = 8
+q_head_num = 32
+head_dim = 128
+block_len = 128
+anchor_num = 1
+cache_seqlen = 8192
+cache_seqlens = torch.tensor([cache_seqlen], dtype=torch.int32, device="cpu")
+seqlens_zero = torch.zeros((1,), dtype=torch.int32, device="cpu")
+anchor_type = kt_kernel_ext.kvcache.AnchorType.DYNAMIC
+kv_type = kt_kernel_ext.kvcache.ggml_type.FP16
+retrieval_type = kt_kernel_ext.kvcache.RetrievalType.LAYER
+layer_step: int = 1
+token_step: int = 1
+layer_offset: int = 0
+max_thread_num: int = 2
+max_batch_size: int = 1
+max_block_num: int = 512
+CPUInfer = kt_kernel_ext.CPUInfer(max_thread_num)
+validation_iter = 100
+
+with torch.inference_mode(mode=True):
+    config = kt_kernel_ext.kvcache.KVCacheConfig(
+        layer_num,
+        kv_head_num,
+        q_head_num,
+        head_dim,
+        block_len,
+        anchor_num,
+        anchor_type,
+        kv_type,
+        retrieval_type,
+        layer_step,
+        token_step,
+        layer_offset,
+        max_block_num,
+        max_batch_size,
+        max_thread_num,
+    )
+    local_kvcache = kt_kernel_ext.kvcache.KVCache(config)
+
+    kvcaches = []
+    block_table = torch.arange(max_block_num, dtype=torch.int32, device="cpu").contiguous().view(1, -1)
+
+    for layer_idx in range(layer_num):
+        k_cache = torch.randn((1, cache_seqlen, kv_head_num, head_dim), dtype=torch.float16, device="cpu").contiguous()
+        v_cache = torch.randn((1, cache_seqlen, kv_head_num, head_dim), dtype=torch.float16, device="cpu").contiguous()
+
+        CPUInfer.submit(
+            local_kvcache.update_kvcache_fp16(
+                k_cache.data_ptr(),
+                v_cache.data_ptr(),
+                layer_idx,
+                block_table.data_ptr(),
+                1,
+                max_block_num,
+                seqlens_zero.data_ptr(),
+                cache_seqlen,
+            )
+        )
+        CPUInfer.sync()
+
+        kvcaches.append((k_cache.to("cuda"), v_cache.to("cuda")))
+
+    # validation
+    for i in range(validation_iter):
+
+        k_cache = kvcaches[i % layer_num][0]
+        v_cache = kvcaches[i % layer_num][1]
+        input = torch.randn((1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu").contiguous()
+        output = torch.empty((1, 1, q_head_num, head_dim), dtype=torch.float16, device="cpu").contiguous()
+
+        # attn_lse: (bsz, q_len, q_head_num)
+        attn_lse = torch.empty((1, 1, q_head_num), dtype=torch.float32, device="cpu").contiguous()
+        input = input / 100
+
+        CPUInfer.submit(
+            local_kvcache.attn(
+                input.data_ptr(),
+                output.data_ptr(),
+                attn_lse.data_ptr(),
+                i % layer_num,
+                0,
+                1,
+                1,
+                max_block_num,
+                block_table.data_ptr(),
+                cache_seqlens.data_ptr(),
+                -1,
+                -1,
+                -1,
+            )
+        )
+        CPUInfer.sync()
+        # print("cpuinfer output", output)
+
+        t_output = flash_attn_with_kvcache(
+            q=input.to("cuda"),
+            k_cache=k_cache,
+            v_cache=v_cache,
+            cache_seqlens=cache_seqlens.to("cuda"),
+        )
+        # print("torch output", t_output)
+
+        diff = torch.mean(torch.abs(output.to("cuda") - t_output)) / torch.mean(torch.abs(t_output))
+        print("diff = ", diff)
+        assert diff < 0.001

kt-kernel/examples/test_awq_moe_amx.py

@@ -0,0 +1,622 @@
+import os, sys
+
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+
+from kt_kernel import kt_kernel_ext
+import torch
+
+# Set fixed seed for reproducible results
+torch.manual_seed(42)
+
+# Constants for 4-bit packing
+Q_BITS = 4
+STORAGE_BITS = 32
+PACK_NUM = STORAGE_BITS // Q_BITS  # 8
+
+
+def pack(imatrix: torch.Tensor, direction: str = "row"):
+    """
+    Packs a 4-bit integer matrix into a packed 32-bit integer matrix.
+    Packing order: 7 6 5 4 3 2 1 0 (MSB to LSB, original order)
+    Args:
+        imatrix (torch.Tensor): matrix of integers
+
+    Returns:
+        qmatrix (torch.Tensor): packed matrix of integers
+    """
+    shifts = torch.arange(0, STORAGE_BITS, Q_BITS, device=imatrix.device)
+
+    imatrix = imatrix.to(torch.int8)
+    imatrix = torch.bitwise_and(imatrix, 0x0F)  # eventually correct overflow
+
+    if direction == "column":
+        imatrix = imatrix.view(-1, imatrix.shape[1] // PACK_NUM, PACK_NUM)
+        qmatrix = torch.bitwise_left_shift(imatrix, shifts[None, None, :]).sum(dim=-1)
+
+    elif direction == "row":
+        imatrix = imatrix.view(imatrix.shape[0] // PACK_NUM, PACK_NUM, -1)
+        qmatrix = torch.bitwise_left_shift(imatrix, shifts[None, :, None]).sum(dim=1)
+
+    qmatrix = qmatrix.to(torch.int32)
+
+    return qmatrix
+
+
+expert_num = 16
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 25600
+num_experts_per_tok = 8
+qlen = 1
+layer_num = 1
+CPUInfer = kt_kernel_ext.CPUInfer(40)
+validation_iter = 10
+k_group_size = 64
+debug_print_count = 16
+
+physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous()
+
+
+def act_fn(x):
+    return x / (1.0 + torch.exp(-x))
+
+
+def generate_original_weights():
+    """Generate original FP16/BF16 weights for online quantization testing"""
+    # Set seed to ensure consistency between online and offline quantization
+    torch.manual_seed(42)
+
+    # Generate weights in the same format as test_moe_amx.py (bfloat16)
+    gate_proj_bf16 = (
+        torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cuda")
+        .to("cpu")
+        .contiguous()
+    )
+    up_proj_bf16 = (
+        torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cuda")
+        .to("cpu")
+        .contiguous()
+    )
+    down_proj_bf16 = (
+        torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.bfloat16, device="cuda")
+        .to("cpu")
+        .contiguous()
+    )
+
+    # Print first row of gate_proj for expert 0 (first debug_print_count elements)
+    print(
+        f"[DEBUG] Online quantization gate_proj expert 0, row 0, first {debug_print_count} elements: {gate_proj_bf16[0, 0, :debug_print_count]}"
+    )
+
+    return gate_proj_bf16, up_proj_bf16, down_proj_bf16
+
+
+def generate_awq_quantized_weights():
+    """Generate AWQ quantized weights (qweight, scales, qzeros) for testing"""
+    # Reset seed to ensure same weights as online quantization
+    torch.manual_seed(42)
+
+    # Generate original FP16 weights (convert from same random values as online version)
+    gate_proj_fp16 = (
+        torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cuda")
+        .to("cpu")
+        .to(torch.float16)
+        .contiguous()
+    )
+    up_proj_fp16 = (
+        torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cuda")
+        .to("cpu")
+        .to(torch.float16)
+        .contiguous()
+    )
+    down_proj_fp16 = (
+        torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.bfloat16, device="cuda")
+        .to("cpu")
+        .to(torch.float16)
+        .contiguous()
+    )
+
+    # Print first row of gate_proj for expert 0 (first debug_print_count elements)
+    print(
+        f"[DEBUG] Offline AWQ gate_proj expert 0, row 0, first {debug_print_count} elements: {gate_proj_fp16[0, 0, :debug_print_count]}"
+    )
+
+    # Calculate quantization parameters per group
+    def quantize_tensor_awq(weight, group_size=128):
+        """Simple AWQ-style quantization simulation with interleaving"""
+
+        w_orig_shape = weight.shape
+        expert_num, col, row = weight.shape
+        group_num = (row + group_size - 1) // group_size
+
+        # 1. reshape into groups along row dimension
+        weight_grouped = weight.view(expert_num, col, group_num, group_size)  # [E, G, group_size, C]
+
+        # 2. calculate scales per group (max abs value / 7.0 for 4-bit signed)
+        max_val = torch.max(weight_grouped, dim=3).values
+        min_val = torch.min(weight_grouped, dim=3).values
+        scales = (max_val - min_val).clamp(min=1e-5) / 15.0  # [E, G, C]
+        zeros = (-torch.round(min_val / scales)).clamp_(0, 15).to(torch.int8)
+
+        # 5. quantize weights
+        qweight_int = torch.clamp(
+            torch.round((weight_grouped - min_val.unsqueeze(-1)) / scales.unsqueeze(-1)), 0, 15
+        ).to(torch.int8)
+
+        qweight_int = qweight_int.view(w_orig_shape)
+
+        # 6. pack qweight along row (group_size) using helper
+        qweight_packed_list = []
+        for e in range(expert_num):
+            packed = pack(qweight_int[e], direction="column")  # [1, ? , col] or similar
+            qweight_packed_list.append(packed)
+        qweight_packed = torch.stack(qweight_packed_list, dim=0)  # [E, row, col / 8]
+
+        # 7. pack zeros along group dimension (row) using helper
+        zeros_packed_list = []
+        for e in range(expert_num):
+            zeros_packed_list.append(pack(zeros[e].transpose(0, 1), direction="column"))  # [blocks, col]
+        qzeros_packed = torch.stack(zeros_packed_list, dim=0)
+
+        scales = scales.transpose(1, 2).to(torch.float16)
+        print(scales.shape)
+        scales = scales.flatten().contiguous()
+
+        min_val = min_val.transpose(1, 2).to(torch.float16).flatten().contiguous()
+
+        zeros = zeros.transpose(1, 2).flatten().contiguous()
+
+        qzeros_packed = qzeros_packed.flatten().contiguous()
+
+        qweight_packed = qweight_packed.flatten().contiguous()
+
+        return {
+            "qweight": qweight_packed,  # Same for both torch and AWQ-MoE
+            "scales": scales,  # Same for both torch and AWQ-MoE
+            "qzeros": qzeros_packed,  # Same for both torch and AWQ-MoE
+            "mins": min_val,  # scales * zeros for comparison
+        }
+
+    # Quantize each projection
+    gate_data = quantize_tensor_awq(gate_proj_fp16, k_group_size)
+    up_data = quantize_tensor_awq(up_proj_fp16, k_group_size)
+    down_data = quantize_tensor_awq(down_proj_fp16, k_group_size)
+
+    return {
+        # Data for both torch and AWQ-MoE (no interleaving)
+        "gate_qweight": gate_data["qweight"],
+        "gate_scales": gate_data["scales"],
+        "gate_qzeros": gate_data["qzeros"],
+        "gate_mins": gate_data["mins"],
+        "up_qweight": up_data["qweight"],
+        "up_scales": up_data["scales"],
+        "up_qzeros": up_data["qzeros"],
+        "up_mins": up_data["mins"],
+        "down_qweight": down_data["qweight"],
+        "down_scales": down_data["scales"],
+        "down_qzeros": down_data["qzeros"],
+        "down_mins": down_data["mins"],
+        "original_fp16": {"gate_proj": gate_proj_fp16, "up_proj": up_proj_fp16, "down_proj": down_proj_fp16},
+    }
+
+
+def mlp_torch(input, gate_proj, up_proj, down_proj, debug_expert_id=None, debug_print=False):
+    gate_buf = torch.mm(input, gate_proj.t())
+    up_buf = torch.mm(input, up_proj.t())
+
+    if debug_print and debug_expert_id is not None:
+        print(f"[TORCH FP16 DEBUG] Expert {debug_expert_id}:")
+        print(f"  gate_buf[:{debug_print_count}] = {gate_buf.flatten()[:debug_print_count]}")
+        print(f"  up_buf[:{debug_print_count}] = {up_buf.flatten()[:debug_print_count]}")
+
+    intermediate = act_fn(gate_buf) * up_buf
+
+    if debug_print and debug_expert_id is not None:
+        print(f"  intermediate[:{debug_print_count}] = {intermediate.flatten()[:debug_print_count]}")
+
+    ret = torch.mm(intermediate, down_proj.t())
+
+    if debug_print and debug_expert_id is not None:
+        print(f"  down_output[:{debug_print_count}] = {ret.flatten()[:debug_print_count]}")
+
+    return ret
+
+
+def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj, debug_print=False):
+    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]]
+
+    # Get the first expert from expert_ids array to match AWQ-MoE behavior
+    target_debug_expert = expert_ids[0, 0].item()  # First expert in expert_ids array
+
+    outputs = []
+    start_idx = 0
+    activated_experts = []
+
+    for i, num_tokens in enumerate(tokens_per_expert):
+        end_idx = start_idx + num_tokens
+        if num_tokens == 0:
+            continue
+        activated_experts.append(i)
+        tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
+        # Only debug the target expert that matches AWQ-MoE's first expert
+        should_debug = debug_print and i == target_debug_expert
+        if gate_proj[i].dtype == torch.float16:
+            expert_out = mlp_torch(
+                tokens_for_this_expert.to(torch.float16),
+                gate_proj[i],
+                up_proj[i],
+                down_proj[i],
+                debug_expert_id=i,
+                debug_print=should_debug,
+            )
+        else:
+            expert_out = mlp_torch(
+                tokens_for_this_expert,
+                gate_proj[i],
+                up_proj[i],
+                down_proj[i],
+                debug_expert_id=i,
+                debug_print=should_debug,
+            )
+        outputs.append(expert_out)
+        start_idx = end_idx
+
+    if debug_print:
+        print(f"[TORCH DEBUG] Processing activated experts: {activated_experts}")
+        print(f"[TORCH DEBUG] Target debug expert (matches AWQ): {target_debug_expert}")
+
+    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)
+    )
+
+    if debug_print:
+        print(f"[TORCH DEBUG] Final MoE output[:{debug_print_count}] = {t_output.flatten()[:debug_print_count]}")
+
+    return t_output
+
+
+def test_online_int4_kgroup_moe():
+    """Test online Int4LowKGroup quantization (reference implementation)"""
+    print("Testing Online Int4LowKGroup quantization (reference)...")
+
+    # Generate original weights for online quantization
+    gate_proj, up_proj, down_proj = generate_original_weights()
+
+    with torch.inference_mode(mode=True):
+        moes = []
+        gate_projs = []
+        up_projs = []
+        down_projs = []
+
+        for _ in range(layer_num):
+            # Create Int4LowKGroup configuration (online quantization)
+            config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size, 0)
+            config.max_len = max_len
+            config.gate_proj = gate_proj.data_ptr()
+            config.up_proj = up_proj.data_ptr()
+            config.down_proj = down_proj.data_ptr()
+            config.gate_scale = 0
+            config.pool = CPUInfer.backend_
+
+            # Set quantization config for Int4LowKGroup (matches test_moe_amx.py)
+            config.quant_config.bits = 4
+            config.quant_config.group_size = k_group_size
+            config.quant_config.zero_point = True
+
+            # Enable weight dumping for comparison
+            config.save = True
+            config.path = "./awq_dump_online"
+
+            # Create Int4LowKGroup MoE (online quantization during load_weights)
+            moe = kt_kernel_ext.moe.AMXInt4_1KGroup_MOE(config)
+
+            # Load weights (performs online quantization)
+            print(f"Physical Map: {physical_to_logical_map.data_ptr()}")
+            CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+            CPUInfer.sync()
+
+            # Warm up
+            CPUInfer.submit(moe.warm_up_task())
+            CPUInfer.sync()
+
+            gate_projs.append(gate_proj)
+            up_projs.append(up_proj)
+            down_projs.append(down_proj)
+            moes.append(moe)
+
+        print("Online Int4LowKGroup MoE created and loaded successfully!")
+
+        # Run validation tests
+        results_online = []
+        for i in range(validation_iter):
+            # Reset seed for reproducible expert_ids and weights
+            torch.manual_seed(100 + i)  # Different seed to avoid same random values
+
+            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.rand((qlen, num_experts_per_tok), dtype=torch.float32).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            # input = torch.tensor(
+            #     data=torch.cat([torch.ones(qlen, 1), torch.zeros(qlen, hidden_size - 1)], dim=1),
+            #     dtype=torch.bfloat16
+            # )
+            input = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous() / 100
+
+            moe = moes[i % layer_num]
+
+            # Enable debug for first few iterations
+            enable_debug = i < 2
+            if enable_debug:
+                print(f"\n=== Online Int4LowKGroup Test Iteration {i} ===")
+                print(f"input[:{debug_print_count}] = {input.flatten()[:debug_print_count]}")
+                print(f"expert_ids = {expert_ids}")
+                print(f"weights = {weights}")
+
+            # Run online quantized MoE forward
+            CPUInfer.submit(
+                moe.forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids.data_ptr(),
+                    weights.data_ptr(),
+                    input.data_ptr(),
+                    output.data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+
+            if enable_debug:
+                print(f"[ONLINE DEBUG] AMX output[:{debug_print_count}] = {output.flatten()[:debug_print_count]}")
+
+            # Compare with FP16 reference
+            gate_proj_ref = gate_projs[i % layer_num]
+            up_proj_ref = up_projs[i % layer_num]
+            down_proj_ref = down_projs[i % layer_num]
+
+            t_output_online = moe_torch(
+                input, expert_ids, weights, gate_proj_ref, up_proj_ref, down_proj_ref, debug_print=enable_debug
+            )
+
+            # Calculate differences
+            diff_online = torch.mean(torch.abs(output - t_output_online)) / torch.mean(torch.abs(t_output_online))
+            results_online.append(output.clone())
+
+            print(f"Online Iteration {i}: Int4LowKGroup vs FP16 = {diff_online:.6f}")
+
+            if enable_debug:
+                abs_diff_online = torch.abs(output - t_output_online)
+                print(f"[COMPARE] Online Int4LowKGroup vs FP16:")
+                print(f"  Max abs diff = {torch.max(abs_diff_online):.6f}")
+                print(f"  Mean abs diff = {torch.mean(abs_diff_online):.6f}")
+                print(f"  Relative diff = {diff_online:.6f}")
+                print("=" * 70)
+
+        print("\n✅ Online Int4LowKGroup tests passed!")
+        return results_online
+
+
+def test_awq_moe():
+    print("Testing AWQ MoE with Int4_1LowKGroup quantization...")
+
+    # Generate AWQ quantized weights
+    awq_data = generate_awq_quantized_weights()
+
+    with torch.inference_mode(mode=True):
+        moes = []
+
+        for _ in range(layer_num):
+            # Create AWQ MoE configuration
+            config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size, 0)
+            config.max_len = max_len
+
+            # Set quantization config for Int4_1LowKGroup
+            config.quant_config.bits = 4
+            config.quant_config.group_size = k_group_size
+            config.quant_config.zero_point = True
+
+            # Enable weight dumping for comparison
+            config.save = True
+            config.path = "./awq_dump_offline"
+
+            # Set pointers to AWQ quantized data (no interleaving)
+            config.gate_proj = awq_data["gate_qweight"].data_ptr()
+            config.up_proj = awq_data["up_qweight"].data_ptr()
+            config.down_proj = awq_data["down_qweight"].data_ptr()
+
+            config.gate_scale = awq_data["gate_scales"].data_ptr()
+            config.up_scale = awq_data["up_scales"].data_ptr()
+            config.down_scale = awq_data["down_scales"].data_ptr()
+
+            config.gate_zeros = awq_data["gate_qzeros"].data_ptr()
+            config.up_zeros = awq_data["up_qzeros"].data_ptr()
+            config.down_zeros = awq_data["down_qzeros"].data_ptr()
+
+            config.pool = CPUInfer.backend_
+
+            # Create Int4_1LowKGroup MoE
+            moe = kt_kernel_ext.moe.AMXInt4_1KGroup_MOE(config)
+
+            # Load weights
+            CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+            CPUInfer.sync()
+
+            # Warm up
+            CPUInfer.submit(moe.warm_up_task())
+            CPUInfer.sync()
+
+            moes.append(moe)
+
+        print("AWQ MoE Int4_1LowKGroup created and loaded successfully!")
+
+        # Run validation tests
+        results_awq = []
+        for i in range(validation_iter):
+            # Reset seed for reproducible expert_ids and weights (same as online test)
+            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.rand((qlen, num_experts_per_tok), dtype=torch.float32).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            # input = torch.tensor(
+            #     data=torch.cat([torch.ones(qlen, 1), torch.zeros(qlen, hidden_size - 1)], dim=1),
+            #     dtype=torch.bfloat16
+            # )
+            input = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous() / 100
+
+            moe = moes[i % layer_num]
+
+            # Enable debug for first few iterations
+            enable_debug = i < 2
+            if enable_debug:
+                print(f"\n=== AWQ MoE Int4_1LowKGroup Test Iteration {i} ===")
+                print(f"input[:{debug_print_count}] = {input.flatten()[:debug_print_count]}")
+                print(f"expert_ids = {expert_ids}")
+                print(f"weights = {weights}")
+
+                # Print which experts will be activated
+                activated_experts = []
+                for token in range(expert_ids.shape[0]):
+                    for expert_idx in range(expert_ids.shape[1]):
+                        expert_id = expert_ids[token][expert_idx].item()
+                        if expert_id not in activated_experts:
+                            activated_experts.append(expert_id)
+                print(f"[TORCH DEBUG] Activated experts: {sorted(activated_experts)}")
+                print(f"[TORCH DEBUG] First expert from expert_ids array: {expert_ids[0, 0].item()}")
+
+            # Run AWQ MoE forward
+            CPUInfer.submit(
+                moe.forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids.data_ptr(),
+                    weights.data_ptr(),
+                    input.data_ptr(),
+                    output.data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+
+            if enable_debug:
+                print(f"[AWQ-MoE DEBUG] AMX output[:{debug_print_count}] = {output.flatten()[:debug_print_count]}")
+
+            # Compare with FP16 reference
+            original_weights = awq_data["original_fp16"]
+            gate_proj = original_weights["gate_proj"].to(torch.float16)
+            up_proj = original_weights["up_proj"].to(torch.float16)
+            down_proj = original_weights["down_proj"].to(torch.float16)
+
+            t_output_fp16 = moe_torch(
+                input, expert_ids, weights, gate_proj, up_proj, down_proj, debug_print=enable_debug
+            )
+
+            # Calculate differences
+            diff_fp16 = torch.mean(torch.abs(output - t_output_fp16)) / torch.mean(torch.abs(t_output_fp16))
+            results_awq.append(output.clone())
+
+            print(f"AWQ Iteration {i}: AWQ-MoE vs FP16 = {diff_fp16:.6f}")
+
+            if enable_debug:
+                abs_diff_fp16 = torch.abs(output - t_output_fp16)
+                print(f"[COMPARE] AWQ-MoE vs FP16:")
+                print(f"  Max abs diff = {torch.max(abs_diff_fp16):.6f}")
+                print(f"  Mean abs diff = {torch.mean(abs_diff_fp16):.6f}")
+                print(f"  Relative diff = {diff_fp16:.6f}")
+                print("=" * 70)
+
+            # AWQ quantization typically has higher error tolerance due to 4-bit quantization vs FP16
+            # assert(diff_fp16 < 0.5), f"AWQ-MoE vs FP16 error too large: {diff_fp16:.6f}"
+
+        print("\n✅ All AWQ MoE tests passed!")
+        return results_awq
+
+
+def compare_quantization_methods():
+    """Compare online and offline quantization methods"""
+    print("=" * 70)
+    print("Comparing Online vs Offline Quantization Methods")
+    print("=" * 70)
+
+    # Run online quantization test (reference)
+    print("\n" + "=" * 70)
+    print("PHASE 1: Online Int4LowKGroup Quantization (Reference)")
+    print("=" * 70)
+    results_online = test_online_int4_kgroup_moe()
+
+    # Run offline AWQ quantization test
+    print("\n" + "=" * 70)
+    print("PHASE 2: Offline AWQ Int4_1LowKGroup Quantization")
+    print("=" * 70)
+    results_awq = test_awq_moe()
+
+    # Compare the results
+    print("\n" + "=" * 70)
+    print("PHASE 3: Comparison Results")
+    print("=" * 70)
+
+    if len(results_online) != len(results_awq):
+        print(f"❌ Different number of results: Online={len(results_online)}, AWQ={len(results_awq)}")
+        return
+
+    print("Comparing Online Int4LowKGroup vs Offline AWQ results:")
+    total_diff = 0.0
+    max_diff = 0.0
+
+    for i in range(len(results_online)):
+        diff = torch.mean(torch.abs(results_online[i] - results_awq[i]))
+        rel_diff = diff / torch.mean(torch.abs(results_online[i]))
+        total_diff += rel_diff
+        max_diff = max(max_diff, diff.item())
+
+        if i < 3:  # Show detailed comparison for first 3 iterations
+            print(f"  Iteration {i}:")
+            print(f"    Absolute diff: {diff:.6f}")
+            print(f"    Relative diff: {rel_diff:.6f}")
+            print(f"    Online output[:{debug_print_count//2}]:  {results_online[i].flatten()[:debug_print_count//2]}")
+            print(f"    AWQ output[:{debug_print_count//2}]:     {results_awq[i].flatten()[:debug_print_count//2]}")
+        else:
+            print(f"  Iteration {i}: Relative diff = {rel_diff:.6f}")
+
+    avg_diff = total_diff / len(results_online)
+    print(f"\nOverall comparison:")
+    print(f"  Average relative difference: {avg_diff:.6f}")
+    print(f"  Maximum absolute difference: {max_diff:.6f}")
+
+    # Determine if results match within acceptable tolerance
+    tolerance = 0.01  # 1% tolerance
+    if avg_diff < tolerance:
+        print(f"✅ Results match within {tolerance:.1%} tolerance!")
+        print("   Your offline AWQ quantization implementation appears to be correct.")
+    else:
+        print(f"❌ Results differ by more than {tolerance:.1%} tolerance.")
+        print("   There may be differences between online and offline quantization.")
+
+
+if __name__ == "__main__":
+    print("=" * 70)
+    print("AWQ MoE AMX Test - Online vs Offline Quantization Comparison")
+    print("=" * 70)
+
+    compare_quantization_methods()
+
+    print("\n" + "=" * 70)
+    print("Test completed successfully!")
+    print("=" * 70)

kt-kernel/examples/test_bf16_moe.py

@@ -0,0 +1,245 @@
+"""
+Test script for AMX_BF16_MOE_TP (native BF16 MoE) kernel validation.
+
+This script:
+1. Generates random BF16 weights
+2. Runs the BF16 MoE kernel
+3. Compares results with PyTorch reference
+
+BF16 format notes:
+- Weight: BF16 stored as ggml_bf16_t, shape [expert_num, n, k]
+- No scales needed (native BF16 precision)
+"""
+
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+
+import torch
+from kt_kernel import kt_kernel_ext
+
+torch.manual_seed(42)
+
+# Model config
+hidden_size = 2048
+intermediate_size = 768
+max_len = 25600
+
+expert_num = 128
+num_experts_per_tok = 8
+
+qlen = 1
+layer_num = 5
+CPUInfer = kt_kernel_ext.CPUInfer(3)
+validation_iter = 5
+debug_print_count = 16
+
+physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous()
+
+
+def act_fn(x):
+    """SiLU activation function"""
+    return x / (1.0 + torch.exp(-x))
+
+
+def mlp_torch(input, gate_proj, up_proj, down_proj):
+    """Reference MLP computation in PyTorch"""
+    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):
+    """Reference MoE computation in PyTorch"""
+    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 build_bf16_weights():
+    """
+    Generate random BF16 weights.
+
+    Returns:
+        dict with BF16 weights for gate, up, down projections
+    """
+    torch.manual_seed(42)
+
+    # Generate random BF16 weights with small values
+    gate_proj = (
+        (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 100.0)
+        .to(torch.bfloat16)
+        .contiguous()
+    )
+    up_proj = (
+        (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 100.0)
+        .to(torch.bfloat16)
+        .contiguous()
+    )
+    down_proj = (
+        (torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32) / 100.0)
+        .to(torch.bfloat16)
+        .contiguous()
+    )
+
+    print(f"BF16 weights shape: gate={gate_proj.shape}, up={up_proj.shape}, down={down_proj.shape}")
+
+    # Debug: Print BF16 weight info for expert 0
+    print("\n=== DEBUG: BF16 Weight Info (Expert 0) ===")
+    print(f"gate_proj[0] first 8 values: {gate_proj[0, 0, :8]}")
+    print(f"gate_proj[0] stats: min={gate_proj[0].min()}, max={gate_proj[0].max()}")
+    print(f"up_proj[0] first 8 values: {up_proj[0, 0, :8]}")
+    print(f"down_proj[0] first 8 values: {down_proj[0, 0, :8]}")
+
+    return {
+        "gate_proj": gate_proj,
+        "up_proj": up_proj,
+        "down_proj": down_proj,
+    }
+
+
+def build_moes_from_bf16_data(bf16_data: dict):
+    """
+    Build BF16 MoE modules from BF16 weight data.
+    """
+    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
+
+            # Set BF16 weight pointers (no scales needed)
+            config.gate_proj = bf16_data["gate_proj"].data_ptr()
+            config.up_proj = bf16_data["up_proj"].data_ptr()
+            config.down_proj = bf16_data["down_proj"].data_ptr()
+
+            # No scales for BF16
+            config.gate_scale = 0
+            config.up_scale = 0
+            config.down_scale = 0
+            config.pool = CPUInfer.backend_
+
+            moe = kt_kernel_ext.moe.AMXBF16_MOE(config)
+            CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+            CPUInfer.sync()
+            moes.append(moe)
+    return moes
+
+
+def run_bf16_moe_test():
+    """
+    Run BF16 MoE validation test.
+    """
+    print("\n" + "=" * 70)
+    print("BF16 MoE Kernel Validation Test")
+    print("=" * 70)
+
+    # Build BF16 weights
+    print("\nGenerating BF16 weights...")
+    bf16_data = build_bf16_weights()
+
+    # Build MoE modules
+    print("\nBuilding BF16 MoE modules...")
+    moes = build_moes_from_bf16_data(bf16_data)
+
+    # Get weights for reference computation
+    gate_proj = bf16_data["gate_proj"]
+    up_proj = bf16_data["up_proj"]
+    down_proj = bf16_data["down_proj"]
+
+    diffs = []
+    with torch.inference_mode(mode=True):
+        for i in range(validation_iter):
+            torch.manual_seed(114514 + 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() / 10
+            input_tensor = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous() * 3
+            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()
+
+            assert not torch.isnan(output).any(), "NaN values detected in CPU expert output."
+            assert not torch.isinf(output).any(), "Inf values detected in CPU expert output."
+
+            # Reference computation using BF16 weights
+            t_output = moe_torch(input_tensor, expert_ids, weights, gate_proj, up_proj, down_proj)
+
+            t_output_flat = t_output.flatten()
+            output_flat = output.flatten()
+
+            diff = torch.mean(torch.abs(output_flat - t_output_flat)) / (torch.mean(torch.abs(t_output_flat)) + 1e-12)
+            diffs.append(diff.item())
+            print(f"Iteration {i}: relative L1 diff = {diff:.6f}")
+
+            if i < 3:  # Print detailed output for first few iterations
+                print(f"  kernel output: {output_flat[:debug_print_count]}")
+                print(f"  torch output:  {t_output_flat[:debug_print_count]}")
+
+    mean_diff = float(sum(diffs) / len(diffs))
+    max_diff = float(max(diffs))
+    min_diff = float(min(diffs))
+
+    print("\n" + "=" * 70)
+    print("BF16 MoE Test Results")
+    print("=" * 70)
+    print(f"Mean relative L1 diff: {mean_diff*100:.4f}%")
+    print(f"Max relative L1 diff:  {max_diff*100:.4f}%")
+    print(f"Min relative L1 diff:  {min_diff*100:.4f}%")
+
+    # Pass/Fail criteria (BF16 should be very accurate, <5% error)
+    threshold = 5.0
+    if mean_diff * 100 < threshold:
+        print(f"\nPASS: Mean error {mean_diff*100:.4f}% < {threshold}% threshold")
+    else:
+        print(f"\nFAIL: Mean error {mean_diff*100:.4f}% >= {threshold}% threshold")
+
+    return {"mean": mean_diff, "max": max_diff, "min": min_diff}
+
+
+if __name__ == "__main__":
+    run_bf16_moe_test()

kt-kernel/examples/test_deepseekv3.py

@@ -0,0 +1,469 @@
+import os, sys
+import time
+
+os.environ["BLAS_NUM_THREADS"] = "1"
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+from kt_kernel import kt_kernel_ext
+from kt_kernel_ext.kvcache import ggml_type
+import torch
+import logging
+import sys
+import json
+from pathlib import Path
+from transformers import (
+    AutoTokenizer,
+    AutoConfig,
+    AutoModelForCausalLM,
+    GenerationConfig,
+    TextStreamer,
+)
+
+logger = logging.getLogger("reader")
+
+from gguf.gguf_reader import GGUFReader
+
+# load_layers = 6
+load_layers = None
+CPUInfer = kt_kernel_ext.CPUInfer(304)
+max_qlen = 4096
+max_kvlen = 4096
+page_size = 256
+pages_count = 200
+
+
+def read_gguf_file(gguf_file_path):
+    """
+    Reads and prints key-value pairs and tensor information from a GGUF file in an improved format.
+
+    Parameters:
+    - gguf_file_path: Path to the GGUF file.
+    """
+
+    reader = GGUFReader(gguf_file_path)
+
+    # List all key-value pairs in a columnized format
+    # print("Key-Value Pairs:") # noqa: NP100
+    # max_key_length = max(len(key) for key in reader.fields.keys())
+    for key, field in reader.fields.items():
+        value = field.parts[field.data[0]]
+        # print(f"{key:{max_key_length}} : {value}") # noqa: NP100
+    # print("----") # noqa: NP100
+
+    # List all tensors
+    # print("Tensors:") # noqa: NP100
+    # tensor_info_format = "{:<30} | Shape: {:<15} | Size: {:<12} | Quantization: {}"
+    # print(tensor_info_format.format("Tensor Name", "Shape", "Size", "Quantization")) # noqa: NP100
+    # print("-" * 80) # noqa: NP100
+    re = []
+    for tensor in reader.tensors:
+        shape_str = "x".join(map(str, tensor.shape))
+        size_str = str(tensor.n_elements)
+        quantization_str = tensor.tensor_type.name
+        # print(tensor_info_format.format(tensor.name, shape_str, size_str, quantization_str)) # noqa: NP100
+        re.append(tensor)
+    return re
+
+
+def read_gguf_directory(directory):
+    """
+    Reads all GGUF files in a directory and prints their contents.
+
+    Parameters:
+    - directory: Path to the directory containing GGUF files.
+    """
+    if not os.path.isdir(directory):
+        logger.error(f"Directory {directory} does not exist.")
+        return
+
+    # List all GGUF files in the directory
+    files = [f for f in os.listdir(directory) if f.endswith(".gguf")]
+    if not files:
+        logger.info(f"No GGUF files found in {directory}.")
+        return
+
+    re = []
+    for file in files:
+        file_path = os.path.join(directory, file)
+        # print(f"Reading {file_path}:") # noqa: NP100
+        # print("\n") # noqa: NP100
+        re.extend(read_gguf_file(file_path))
+    re = {r.name: r for r in re}
+    return re
+
+
+def find_weights(name, weights):
+    """
+    Finds and returns the weights for a given name from the list of weights.
+
+    Parameters:
+    - name: The name of the weights to find.
+    - weights: List of weight tensors.
+
+    Returns:
+    - The weight tensor if found, otherwise None.
+    """
+    for weight in weights:
+        if weight.name == name:
+            return weight
+    raise ValueError(f"Weight with name {name} not found in the provided weights list.")
+
+
+def get_torch_tensor_from_gguf(gguf_weights, name):
+    return torch.from_numpy(gguf_weights[name].data).contiguous()
+
+
+def get_torch_tensor_and_type_from_gguf(gguf_weights, name):
+    return torch.from_numpy(gguf_weights[name].data).contiguous(), gguf_weights[name].tensor_type.name
+
+
+def type_to_ggml_type(type):
+    if type == "F32":
+        return ggml_type.FP32
+    elif type == "F16":
+        return ggml_type.FP16
+    elif type == "BF16":
+        return ggml_type.BF16
+    else:
+        raise ValueError(f"Unsupported data type: {type}")
+
+
+def build_mla(layer_idx, json_config, gguf_weights):
+    hidden_size = json_config["hidden_size"]
+    num_heads = json_config["num_attention_heads"]
+    q_lora_rank = json_config["q_lora_rank"]
+    kv_lora_rank = json_config["kv_lora_rank"]
+    nope_size = json_config["qk_nope_head_dim"]
+    rope_size = json_config["qk_rope_head_dim"]
+    max_position_embeddings = json_config["max_position_embeddings"]
+    rope_theta = json_config["rope_theta"]
+    rope_scaling = json_config["rope_scaling"]
+
+    config = kt_kernel_ext.mla.MLAConfig(
+        hidden_size,
+        q_lora_rank,
+        kv_lora_rank,
+        num_heads,
+        nope_size,
+        rope_size,
+    )
+    config.max_qlen = max_qlen
+    config.max_kvlen = max_kvlen
+    config.max_position_embeddings = max_position_embeddings
+    config.rope_scaling_factor = rope_scaling["factor"]
+    config.rope_theta = rope_theta
+    config.rope_scaling_beta_fast = rope_scaling["beta_fast"]
+    config.rope_scaling_beta_slow = rope_scaling["beta_slow"]
+    config.rope_scaling_mscale = rope_scaling["mscale"]
+    config.rope_scaling_mscale_all_dim = rope_scaling["mscale_all_dim"]
+    config.rope_scaling_original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
+
+    q_a_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_a.weight")
+    config.q_a_proj = q_a_proj_weight.data_ptr()
+    config.q_a_proj_type = type_to_ggml_type(type)
+    q_a_type = type
+
+    q_a_norm_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_a_norm.weight")
+    config.q_a_norm = q_a_norm_weight.data_ptr()
+    config.q_a_norm_type = type_to_ggml_type(type)
+
+    q_b_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_b.weight")
+    config.q_b_proj = q_b_proj_weight.data_ptr()
+    config.q_b_proj_type = type_to_ggml_type(type)
+
+    kv_a_proj_with_mqa_weight, type = get_torch_tensor_and_type_from_gguf(
+        gguf_weights, f"blk.{layer_idx}.attn_kv_a_mqa.weight"
+    )
+    config.kv_a_proj_with_mqa = kv_a_proj_with_mqa_weight.data_ptr()
+    config.kv_a_proj_with_mqa_type = type_to_ggml_type(type)
+
+    kv_a_norm_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_kv_a_norm.weight")
+    config.kv_a_norm = kv_a_norm_weight.data_ptr()
+    config.kv_a_norm_type = type_to_ggml_type(type)
+
+    kv_b_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_kv_b.weight")
+    config.kv_b_proj = kv_b_proj_weight.data_ptr()
+    config.kv_b_proj_type = type_to_ggml_type(type)
+
+    o_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_output.weight")
+    config.o_proj = o_proj_weight.data_ptr()
+    config.w_o_type = type_to_ggml_type(type)
+
+    config.layer_idx = layer_idx
+    config.pool = CPUInfer.backend_
+    config.page_count = pages_count
+
+    if q_a_type == "F32":
+        mla = kt_kernel_ext.mla.MLA_F32(config)
+    elif q_a_type == "F16":
+        mla = kt_kernel_ext.mla.MLA_F16(config)
+    elif q_a_type == "BF16":
+        # mla = kt_kernel_ext.mla.MLA_F32(config)
+        mla = kt_kernel_ext.mla.MLA_QUAN_F32(config)
+    else:
+        raise ValueError(f"Unsupported data type: {q_a_type}")
+
+    mla.load_weights()
+    mla.set_local_pages(pages_count)
+    return mla
+
+
+def build_ffn(layer_idx, json_config, gguf_weights):
+    if f"blk.{layer_idx}.ffn_gate.weight" in gguf_weights:  # dense
+        config = kt_kernel_ext.moe.MOEConfig(
+            json_config["num_experts_per_tok"] + json_config["n_shared_experts"],
+            json_config["num_experts_per_tok"] + json_config["n_shared_experts"],
+            json_config["hidden_size"],
+            json_config["moe_intermediate_size"],
+        )
+        config.layer_idx = layer_idx
+        config.max_len = max_qlen
+        config.pool = CPUInfer.backend_
+        gate, gate_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_gate.weight")
+        up, up_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_up.weight")
+        down, down_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_down.weight")
+
+        config.gate_proj = gate.data_ptr()
+        config.gate_type = type_to_ggml_type(gate_type)
+        config.up_proj = up.data_ptr()
+        config.up_type = type_to_ggml_type(up_type)
+        config.down_proj = down.data_ptr()
+        config.down_type = type_to_ggml_type(down_type)
+
+        moe = kt_kernel_ext.moe.KMLInt8_MOE(config)
+        moe.load_weights()
+        return moe
+
+    elif f"blk.{layer_idx}.ffn_gate_exps.weight" in gguf_weights:
+        config = kt_kernel_ext.moe.MOEConfig(
+            json_config["n_routed_experts"] + json_config["n_shared_experts"],
+            json_config["num_experts_per_tok"] + json_config["n_shared_experts"],
+            json_config["hidden_size"],
+            json_config["moe_intermediate_size"],
+        )
+        config.layer_idx = layer_idx
+        config.max_len = max_qlen
+        config.pool = CPUInfer.backend_
+        gate, gate_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_gate_exps.weight")
+        up, up_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_up_exps.weight")
+        down, down_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_down_exps.weight")
+
+        gate_sh, gate_sh_type = get_torch_tensor_and_type_from_gguf(
+            gguf_weights, f"blk.{layer_idx}.ffn_gate_shexp.weight"
+        )
+        up_sh, up_sh_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_up_shexp.weight")
+        down_sh, down_sh_type = get_torch_tensor_and_type_from_gguf(
+            gguf_weights, f"blk.{layer_idx}.ffn_down_shexp.weight"
+        )
+
+        gate_sh_expanded = gate_sh.unsqueeze(0)
+        gate = torch.cat([gate, gate_sh_expanded], dim=0).contiguous()
+        up_sh_expanded = up_sh.unsqueeze(0)
+        up = torch.cat([up, up_sh_expanded], dim=0).contiguous()
+        down_sh_expanded = down_sh.unsqueeze(0)
+        down = torch.cat([down, down_sh_expanded], dim=0).contiguous()
+
+        config.gate_proj = gate.data_ptr()
+        config.gate_type = type_to_ggml_type(gate_type)
+        config.up_proj = up.data_ptr()
+        config.up_type = type_to_ggml_type(up_type)
+        config.down_proj = down.data_ptr()
+        config.down_type = type_to_ggml_type(down_type)
+
+        moe = kt_kernel_ext.moe.KMLInt8_MOE(config)
+        moe.load_weights()
+        return moe
+
+    else:
+        raise ValueError(f"Unsupported FFN type for layer {layer_idx}")
+
+
+def build_moegate(layer_idx, json_config, gguf_weights):
+    config = kt_kernel_ext.gate.GateConfig(
+        json_config["hidden_size"],
+        json_config["num_experts_per_tok"],
+        json_config["n_routed_experts"],
+        json_config["n_group"],
+        json_config["topk_group"],
+    )
+
+    config.routed_scaling_factor = json_config["routed_scaling_factor"]
+
+    config.pool = CPUInfer.backend_
+
+    weight, weight_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_gate_inp.weight")
+    config.weight = weight.data_ptr()
+    config.weight_type = type_to_ggml_type(weight_type)
+
+    bias, bias_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.exp_probs_b.bias")
+    config.e_score_correction_bias = bias.data_ptr()
+    config.e_score_correction_bias_type = type_to_ggml_type(bias_type)
+
+    gate = kt_kernel_ext.gate.MoEGate(config)
+
+    return gate
+
+
+def build_llm(json_config, gguf_weights):
+
+    general_config = kt_kernel_ext.GeneralConfig()
+    general_config.vocab_size = json_config["vocab_size"]
+    general_config.hidden_size = json_config["hidden_size"]
+    general_config.num_experts_per_tok = json_config["num_experts_per_tok"]
+    general_config.n_routed_experts = json_config["n_routed_experts"]
+    general_config.n_shared_experts = json_config["n_shared_experts"]
+    general_config.max_qlen = max_qlen
+
+    lm_heads, lm_heads_type = get_torch_tensor_and_type_from_gguf(gguf_weights, "output.weight")
+    general_config.lm_heads_ptr = lm_heads.data_ptr()
+    general_config.lm_heads_type = type_to_ggml_type(lm_heads_type)
+
+    output_norm, output_norm_type = get_torch_tensor_and_type_from_gguf(gguf_weights, "output_norm.weight")
+    general_config.norm_weights_ptr = output_norm.data_ptr()
+    general_config.norm_weights_type = type_to_ggml_type(output_norm_type)
+
+    token_embd, token_embd_type = get_torch_tensor_and_type_from_gguf(weights, "token_embd.weight")
+    general_config.token_embd_ptr = token_embd.data_ptr()
+    general_config.token_embd_type = type_to_ggml_type(token_embd_type)
+
+    general_config.pool = CPUInfer.backend_
+
+    llm = kt_kernel_ext.DeepseekV3ForCausalLM(general_config)
+    model = kt_kernel_ext.DeepseekV3Model(general_config)
+    llm.model = model
+
+    decoder_layers = []
+    real_load_layers = json_config["num_hidden_layers"] if load_layers is None else load_layers
+
+    for i in range(real_load_layers):
+        layer = kt_kernel_ext.DeepseekV3DecoderLayer(general_config, i)
+        attn_norm, attn_norm_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{i}.attn_norm.weight")
+        ffn_norm, ffn_norm_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{i}.ffn_norm.weight")
+
+        layer.load_norm(
+            attn_norm.data_ptr(),
+            type_to_ggml_type(attn_norm_type),
+            ffn_norm.data_ptr(),
+            type_to_ggml_type(ffn_norm_type),
+        )
+        layer.self_attn = build_mla(i, json_config, gguf_weights)
+        if f"blk.{i}.ffn_gate_inp.weight" in gguf_weights:
+            layer.gate = build_moegate(i, json_config, gguf_weights)
+        layer.ffn = build_ffn(i, json_config, gguf_weights)
+        decoder_layers.append(layer)
+
+    model.layers = decoder_layers
+    return llm
+
+
+safetensor_path = "/home/bd/models/DeepSeek-R1"
+json_path = os.path.join(safetensor_path, "config.json")
+json_config = json.load(open(json_path, "r"))
+print(json_config)
+
+gguf_path = "/home/bd/models/DeepSeek-R1-BF16"
+weights = read_gguf_directory(gguf_path)
+weights = dict(sorted(weights.items()))
+
+
+for name, t in weights.items():
+    # if not name.startswith("blk"):
+    # if name.startswith("blk.10."):
+    # if "ffn_gate." in name:
+    # print(f"Found weight: {t.name}, Shape: {t.shape}, Type: {t.tensor_type.name}, Size: {t.n_elements}")
+    print(f"Found weight: {t.name}, Shape: {t.shape}, Type: {t.tensor_type.name}, Size: {t.n_elements}")
+
+print("Building LLM ...")
+load_start_time = time.perf_counter()
+llm = build_llm(json_config, weights)
+load_end_time = time.perf_counter()
+print(f"Load time: {load_end_time - load_start_time:.4f} seconds")
+
+print("Release Weight Tensors ...")
+weights = None
+print("Loading Configs ...")
+
+
+tokenizer = AutoTokenizer.from_pretrained(safetensor_path, trust_remote_code=True)
+config = AutoConfig.from_pretrained(safetensor_path, trust_remote_code=True)
+
+force_think = False
+
+
+output_logits = torch.zeros((max_qlen, json_config["vocab_size"]), dtype=torch.float32)
+
+
+def start_chat(content=None):
+    if content is None:
+        content = input("Chat: ")
+
+    messages = [{"role": "user", "content": content}]
+    input_tensor = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
+    if force_think:
+        token_thinks = torch.tensor(
+            [tokenizer.encode("<think>\\n", add_special_tokens=False)], device=input_tensor.device
+        )
+        input_tensor = torch.cat([input_tensor, token_thinks], dim=1)
+    input_tensor = input_tensor.squeeze(0)  # Add batch dimension
+
+    print(f"Input tensor: {input_tensor}, type {input_tensor.dtype}, shape {input_tensor.shape}")
+    kvlen = 0
+    step = 2
+    while True or step > 0:
+        step -= 1
+        stream = TextStreamer(tokenizer)
+
+        qlen = input_tensor.shape[0]
+        qlens = [qlen - kvlen]
+        kvlens = [kvlen]
+        page_tables = [list(range(pages_count))]
+        start_time = time.perf_counter()
+        llm.forward(qlens, page_tables, kvlens, input_tensor[kvlen:].data_ptr(), output_logits.data_ptr())
+        end_time = time.perf_counter()
+        print(
+            f"Forward time: {end_time - start_time:.4f} seconds, tps: {qlens[0] / (end_time - start_time)} tokens/sec"
+        )
+
+        logits = output_logits[0]
+        # print(logits)
+        # sample
+        next_token = torch.argmax(logits).item()
+        # print(f"Next token: {next_token}, {tokenizer.decode(next_token)}")
+        kvlen = input_tensor.shape[0]
+        input_tensor = torch.cat((input_tensor, torch.tensor([next_token])), dim=-1)
+
+        if next_token == tokenizer.eos_token_id or tokenizer.decode(next_token) == "<|im_end|>":
+            stream.end()
+            break
+        else:
+            stream.put(torch.tensor([next_token]))
+
+
+job_id = 0
+while True:
+    try:
+        # ---------- 让用户决定是否继续 ----------
+        choice = input("\n【回车】开始对话 | 输入 1 读取文件 | 输入 q/quit/exit 退出程序： ").strip().lower()
+        if choice in {"q", "quit", "exit"}:
+            print("收到退出指令，程序结束。")
+            break
+        elif choice == "1":
+            file_path = input("请输入要读取的文件路径：").strip()
+            if not Path(file_path).is_file():
+                print(f"文件 {file_path} 不存在，请检查路径。")
+                continue
+            with open(file_path, "r", encoding="utf-8") as file:
+                content = file.read()
+            print(f"读取到内容：\n{content}\n")
+            start_chat(content)
+        else:
+            start_chat()
+
+    except KeyboardInterrupt:
+        # 随时 Ctrl-C：放弃当前任务并重启
+        print(f"\n检测到 Ctrl-C，已终止对话 #{job_id}，马上重启…")
+    except Exception as e:
+        # 其他异常：打印错误信息并重启
+        print(f"\n发生错误：{e}\n已终止对话 #{job_id}，马上重启…")
+        logger.error(f"Error in job {job_id}: {e}", exc_info=True)
+    finally:
+        job_id += 1  # 不管中断与否，都给下一任务换编号

kt-kernel/examples/test_deepseekv3_prefill.py

@@ -0,0 +1,462 @@
+import os, sys
+import time
+
+os.environ["BLAS_NUM_THREADS"] = "1"
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+from kt_kernel import kt_kernel_ext
+from kt_kernel_ext.kvcache import ggml_type
+import torch
+import logging
+import sys
+import json
+from pathlib import Path
+from transformers import (
+    AutoTokenizer,
+    AutoConfig,
+    AutoModelForCausalLM,
+    GenerationConfig,
+    TextStreamer,
+)
+
+logger = logging.getLogger("reader")
+
+from gguf.gguf_reader import GGUFReader
+
+CPUInfer = kt_kernel_ext.CPUInfer(304)
+max_qlen = 4096
+max_kvlen = 4096
+page_size = 256
+pages_count = 200
+
+
+def read_gguf_file(gguf_file_path):
+    """
+    Reads and prints key-value pairs and tensor information from a GGUF file in an improved format.
+
+    Parameters:
+    - gguf_file_path: Path to the GGUF file.
+    """
+
+    reader = GGUFReader(gguf_file_path)
+
+    # List all key-value pairs in a columnized format
+    # print("Key-Value Pairs:") # noqa: NP100
+    # max_key_length = max(len(key) for key in reader.fields.keys())
+    for key, field in reader.fields.items():
+        value = field.parts[field.data[0]]
+        # print(f"{key:{max_key_length}} : {value}") # noqa: NP100
+    # print("----") # noqa: NP100
+
+    # List all tensors
+    # print("Tensors:") # noqa: NP100
+    # tensor_info_format = "{:<30} | Shape: {:<15} | Size: {:<12} | Quantization: {}"
+    # print(tensor_info_format.format("Tensor Name", "Shape", "Size", "Quantization")) # noqa: NP100
+    # print("-" * 80) # noqa: NP100
+    re = []
+    for tensor in reader.tensors:
+        shape_str = "x".join(map(str, tensor.shape))
+        size_str = str(tensor.n_elements)
+        quantization_str = tensor.tensor_type.name
+        # print(tensor_info_format.format(tensor.name, shape_str, size_str, quantization_str)) # noqa: NP100
+        re.append(tensor)
+    return re
+
+
+def read_gguf_directory(directory):
+    """
+    Reads all GGUF files in a directory and prints their contents.
+
+    Parameters:
+    - directory: Path to the directory containing GGUF files.
+    """
+    if not os.path.isdir(directory):
+        logger.error(f"Directory {directory} does not exist.")
+        return
+
+    # List all GGUF files in the directory
+    files = [f for f in os.listdir(directory) if f.endswith(".gguf")]
+    if not files:
+        logger.info(f"No GGUF files found in {directory}.")
+        return
+
+    re = []
+    for file in files:
+        file_path = os.path.join(directory, file)
+        # print(f"Reading {file_path}:") # noqa: NP100
+        # print("\n") # noqa: NP100
+        re.extend(read_gguf_file(file_path))
+    re = {r.name: r for r in re}
+    return re
+
+
+def find_weights(name, weights):
+    """
+    Finds and returns the weights for a given name from the list of weights.
+
+    Parameters:
+    - name: The name of the weights to find.
+    - weights: List of weight tensors.
+
+    Returns:
+    - The weight tensor if found, otherwise None.
+    """
+    for weight in weights:
+        if weight.name == name:
+            return weight
+    raise ValueError(f"Weight with name {name} not found in the provided weights list.")
+
+
+def get_torch_tensor_from_gguf(gguf_weights, name):
+    return torch.from_numpy(gguf_weights[name].data).contiguous()
+
+
+def get_torch_tensor_and_type_from_gguf(gguf_weights, name):
+    return torch.from_numpy(gguf_weights[name].data).contiguous(), gguf_weights[name].tensor_type.name
+
+
+def type_to_ggml_type(type):
+    if type == "F32":
+        return ggml_type.FP32
+    elif type == "F16":
+        return ggml_type.FP16
+    elif type == "BF16":
+        return ggml_type.BF16
+    else:
+        raise ValueError(f"Unsupported data type: {type}")
+
+
+def build_mla(layer_idx, json_config, gguf_weights):
+    hidden_size = json_config["hidden_size"]
+    num_heads = json_config["num_attention_heads"]
+    q_lora_rank = json_config["q_lora_rank"]
+    kv_lora_rank = json_config["kv_lora_rank"]
+    nope_size = json_config["qk_nope_head_dim"]
+    rope_size = json_config["qk_rope_head_dim"]
+    max_position_embeddings = json_config["max_position_embeddings"]
+    rope_theta = json_config["rope_theta"]
+    rope_scaling = json_config["rope_scaling"]
+
+    config = kt_kernel_ext.mla.MLAConfig(
+        hidden_size,
+        q_lora_rank,
+        kv_lora_rank,
+        num_heads,
+        nope_size,
+        rope_size,
+    )
+    config.max_qlen = max_qlen
+    config.max_kvlen = max_kvlen
+    config.max_position_embeddings = max_position_embeddings
+    config.rope_scaling_factor = rope_scaling["factor"]
+    config.rope_theta = rope_theta
+    config.rope_scaling_beta_fast = rope_scaling["beta_fast"]
+    config.rope_scaling_beta_slow = rope_scaling["beta_slow"]
+    config.rope_scaling_mscale = rope_scaling["mscale"]
+    config.rope_scaling_mscale_all_dim = rope_scaling["mscale_all_dim"]
+    config.rope_scaling_original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
+
+    q_a_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_a.weight")
+    config.q_a_proj = q_a_proj_weight.data_ptr()
+    config.q_a_proj_type = type_to_ggml_type(type)
+    q_a_type = type
+
+    q_a_norm_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_a_norm.weight")
+    config.q_a_norm = q_a_norm_weight.data_ptr()
+    config.q_a_norm_type = type_to_ggml_type(type)
+
+    q_b_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_b.weight")
+    config.q_b_proj = q_b_proj_weight.data_ptr()
+    config.q_b_proj_type = type_to_ggml_type(type)
+
+    kv_a_proj_with_mqa_weight, type = get_torch_tensor_and_type_from_gguf(
+        gguf_weights, f"blk.{layer_idx}.attn_kv_a_mqa.weight"
+    )
+    config.kv_a_proj_with_mqa = kv_a_proj_with_mqa_weight.data_ptr()
+    config.kv_a_proj_with_mqa_type = type_to_ggml_type(type)
+
+    kv_a_norm_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_kv_a_norm.weight")
+    config.kv_a_norm = kv_a_norm_weight.data_ptr()
+    config.kv_a_norm_type = type_to_ggml_type(type)
+
+    kv_b_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_kv_b.weight")
+    config.kv_b_proj = kv_b_proj_weight.data_ptr()
+    config.kv_b_proj_type = type_to_ggml_type(type)
+
+    o_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_output.weight")
+    config.o_proj = o_proj_weight.data_ptr()
+    config.w_o_type = type_to_ggml_type(type)
+
+    config.layer_idx = layer_idx
+    config.pool = CPUInfer.backend_
+    config.page_count = pages_count
+
+    if q_a_type == "F32":
+        mla = kt_kernel_ext.mla.MLA_F32(config)
+    elif q_a_type == "F16":
+        mla = kt_kernel_ext.mla.MLA_F16(config)
+    elif q_a_type == "BF16":
+        mla = kt_kernel_ext.mla.MLA_QUAN_F32(config)
+        # mla = kt_kernel_ext.mla.MLA_F32(config)
+    else:
+        raise ValueError(f"Unsupported data type: {q_a_type}")
+
+    mla.load_weights()
+    mla.set_local_pages(pages_count)
+    return mla
+
+
+def build_ffn(layer_idx, json_config, gguf_weights):
+    if f"blk.{layer_idx}.ffn_gate.weight" in gguf_weights:  # dense
+        config = kt_kernel_ext.moe.MOEConfig(
+            json_config["num_experts_per_tok"] + json_config["n_shared_experts"],
+            json_config["num_experts_per_tok"] + json_config["n_shared_experts"],
+            json_config["hidden_size"],
+            json_config["moe_intermediate_size"],
+        )
+        config.layer_idx = layer_idx
+        config.max_len = max_qlen
+        config.pool = CPUInfer.backend_
+        gate, gate_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_gate.weight")
+        up, up_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_up.weight")
+        down, down_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_down.weight")
+
+        config.gate_proj = gate.data_ptr()
+        config.gate_type = type_to_ggml_type(gate_type)
+        config.up_proj = up.data_ptr()
+        config.up_type = type_to_ggml_type(up_type)
+        config.down_proj = down.data_ptr()
+        config.down_type = type_to_ggml_type(down_type)
+
+        moe = kt_kernel_ext.moe.KMLInt8_MOE(config)
+        moe.load_weights()
+        return moe
+
+    elif f"blk.{layer_idx}.ffn_gate_exps.weight" in gguf_weights:
+        config = kt_kernel_ext.moe.MOEConfig(
+            json_config["n_routed_experts"] + json_config["n_shared_experts"],
+            json_config["num_experts_per_tok"] + json_config["n_shared_experts"],
+            json_config["hidden_size"],
+            json_config["moe_intermediate_size"],
+        )
+        config.layer_idx = layer_idx
+        config.max_len = max_qlen
+        config.pool = CPUInfer.backend_
+        gate, gate_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_gate_exps.weight")
+        up, up_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_up_exps.weight")
+        down, down_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_down_exps.weight")
+
+        gate_sh, gate_sh_type = get_torch_tensor_and_type_from_gguf(
+            gguf_weights, f"blk.{layer_idx}.ffn_gate_shexp.weight"
+        )
+        up_sh, up_sh_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_up_shexp.weight")
+        down_sh, down_sh_type = get_torch_tensor_and_type_from_gguf(
+            gguf_weights, f"blk.{layer_idx}.ffn_down_shexp.weight"
+        )
+
+        gate_sh_expanded = gate_sh.unsqueeze(0)
+        gate = torch.cat([gate, gate_sh_expanded], dim=0).contiguous()
+        up_sh_expanded = up_sh.unsqueeze(0)
+        up = torch.cat([up, up_sh_expanded], dim=0).contiguous()
+        down_sh_expanded = down_sh.unsqueeze(0)
+        down = torch.cat([down, down_sh_expanded], dim=0).contiguous()
+
+        config.gate_proj = gate.data_ptr()
+        config.gate_type = type_to_ggml_type(gate_type)
+        config.up_proj = up.data_ptr()
+        config.up_type = type_to_ggml_type(up_type)
+        config.down_proj = down.data_ptr()
+        config.down_type = type_to_ggml_type(down_type)
+
+        moe = kt_kernel_ext.moe.KMLInt8_MOE(config)
+        moe.load_weights()
+        return moe
+
+    else:
+        raise ValueError(f"Unsupported FFN type for layer {layer_idx}")
+
+
+def build_moegate(layer_idx, json_config, gguf_weights):
+    config = kt_kernel_ext.gate.GateConfig(
+        json_config["hidden_size"],
+        json_config["num_experts_per_tok"],
+        json_config["n_routed_experts"],
+        json_config["n_group"],
+        json_config["topk_group"],
+    )
+
+    config.routed_scaling_factor = json_config["routed_scaling_factor"]
+
+    config.pool = CPUInfer.backend_
+
+    weight, weight_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_gate_inp.weight")
+    config.weight = weight.data_ptr()
+    config.weight_type = type_to_ggml_type(weight_type)
+
+    bias, bias_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.exp_probs_b.bias")
+    config.e_score_correction_bias = bias.data_ptr()
+    config.e_score_correction_bias_type = type_to_ggml_type(bias_type)
+
+    gate = kt_kernel_ext.gate.MoEGate(config)
+
+    return gate
+
+
+def build_llm(json_config, gguf_weights):
+
+    general_config = kt_kernel_ext.GeneralConfig()
+    general_config.vocab_size = json_config["vocab_size"]
+    general_config.hidden_size = json_config["hidden_size"]
+    general_config.num_experts_per_tok = json_config["num_experts_per_tok"]
+    general_config.n_routed_experts = json_config["n_routed_experts"]
+    general_config.n_shared_experts = json_config["n_shared_experts"]
+    general_config.max_qlen = max_qlen
+
+    lm_heads, lm_heads_type = get_torch_tensor_and_type_from_gguf(gguf_weights, "output.weight")
+    general_config.lm_heads_ptr = lm_heads.data_ptr()
+    general_config.lm_heads_type = type_to_ggml_type(lm_heads_type)
+
+    output_norm, output_norm_type = get_torch_tensor_and_type_from_gguf(gguf_weights, "output_norm.weight")
+    general_config.norm_weights_ptr = output_norm.data_ptr()
+    general_config.norm_weights_type = type_to_ggml_type(output_norm_type)
+
+    token_embd, token_embd_type = get_torch_tensor_and_type_from_gguf(weights, "token_embd.weight")
+    general_config.token_embd_ptr = token_embd.data_ptr()
+    general_config.token_embd_type = type_to_ggml_type(token_embd_type)
+
+    general_config.pool = CPUInfer.backend_
+
+    llm = kt_kernel_ext.DeepseekV3ForCausalLM(general_config)
+    model = kt_kernel_ext.DeepseekV3Model(general_config)
+    llm.model = model
+
+    decoder_layers = []
+    for i in range(json_config["num_hidden_layers"]):
+        # for i in range(6):
+        # for i in [0,1,2,3,4,5,6,7,8,9,10]:
+        layer = kt_kernel_ext.DeepseekV3DecoderLayer(general_config, i)
+        attn_norm, attn_norm_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{i}.attn_norm.weight")
+        ffn_norm, ffn_norm_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{i}.ffn_norm.weight")
+
+        layer.load_norm(
+            attn_norm.data_ptr(),
+            type_to_ggml_type(attn_norm_type),
+            ffn_norm.data_ptr(),
+            type_to_ggml_type(ffn_norm_type),
+        )
+        layer.self_attn = build_mla(i, json_config, gguf_weights)
+        if f"blk.{i}.ffn_gate_inp.weight" in gguf_weights:
+            layer.gate = build_moegate(i, json_config, gguf_weights)
+        layer.ffn = build_ffn(i, json_config, gguf_weights)
+        decoder_layers.append(layer)
+
+    model.layers = decoder_layers
+    return llm
+
+
+safetensor_path = "/home/bd/models/DeepSeek-R1"
+json_path = os.path.join(safetensor_path, "config.json")
+json_config = json.load(open(json_path, "r"))
+print(json_config)
+
+gguf_path = "/home/bd/models/DeepSeek-R1-BF16"
+weights = read_gguf_directory(gguf_path)
+weights = dict(sorted(weights.items()))
+
+
+for name, t in weights.items():
+    # if not name.startswith("blk"):
+    # if name.startswith("blk.10."):
+    # if "ffn_gate." in name:
+    # print(f"Found weight: {t.name}, Shape: {t.shape}, Type: {t.tensor_type.name}, Size: {t.n_elements}")
+    print(f"Found weight: {t.name}, Shape: {t.shape}, Type: {t.tensor_type.name}, Size: {t.n_elements}")
+print("Building LLM ...")
+llm = build_llm(json_config, weights)
+print("Release Weight Tensors ...")
+weights = None
+print("Loading Configs ...")
+
+
+tokenizer = AutoTokenizer.from_pretrained(safetensor_path, trust_remote_code=True)
+config = AutoConfig.from_pretrained(safetensor_path, trust_remote_code=True)
+prompt_file = None
+force_think = False
+
+
+output_logits = torch.zeros((max_qlen, json_config["vocab_size"]), dtype=torch.float32)
+
+
+def start_chat():
+    while True:
+        content = input("Chat: ")
+        if content.startswith('"""'):  # prefix """
+            # multi lines input
+            content = content[3:] + "\n"
+            while True:
+                line = input("")
+                if line.endswith('"""'):
+                    # end multi lines input
+                    line = line[:-3]  # suffix """
+                    if line:
+                        content += line + "\n"
+                    break
+                else:
+                    content += line + "\n"
+
+        if content == "":
+            if prompt_file != None:
+                content = open(prompt_file, "r").read()
+            else:
+                content = "Please write a piece of quicksort code in C++."
+        elif os.path.isfile(content):
+            content = open(content, "r").read()
+
+        messages = [{"role": "user", "content": content}]
+        input_tensor = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
+        if force_think:
+            token_thinks = torch.tensor(
+                [tokenizer.encode("<think>\\n", add_special_tokens=False)], device=input_tensor.device
+            )
+            input_tensor = torch.cat([input_tensor, token_thinks], dim=1)
+        input_tensor = input_tensor.squeeze(0)  # Add batch dimension
+
+        print(f"Input tensor: {input_tensor}, type {input_tensor.dtype}, shape {input_tensor.shape}")
+        while True:
+            stream = TextStreamer(tokenizer)
+
+            qlen = input_tensor.shape[0]
+            qlens = [qlen]
+            kvlens = [0]
+            page_tables = [list(range(pages_count))]
+            llm.forward(qlens, page_tables, kvlens, input_tensor.data_ptr(), output_logits.data_ptr())
+
+            logits = output_logits[0]
+            # print(logits)
+            # sample
+            next_token = torch.argmax(logits).item()
+            # print(f"Next token: {next_token}, {tokenizer.decode(next_token)}")
+            input_tensor = torch.cat((input_tensor, torch.tensor([next_token])), dim=-1)
+
+            if next_token == tokenizer.eos_token_id or tokenizer.decode(next_token) == "<|im_end|>":
+                print(stream.end(), end="", flush=True)
+                break
+            else:
+                print(stream.put(torch.tensor([next_token])), end="", flush=True)
+
+
+job_id = 0
+while True:
+    try:
+        # ---------- 让用户决定是否继续 ----------
+        choice = input("\n【回车】开始对话 | 输入 q/quit/exit 退出程序： ").strip().lower()
+        if choice in {"q", "quit", "exit"}:
+            print("收到退出指令，程序结束。")
+            break
+
+        # ----------------------------------------
+
+        start_chat()  # 启动聊天会话
+    except KeyboardInterrupt:
+        # 随时 Ctrl-C：放弃当前任务并重启
+        print(f"\n检测到 Ctrl-C，已终止对话 #{job_id}，马上重启…")
+    finally:
+        job_id += 1  # 不管中断与否，都给下一任务换编号

kt-kernel/examples/test_deepseekv3_prefill_speed.py

@@ -0,0 +1,475 @@
+import os, sys
+import time
+
+os.environ["BLAS_NUM_THREADS"] = "1"
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+from kt_kernel import kt_kernel_ext
+from kt_kernel_ext.kvcache import ggml_type
+import torch
+import logging
+import sys
+import json
+from pathlib import Path
+from transformers import (
+    AutoTokenizer,
+    AutoConfig,
+    AutoModelForCausalLM,
+    GenerationConfig,
+    TextStreamer,
+)
+
+logger = logging.getLogger("reader")
+
+from gguf.gguf_reader import GGUFReader
+
+# load_layers = 3
+load_layers = None
+worker_config = kt_kernel_ext.WorkerPoolConfig()
+worker_config.subpool_count = 2
+worker_config.subpool_numa_map = [0, 1]
+worker_config.subpool_thread_count = [72, 72]
+CPUInfer = kt_kernel_ext.CPUInfer(worker_config)
+
+max_qlen = 4096
+max_kvlen = 4096
+page_size = 256
+pages_count = 200
+
+
+def read_gguf_file(gguf_file_path):
+    """
+    Reads and prints key-value pairs and tensor information from a GGUF file in an improved format.
+
+    Parameters:
+    - gguf_file_path: Path to the GGUF file.
+    """
+
+    reader = GGUFReader(gguf_file_path)
+
+    # List all key-value pairs in a columnized format
+    # print("Key-Value Pairs:") # noqa: NP100
+    # max_key_length = max(len(key) for key in reader.fields.keys())
+    for key, field in reader.fields.items():
+        value = field.parts[field.data[0]]
+        # print(f"{key:{max_key_length}} : {value}") # noqa: NP100
+    # print("----") # noqa: NP100
+
+    # List all tensors
+    # print("Tensors:") # noqa: NP100
+    # tensor_info_format = "{:<30} | Shape: {:<15} | Size: {:<12} | Quantization: {}"
+    # print(tensor_info_format.format("Tensor Name", "Shape", "Size", "Quantization")) # noqa: NP100
+    # print("-" * 80) # noqa: NP100
+    re = []
+    for tensor in reader.tensors:
+        shape_str = "x".join(map(str, tensor.shape))
+        size_str = str(tensor.n_elements)
+        quantization_str = tensor.tensor_type.name
+        # print(tensor_info_format.format(tensor.name, shape_str, size_str, quantization_str)) # noqa: NP100
+        re.append(tensor)
+    return re
+
+
+def read_gguf_directory(directory):
+    """
+    Reads all GGUF files in a directory and prints their contents.
+
+    Parameters:
+    - directory: Path to the directory containing GGUF files.
+    """
+    if not os.path.isdir(directory):
+        logger.error(f"Directory {directory} does not exist.")
+        return
+
+    # List all GGUF files in the directory
+    files = [f for f in os.listdir(directory) if f.endswith(".gguf")]
+    if not files:
+        logger.info(f"No GGUF files found in {directory}.")
+        return
+
+    re = []
+    for file in files:
+        file_path = os.path.join(directory, file)
+        # print(f"Reading {file_path}:") # noqa: NP100
+        # print("\n") # noqa: NP100
+        re.extend(read_gguf_file(file_path))
+    re = {r.name: r for r in re}
+    return re
+
+
+def find_weights(name, weights):
+    """
+    Finds and returns the weights for a given name from the list of weights.
+
+    Parameters:
+    - name: The name of the weights to find.
+    - weights: List of weight tensors.
+
+    Returns:
+    - The weight tensor if found, otherwise None.
+    """
+    for weight in weights:
+        if weight.name == name:
+            return weight
+    raise ValueError(f"Weight with name {name} not found in the provided weights list.")
+
+
+def get_torch_tensor_from_gguf(gguf_weights, name):
+    return torch.from_numpy(gguf_weights[name].data).contiguous()
+
+
+def get_torch_tensor_and_type_from_gguf(gguf_weights, name):
+    return torch.from_numpy(gguf_weights[name].data).contiguous(), gguf_weights[name].tensor_type.name
+
+
+def type_to_ggml_type(type):
+    if type == "F32":
+        return ggml_type.FP32
+    elif type == "F16":
+        return ggml_type.FP16
+    elif type == "BF16":
+        return ggml_type.BF16
+    else:
+        raise ValueError(f"Unsupported data type: {type}")
+
+
+def build_mla(layer_idx, json_config, gguf_weights):
+    hidden_size = json_config["hidden_size"]
+    num_heads = json_config["num_attention_heads"]
+    q_lora_rank = json_config["q_lora_rank"]
+    kv_lora_rank = json_config["kv_lora_rank"]
+    nope_size = json_config["qk_nope_head_dim"]
+    rope_size = json_config["qk_rope_head_dim"]
+    max_position_embeddings = json_config["max_position_embeddings"]
+    rope_theta = json_config["rope_theta"]
+    rope_scaling = json_config["rope_scaling"]
+
+    config = kt_kernel_ext.mla.MLAConfig(
+        hidden_size,
+        q_lora_rank,
+        kv_lora_rank,
+        num_heads,
+        nope_size,
+        rope_size,
+    )
+    config.max_qlen = max_qlen
+    config.max_kvlen = max_kvlen
+    config.max_position_embeddings = max_position_embeddings
+    config.rope_scaling_factor = rope_scaling["factor"]
+    config.rope_theta = rope_theta
+    config.rope_scaling_beta_fast = rope_scaling["beta_fast"]
+    config.rope_scaling_beta_slow = rope_scaling["beta_slow"]
+    config.rope_scaling_mscale = rope_scaling["mscale"]
+    config.rope_scaling_mscale_all_dim = rope_scaling["mscale_all_dim"]
+    config.rope_scaling_original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
+
+    q_a_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_a.weight")
+    config.q_a_proj = q_a_proj_weight.data_ptr()
+    config.q_a_proj_type = type_to_ggml_type(type)
+    q_a_type = type
+
+    q_a_norm_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_a_norm.weight")
+    config.q_a_norm = q_a_norm_weight.data_ptr()
+    config.q_a_norm_type = type_to_ggml_type(type)
+
+    q_b_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_b.weight")
+    config.q_b_proj = q_b_proj_weight.data_ptr()
+    config.q_b_proj_type = type_to_ggml_type(type)
+
+    kv_a_proj_with_mqa_weight, type = get_torch_tensor_and_type_from_gguf(
+        gguf_weights, f"blk.{layer_idx}.attn_kv_a_mqa.weight"
+    )
+    config.kv_a_proj_with_mqa = kv_a_proj_with_mqa_weight.data_ptr()
+    config.kv_a_proj_with_mqa_type = type_to_ggml_type(type)
+
+    kv_a_norm_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_kv_a_norm.weight")
+    config.kv_a_norm = kv_a_norm_weight.data_ptr()
+    config.kv_a_norm_type = type_to_ggml_type(type)
+
+    kv_b_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_kv_b.weight")
+    config.kv_b_proj = kv_b_proj_weight.data_ptr()
+    config.kv_b_proj_type = type_to_ggml_type(type)
+
+    o_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_output.weight")
+    config.o_proj = o_proj_weight.data_ptr()
+    config.w_o_type = type_to_ggml_type(type)
+
+    config.layer_idx = layer_idx
+    config.pool = CPUInfer.backend_
+    config.page_count = pages_count
+
+    if q_a_type == "F32":
+        mla = kt_kernel_ext.mla.MLA_F32(config)
+    elif q_a_type == "F16":
+        mla = kt_kernel_ext.mla.MLA_F16(config)
+    elif q_a_type == "BF16":
+        # mla = kt_kernel_ext.mla.MLA_F32(config)
+        mla = kt_kernel_ext.mla.MLA_QUAN_F32(config)
+    else:
+        raise ValueError(f"Unsupported data type: {q_a_type}")
+
+    mla.load_weights()
+    mla.set_local_pages(pages_count)
+    return mla
+
+
+def build_ffn(layer_idx, json_config, gguf_weights):
+    if f"blk.{layer_idx}.ffn_gate.weight" in gguf_weights:  # dense
+        config = kt_kernel_ext.moe.MOEConfig(
+            json_config["num_experts_per_tok"] + json_config["n_shared_experts"],
+            json_config["num_experts_per_tok"] + json_config["n_shared_experts"],
+            json_config["hidden_size"],
+            json_config["moe_intermediate_size"],
+        )
+        config.layer_idx = layer_idx
+        config.max_len = max_qlen
+        config.pool = CPUInfer.backend_
+        gate, gate_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_gate.weight")
+        up, up_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_up.weight")
+        down, down_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_down.weight")
+
+        config.gate_proj = gate.data_ptr()
+        config.gate_type = type_to_ggml_type(gate_type)
+        config.up_proj = up.data_ptr()
+        config.up_type = type_to_ggml_type(up_type)
+        config.down_proj = down.data_ptr()
+        config.down_type = type_to_ggml_type(down_type)
+
+        moe = kt_kernel_ext.moe.KMLInt8_MOE(config)
+        moe.load_weights()
+        return moe
+
+    elif f"blk.{layer_idx}.ffn_gate_exps.weight" in gguf_weights:
+        config = kt_kernel_ext.moe.MOEConfig(
+            json_config["n_routed_experts"] + json_config["n_shared_experts"],
+            json_config["num_experts_per_tok"] + json_config["n_shared_experts"],
+            json_config["hidden_size"],
+            json_config["moe_intermediate_size"],
+        )
+        config.layer_idx = layer_idx
+        config.max_len = max_qlen
+        config.pool = CPUInfer.backend_
+        gate, gate_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_gate_exps.weight")
+        up, up_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_up_exps.weight")
+        down, down_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_down_exps.weight")
+
+        gate_sh, gate_sh_type = get_torch_tensor_and_type_from_gguf(
+            gguf_weights, f"blk.{layer_idx}.ffn_gate_shexp.weight"
+        )
+        up_sh, up_sh_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_up_shexp.weight")
+        down_sh, down_sh_type = get_torch_tensor_and_type_from_gguf(
+            gguf_weights, f"blk.{layer_idx}.ffn_down_shexp.weight"
+        )
+
+        gate_sh_expanded = gate_sh.unsqueeze(0)
+        gate = torch.cat([gate, gate_sh_expanded], dim=0).contiguous()
+        up_sh_expanded = up_sh.unsqueeze(0)
+        up = torch.cat([up, up_sh_expanded], dim=0).contiguous()
+        down_sh_expanded = down_sh.unsqueeze(0)
+        down = torch.cat([down, down_sh_expanded], dim=0).contiguous()
+
+        config.gate_proj = gate.data_ptr()
+        config.gate_type = type_to_ggml_type(gate_type)
+        config.up_proj = up.data_ptr()
+        config.up_type = type_to_ggml_type(up_type)
+        config.down_proj = down.data_ptr()
+        config.down_type = type_to_ggml_type(down_type)
+
+        moe = kt_kernel_ext.moe.KMLInt8_MOE(config)
+        moe.load_weights()
+        return moe
+
+    else:
+        raise ValueError(f"Unsupported FFN type for layer {layer_idx}")
+
+
+def build_moegate(layer_idx, json_config, gguf_weights):
+    config = kt_kernel_ext.gate.GateConfig(
+        json_config["hidden_size"],
+        json_config["num_experts_per_tok"],
+        json_config["n_routed_experts"],
+        json_config["n_group"],
+        json_config["topk_group"],
+    )
+
+    config.routed_scaling_factor = json_config["routed_scaling_factor"]
+
+    config.pool = CPUInfer.backend_
+
+    weight, weight_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.ffn_gate_inp.weight")
+    config.weight = weight.data_ptr()
+    config.weight_type = type_to_ggml_type(weight_type)
+
+    bias, bias_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.exp_probs_b.bias")
+    config.e_score_correction_bias = bias.data_ptr()
+    config.e_score_correction_bias_type = type_to_ggml_type(bias_type)
+
+    gate = kt_kernel_ext.gate.MoEGate(config)
+
+    return gate
+
+
+def build_llm(json_config, gguf_weights):
+
+    general_config = kt_kernel_ext.GeneralConfig()
+    general_config.vocab_size = json_config["vocab_size"]
+    general_config.hidden_size = json_config["hidden_size"]
+    general_config.num_experts_per_tok = json_config["num_experts_per_tok"]
+    general_config.n_routed_experts = json_config["n_routed_experts"]
+    general_config.n_shared_experts = json_config["n_shared_experts"]
+    general_config.max_qlen = max_qlen
+
+    lm_heads, lm_heads_type = get_torch_tensor_and_type_from_gguf(gguf_weights, "output.weight")
+    general_config.lm_heads_ptr = lm_heads.data_ptr()
+    general_config.lm_heads_type = type_to_ggml_type(lm_heads_type)
+
+    output_norm, output_norm_type = get_torch_tensor_and_type_from_gguf(gguf_weights, "output_norm.weight")
+    general_config.norm_weights_ptr = output_norm.data_ptr()
+    general_config.norm_weights_type = type_to_ggml_type(output_norm_type)
+
+    token_embd, token_embd_type = get_torch_tensor_and_type_from_gguf(weights, "token_embd.weight")
+    general_config.token_embd_ptr = token_embd.data_ptr()
+    general_config.token_embd_type = type_to_ggml_type(token_embd_type)
+
+    general_config.pool = CPUInfer.backend_
+
+    llm = kt_kernel_ext.DeepseekV3ForCausalLM(general_config)
+    model = kt_kernel_ext.DeepseekV3Model(general_config)
+    llm.model = model
+
+    decoder_layers = []
+    real_load_layers = json_config["num_hidden_layers"] if load_layers is None else load_layers
+
+    for i in range(real_load_layers):
+        # for i in [2,3]:
+        layer = kt_kernel_ext.DeepseekV3DecoderLayer(general_config, i)
+        attn_norm, attn_norm_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{i}.attn_norm.weight")
+        ffn_norm, ffn_norm_type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{i}.ffn_norm.weight")
+
+        layer.load_norm(
+            attn_norm.data_ptr(),
+            type_to_ggml_type(attn_norm_type),
+            ffn_norm.data_ptr(),
+            type_to_ggml_type(ffn_norm_type),
+        )
+        layer.self_attn = build_mla(i, json_config, gguf_weights)
+        if f"blk.{i}.ffn_gate_inp.weight" in gguf_weights:
+            layer.gate = build_moegate(i, json_config, gguf_weights)
+        layer.ffn = build_ffn(i, json_config, gguf_weights)
+        decoder_layers.append(layer)
+
+    model.layers = decoder_layers
+    return llm
+
+
+safetensor_path = "/home/bd/models/DeepSeek-R1"
+json_path = os.path.join(safetensor_path, "config.json")
+json_config = json.load(open(json_path, "r"))
+print(json_config)
+
+gguf_path = "/home/bd/models/DeepSeek-R1-BF16"
+weights = read_gguf_directory(gguf_path)
+weights = dict(sorted(weights.items()))
+
+
+# for name, t in weights.items():
+# if not name.startswith("blk"):
+# if name.startswith("blk.10."):
+# if "ffn_gate." in name:
+# print(f"Found weight: {t.name}, Shape: {t.shape}, Type: {t.tensor_type.name}, Size: {t.n_elements}")
+# print(f"Found weight: {t.name}, Shape: {t.shape}, Type: {t.tensor_type.name}, Size: {t.n_elements}")
+
+print("Building LLM ...")
+load_start_time = time.perf_counter()
+llm = build_llm(json_config, weights)
+load_end_time = time.perf_counter()
+print(f"Load time: {load_end_time - load_start_time:.4f} seconds")
+
+print("Release Weight Tensors ...")
+weights = None
+print("Loading Configs ...")
+
+
+tokenizer = AutoTokenizer.from_pretrained(safetensor_path, trust_remote_code=True)
+config = AutoConfig.from_pretrained(safetensor_path, trust_remote_code=True)
+
+force_think = False
+
+
+output_logits = torch.zeros((max_qlen, json_config["vocab_size"]), dtype=torch.float32)
+
+
+def start_chat(content=None):
+    if content is None:
+        content = input("Chat: ")
+
+    messages = [{"role": "user", "content": content}]
+    input_tensor = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
+    if force_think:
+        token_thinks = torch.tensor(
+            [tokenizer.encode("<think>\\n", add_special_tokens=False)], device=input_tensor.device
+        )
+        input_tensor = torch.cat([input_tensor, token_thinks], dim=1)
+    input_tensor = input_tensor.squeeze(0)  # Add batch dimension
+
+    print(f"Input tensor: {input_tensor}, type {input_tensor.dtype}, shape {input_tensor.shape}")
+    kvlen = 0
+    step = 2
+    while True or step > 0:
+        step -= 1
+        stream = TextStreamer(tokenizer)
+
+        qlen = input_tensor.shape[0]
+        qlens = [qlen]
+        kvlens = [0]
+        page_tables = [list(range(pages_count))]
+        start_time = time.perf_counter()
+        llm.forward(qlens, page_tables, kvlens, input_tensor.data_ptr(), output_logits.data_ptr())
+        end_time = time.perf_counter()
+        print(
+            f"Forward time: {end_time - start_time:.4f} seconds, tps: {qlens[0] / (end_time - start_time)} tokens/sec"
+        )
+
+        logits = output_logits[0]
+        # print(logits)
+        # sample
+        next_token = torch.argmax(logits).item()
+        # print(f"Next token: {next_token}, {tokenizer.decode(next_token)}")
+        # kvlen = input_tensor.shape[0]
+        input_tensor = torch.cat((input_tensor, torch.tensor([next_token])), dim=-1)
+
+        if next_token == tokenizer.eos_token_id or tokenizer.decode(next_token) == "<|im_end|>":
+            stream.end()
+            break
+        else:
+            stream.put(torch.tensor([next_token]))
+
+
+job_id = 0
+while True:
+    try:
+        # ---------- 让用户决定是否继续 ----------
+        choice = input("\n【回车】开始对话 | 输入 1 读取文件 | 输入 q/quit/exit 退出程序： ").strip().lower()
+        if choice in {"q", "quit", "exit"}:
+            print("收到退出指令，程序结束。")
+            break
+        elif choice == "1":
+            file_path = input("请输入要读取的文件路径：").strip()
+            if not Path(file_path).is_file():
+                print(f"文件 {file_path} 不存在，请检查路径。")
+                continue
+            with open(file_path, "r", encoding="utf-8") as file:
+                content = file.read()
+            print(f"读取到内容：\n{content}\n")
+            start_chat(content)
+        else:
+            start_chat()
+
+    except KeyboardInterrupt:
+        # 随时 Ctrl-C：放弃当前任务并重启
+        print(f"\n检测到 Ctrl-C，已终止对话 #{job_id}，马上重启…")
+    except Exception as e:
+        # 其他异常：打印错误信息并重启
+        print(f"\n发生错误：{e}\n已终止对话 #{job_id}，马上重启…")
+        logger.error(f"Error in job {job_id}: {e}", exc_info=True)
+    finally:
+        job_id += 1  # 不管中断与否，都给下一任务换编号

kt-kernel/examples/test_fp4_moe_amx.py

@@ -0,0 +1,354 @@
+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()

kt-kernel/examples/test_fp4_moe_avx2.py

@@ -0,0 +1,236 @@
+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
+
+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 = 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)
+
+
+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
+    return torch.mm(intermediate, down_proj.t())
+
+
+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
+        expert_out = mlp_torch(sorted_tokens[start_idx:end_idx], 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
+    return (
+        new_x.view(*expert_ids.shape, -1)
+        .type(weights.dtype)
+        .mul_(weights.unsqueeze(dim=-1))
+        .sum(dim=1)
+        .type(new_x.dtype)
+    )
+
+
+def quantize_mxfp4_tensor(weights: torch.Tensor, group_size: int):
+    weights_f32 = weights.to(torch.float32)
+    e, rows, cols = weights_f32.shape
+    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)
+    normalized = reshaped / scales.unsqueeze(-1)
+    e2m1_vals = E2M1_VALUES.view(1, 1, 1, 1, 16)
+    normalized_expanded = normalized.unsqueeze(-1)
+    distances = torch.abs(normalized_expanded - e2m1_vals)
+    closest_indices = distances.argmin(dim=-1)
+    dequant = E2M1_VALUES[closest_indices].to(torch.float32) * scales.unsqueeze(-1)
+    dequant = dequant.view(e, rows, cols)
+    nibbles = closest_indices.to(torch.uint8)
+    nibbles = nibbles.view(e, rows, cols // 2, 2)
+    lo = nibbles[..., 0]
+    hi = nibbles[..., 1]
+    packed_bytes = (hi << 4) | lo
+    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
+
+
+WEIGHT_PATTERNS = {
+    "uniform_scale":     ("All k-groups share the same abs max / scale",   lambda g: torch.full((g,), 0.02, dtype=torch.float32)),
+    "alternating_scale": ("Alternate small / large abs max per k-group",   lambda g: torch.where(torch.arange(g) % 2 == 0, torch.full((g,), 0.015), torch.full((g,), 0.03))),
+    "ramp_scale":        ("Linearly increasing abs max per k-group",       lambda g: torch.linspace(0.005, 0.04, steps=g, dtype=torch.float32)),
+    "random":            ("Random bf16 weights (baseline)",                None),
+}
+
+
+def build_structured_tensor(shape, pattern):
+    if pattern == "random":
+        torch.manual_seed(42)
+        return (torch.randn(shape, dtype=torch.bfloat16) / 100.0).contiguous()
+    e, rows, cols = shape
+    groups = cols // k_group_size
+    group_vals = WEIGHT_PATTERNS[pattern][1](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_weights(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_s, gate_dq = quantize_mxfp4_tensor(gate_proj, k_group_size)
+    up_q,   up_s,   up_dq   = quantize_mxfp4_tensor(up_proj,   k_group_size)
+    down_q, down_s, down_dq = 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_s.contiguous(), "up_scales":  up_s.contiguous(), "down_scales":  down_s.contiguous(),
+        "dequantized":  {"gate_proj": gate_dq.to(torch.bfloat16), "up_proj": up_dq.to(torch.bfloat16), "down_proj": down_dq.to(torch.bfloat16)},
+    }
+
+
+def build_moes(quant_data):
+    AVX2MXFP4_MOE = getattr(kt_kernel_ext.moe, "AVX2MXFP4_MOE", None)
+    if AVX2MXFP4_MOE is None:
+        raise RuntimeError("AVX2MXFP4_MOE not found — rebuild kt-kernel with the AVX2 MXFP4 path")
+    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 = AVX2MXFP4_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, qlen):
+    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_weights(pattern)
+    moes = build_moes(quant_data)
+    dq = quant_data["dequantized"]
+
+    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()
+
+            t_output = moe_torch(input_tensor.to(torch.bfloat16), expert_ids, weights,
+                                 dq["gate_proj"], dq["up_proj"], dq["down_proj"]).to(torch.bfloat16)
+
+            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}] iter {i}: rel-L1 = {diff:.4f}")
+            print(f"  output   {output.flatten()[:debug_print_count]}")
+            print(f"  t_output {t_output.flatten()[:debug_print_count]}")
+
+    return {"case": pattern, "description": desc,
+            "mean": sum(diffs)/len(diffs) if diffs else 0.0,
+            "max": max(diffs) if diffs else 0.0,
+            "min": min(diffs) if diffs else 0.0}
+
+
+def run_fp4_moe_avx2_test():
+    summary = []
+    for qlen in QLEN_LIST:
+        path = "mat-vec" if qlen <= DISPATCH_THRESHOLD else "mat-mat"
+        print(f"\n##### qlen={qlen}  path={path} #####")
+        for pattern in WEIGHT_PATTERNS:
+            r = run_case(pattern, qlen)
+            r.update({"qlen": qlen, "path": path})
+            summary.append(r)
+
+    print("\n=== AVX2 MXFP4 — Relative Error Summary ===")
+    print(f"{'Case':<20} {'qlen':>5} {'path':<8} {'Mean':>10} {'Max':>10} {'Min':>10}")
+    for r in summary:
+        print(f"{r['case']:<20} {r['qlen']:>5} {r['path']:<8} "
+              f"{r['mean']*100:9.2f}% {r['max']*100:9.2f}% {r['min']*100:9.2f}%")
+
+
+if __name__ == "__main__":
+    run_fp4_moe_avx2_test()

kt-kernel/examples/test_fp4_moe_v4.py

@@ -0,0 +1,178 @@
+"""End-to-end MXFP4 MoE validation against the native DeepSeek-V4-Flash ckpt.
+
+Loads layer-`LAYER_ID` experts via :class:`MXFP4SafeTensorLoader`, runs the AMX
+FP4 backend, and compares against a torch reference that dequantizes the same
+nibble-packed weights with the OCP E2M1 LUT.
+
+Usage:
+    python test_fp4_moe_v4.py --weight-path /path/to/DeepSeek-V4-Flash [--layer 1]
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from typing import Tuple
+
+import torch
+
+# Allow running from kt-kernel/examples without install.
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/build")
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/python")
+
+from kt_kernel import kt_kernel_ext  # noqa: E402
+from kt_kernel.utils.loader import MXFP4SafeTensorLoader  # noqa: E402
+
+# OCP E2M1 codepoints in our LUT order (matches operators/amx/fp4-moe.hpp).
+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,
+)
+
+
+def dequantize_mxfp4(weight_u8: torch.Tensor, scale_bf16: torch.Tensor, group_size: int) -> torch.Tensor:
+    """Decode a [N, K/2] uint8 tensor of nibble-packed E2M1 with [N, K/gs] bf16
+    scales into a [N, K] bf16 weight tensor.
+
+    Layout (matches kernel's mxfp4_to_bf16_32): byte `b` low nibble = element K=2b,
+    high nibble = element K=2b+1.
+    """
+    n, k_packed = weight_u8.shape
+    k = k_packed * 2
+    assert k % group_size == 0, f"K={k} must be divisible by group_size={group_size}"
+    assert scale_bf16.shape == (n, k // group_size)
+
+    lo = (weight_u8 & 0x0F).to(torch.long)
+    hi = ((weight_u8 >> 4) & 0x0F).to(torch.long)
+    nibbles = torch.stack([lo, hi], dim=-1).view(n, k)  # interleave back to K order
+    decoded = E2M1_VALUES.to(weight_u8.device)[nibbles]  # [N, K] fp32
+
+    scale_fp32 = scale_bf16.to(torch.float32)
+    scale_full = scale_fp32.repeat_interleave(group_size, dim=-1)  # [N, K]
+    return (decoded * scale_full).to(torch.bfloat16).contiguous()
+
+
+def reference_mlp(x: torch.Tensor, gate: torch.Tensor, up: torch.Tensor, down: torch.Tensor) -> torch.Tensor:
+    g = torch.mm(x, gate.t())
+    u = torch.mm(x, up.t())
+    silu = g / (1.0 + torch.exp(-g.float())).to(g.dtype)
+    return torch.mm(silu * u, down.t())
+
+
+def reference_moe(
+    hidden: torch.Tensor,
+    expert_ids: torch.Tensor,
+    weights: torch.Tensor,
+    gate_w: torch.Tensor,  # [E, N, K]
+    up_w: torch.Tensor,
+    down_w: torch.Tensor,
+) -> torch.Tensor:
+    out = torch.zeros_like(hidden, dtype=torch.float32)
+    for tok in range(hidden.shape[0]):
+        for slot in range(expert_ids.shape[1]):
+            eid = int(expert_ids[tok, slot])
+            w = float(weights[tok, slot])
+            x = hidden[tok : tok + 1]
+            y = reference_mlp(x, gate_w[eid], up_w[eid], down_w[eid])
+            out[tok] += w * y[0].float()
+    return out.to(hidden.dtype)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser()
+    p.add_argument("--weight-path", required=True, help="Path to DeepSeek-V4-Flash safetensors directory.")
+    p.add_argument("--layer", type=int, default=1, help="Layer index to validate (default: 1).")
+    p.add_argument("--qlen", type=int, default=1, help="Number of tokens to test.")
+    p.add_argument("--top-k", type=int, default=6, help="num_experts_per_tok (V4 default 6).")
+    p.add_argument("--cpu-threads", type=int, default=32)
+    p.add_argument("--max-experts", type=int, default=0, help="Cap number of experts loaded (0 = all).")
+    return p.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    torch.manual_seed(0)
+
+    print(f"[V4-MXFP4] Loading layer {args.layer} from {args.weight_path}")
+    loader = MXFP4SafeTensorLoader(args.weight_path)
+    weights = loader.load_experts(f"model.layers.{args.layer}")
+
+    expert_num = len(weights["gate"])
+    if args.max_experts and args.max_experts < expert_num:
+        for k in ("gate", "up", "down", "gate_scale", "up_scale", "down_scale"):
+            weights[k] = weights[k][: args.max_experts]
+        expert_num = args.max_experts
+    print(f"[V4-MXFP4] expert_num={expert_num}")
+
+    gate0 = weights["gate"][0]
+    down0 = weights["down"][0]
+    intermediate_size = gate0.shape[0]
+    hidden_size = gate0.shape[1] * 2  # nibble-packed K
+    assert down0.shape == (hidden_size, intermediate_size // 2), f"unexpected down shape {down0.shape}"
+
+    group_size = hidden_size // weights["gate_scale"][0].shape[1]
+    print(f"[V4-MXFP4] hidden={hidden_size} inter={intermediate_size} gs={group_size}")
+    assert group_size == 32, "MXFP4 backend hard-codes group_size=32"
+
+    physical_to_logical = torch.arange(expert_num, dtype=torch.int64).contiguous()
+
+    # ----- AMX FP4 forward -----
+    cpu_infer = kt_kernel_ext.CPUInfer(args.cpu_threads)
+    cfg = kt_kernel_ext.moe.MOEConfig(expert_num, args.top_k, hidden_size, intermediate_size, 0)
+    cfg.layer_idx = args.layer
+    cfg.max_len = max(args.qlen, 1)
+    cfg.pool = cpu_infer.backend_
+    cfg.quant_config.bits = 4
+    cfg.quant_config.group_size = group_size
+    cfg.quant_config.zero_point = False
+
+    cfg.gate_projs = [[t.data_ptr() for t in weights["gate"]]]
+    cfg.up_projs = [[t.data_ptr() for t in weights["up"]]]
+    cfg.down_projs = [[t.data_ptr() for t in weights["down"]]]
+    cfg.gate_scales = [[t.data_ptr() for t in weights["gate_scale"]]]
+    cfg.up_scales = [[t.data_ptr() for t in weights["up_scale"]]]
+    cfg.down_scales = [[t.data_ptr() for t in weights["down_scale"]]]
+
+    moe = kt_kernel_ext.moe.AMXFP4_KGroup_MOE(cfg)
+    cpu_infer.submit(moe.load_weights_task(physical_to_logical.data_ptr()))
+    cpu_infer.sync()
+
+    qlen = args.qlen
+    top_k = args.top_k
+    bsz = torch.tensor([qlen], dtype=torch.int32)
+    expert_ids = torch.stack([torch.randperm(expert_num)[:top_k] for _ in range(qlen)]).to(torch.int32).contiguous()
+    routing = torch.randn((qlen, top_k), dtype=torch.float32).contiguous()
+    x = (torch.randn((qlen, hidden_size), dtype=torch.bfloat16) / 100).contiguous()
+    y_amx = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+
+    cpu_infer.submit(
+        moe.forward_task(
+            bsz.data_ptr(), top_k, expert_ids.data_ptr(), routing.data_ptr(),
+            x.data_ptr(), y_amx.data_ptr(), False,
+        )
+    )
+    cpu_infer.sync()
+
+    # ----- Torch reference (dequantize same nibbles + scales) -----
+    print("[V4-MXFP4] Building torch reference (dequantizing all loaded experts)…")
+    gate_bf16 = torch.stack([dequantize_mxfp4(weights["gate"][i], weights["gate_scale"][i], group_size) for i in range(expert_num)])
+    up_bf16 = torch.stack([dequantize_mxfp4(weights["up"][i], weights["up_scale"][i], group_size) for i in range(expert_num)])
+    down_bf16 = torch.stack([dequantize_mxfp4(weights["down"][i], weights["down_scale"][i], group_size) for i in range(expert_num)])
+
+    y_ref = reference_moe(x, expert_ids, routing, gate_bf16, up_bf16, down_bf16)
+
+    diff = (y_amx.float() - y_ref.float()).abs()
+    rel = diff.mean() / (y_ref.float().abs().mean() + 1e-12)
+    print(f"[V4-MXFP4] mean abs diff = {diff.mean().item():.4e}")
+    print(f"[V4-MXFP4] max  abs diff = {diff.max().item():.4e}")
+    print(f"[V4-MXFP4] rel mean diff = {rel.item()*100:.3f}%")
+    print(f"[V4-MXFP4] amx[:8]  = {y_amx.flatten()[:8]}")
+    print(f"[V4-MXFP4] ref[:8]  = {y_ref.flatten()[:8]}")
+
+    return 0 if rel.item() < 0.10 else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

kt-kernel/examples/test_fp8_moe.py

@@ -0,0 +1,458 @@
+"""
+Test script for GemmKernel224FP8 (FP8 MoE) kernel validation.
+
+This script:
+1. Generates random BF16 weights
+2. Quantizes them to FP8 format with 128x128 block-wise scales
+3. Runs the FP8 MoE kernel
+4. Compares results with PyTorch reference using dequantized BF16 weights
+
+FP8 format notes:
+- Weight: FP8 (E4M3) stored as uint8, shape [expert_num, n, k]
+- Scale: FP32, shape [expert_num, n // group_size, k // group_size], group_size=128
+"""
+
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+
+import torch
+import kt_kernel
+from kt_kernel import kt_kernel_ext
+
+torch.manual_seed(42)
+
+# Model config
+hidden_size = 3072
+intermediate_size = 1536
+max_len = 25600
+
+expert_num = 16
+num_experts_per_tok = 8
+
+qlen = 100
+layer_num = 1
+CPUInfer = kt_kernel_ext.CPUInfer(40)
+validation_iter = 1
+fp8_group_size = 128  # FP8 uses 128x128 block quantization
+debug_print_count = 16
+
+physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous()
+
+
+def act_fn(x):
+    """SiLU activation function"""
+    return x / (1.0 + torch.exp(-x))
+
+
+def mlp_torch(input, gate_proj, up_proj, down_proj):
+    """Reference MLP computation in PyTorch"""
+    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):
+    """Reference MoE computation in PyTorch"""
+    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
+
+
+# FP8 E4M3 constants
+FP8_E4M3_MAX = 448.0  # Maximum representable value in FP8 E4M3
+
+
+def fp8_e4m3_to_float(fp8_val: int) -> float:
+    """
+    Convert FP8 E4M3 value to float.
+    FP8 E4M3 format: 1 sign bit, 4 exponent bits, 3 mantissa bits
+    """
+    sign = (fp8_val >> 7) & 1
+    exp = (fp8_val >> 3) & 0xF
+    mant = fp8_val & 0x7
+
+    if exp == 0:
+        # Subnormal or zero
+        if mant == 0:
+            return -0.0 if sign else 0.0
+        # Subnormal: value = (-1)^sign * 2^(-6) * (0.mant)
+        return ((-1) ** sign) * (2**-6) * (mant / 8.0)
+    elif exp == 15:
+        # NaN (FP8 E4M3 doesn't have Inf, all exp=15 are NaN)
+        return float("nan")
+    else:
+        # Normal: value = (-1)^sign * 2^(exp-7) * (1.mant)
+        return ((-1) ** sign) * (2 ** (exp - 7)) * (1.0 + mant / 8.0)
+
+
+def float_to_fp8_e4m3(val: float) -> int:
+    """
+    Convert float to FP8 E4M3 value.
+    """
+    if val != val:  # NaN
+        return 0x7F  # NaN representation
+
+    sign = 1 if val < 0 else 0
+    val = abs(val)
+
+    if val == 0:
+        return sign << 7
+
+    # Clamp to max representable value
+    val = min(val, FP8_E4M3_MAX)
+
+    # Find exponent
+    import math
+
+    if val < 2**-9:  # Subnormal threshold
+        # Subnormal
+        mant = int(round(val / (2**-9)))
+        mant = min(mant, 7)
+        return (sign << 7) | mant
+
+    exp = int(math.floor(math.log2(val))) + 7
+    exp = max(1, min(exp, 14))  # Clamp exponent to valid range
+
+    # Calculate mantissa
+    mant = int(round((val / (2 ** (exp - 7)) - 1.0) * 8))
+    mant = max(0, min(mant, 7))
+
+    # Handle overflow to next exponent
+    if mant > 7:
+        mant = 0
+        exp += 1
+        if exp > 14:
+            exp = 14
+            mant = 7
+
+    return (sign << 7) | (exp << 3) | mant
+
+
+def quantize_to_fp8_blockwise(weights: torch.Tensor, group_size: int = 128):
+    """
+    Quantize BF16/FP32 weights to FP8 with block-wise scaling.
+
+    Args:
+        weights: [expert_num, n, k] tensor in BF16/FP32
+        group_size: Block size for quantization (default 128 for DeepSeek)
+
+    Returns:
+        fp8_weights: [expert_num, n, k] uint8 tensor
+        scales: [expert_num, n // group_size, k // group_size] BF16 tensor (scale_inv)
+    """
+    weights_f32 = weights.to(torch.float32)
+    e, n, k = weights_f32.shape
+
+    assert n % group_size == 0, f"n ({n}) must be divisible by group_size ({group_size})"
+    assert k % group_size == 0, f"k ({k}) must be divisible by group_size ({group_size})"
+
+    n_blocks = n // group_size
+    k_blocks = k // group_size
+
+    # Reshape to [e, n_blocks, group_size, k_blocks, group_size]
+    reshaped = weights_f32.view(e, n_blocks, group_size, k_blocks, group_size)
+    # Move to [e, n_blocks, k_blocks, group_size, group_size] for block processing
+    reshaped = reshaped.permute(0, 1, 3, 2, 4)
+
+    # Calculate max abs per block
+    max_abs = reshaped.abs().amax(dim=(-2, -1), keepdim=True)
+    max_abs = torch.clamp(max_abs, min=1e-12)
+
+    # Scale to FP8 range: scale = max_abs / FP8_MAX
+    # We store scale_inv = scale (for dequantization: fp8 * scale)
+    scales = (max_abs / FP8_E4M3_MAX).squeeze(-1).squeeze(-1)  # [e, n_blocks, k_blocks]
+
+    # Quantize: q = round(val / scale)
+    scaled = reshaped / (scales.unsqueeze(-1).unsqueeze(-1) + 1e-12)
+
+    # Convert to FP8 E4M3 using vectorized approach
+    # Clamp to FP8 representable range
+    scaled = scaled.clamp(-FP8_E4M3_MAX, FP8_E4M3_MAX)
+
+    # Simple quantization: round to nearest representable FP8 value
+    # For simplicity, we use a lookup table approach
+    fp8_q = torch.zeros_like(scaled, dtype=torch.uint8)
+
+    # Vectorized FP8 quantization
+    sign_mask = (scaled < 0).to(torch.uint8) << 7
+    abs_scaled = scaled.abs()
+
+    # Handle different ranges
+    # Subnormal: 0 < |x| < 2^-6
+    subnormal_mask = (abs_scaled > 0) & (abs_scaled < 2**-6)
+    subnormal_mant = (abs_scaled / (2**-9)).round().clamp(0, 7).to(torch.uint8)
+
+    # Normal values
+    normal_mask = abs_scaled >= 2**-6
+    log2_val = torch.log2(abs_scaled.clamp(min=2**-9))
+    exp = (log2_val.floor() + 7).clamp(1, 14).to(torch.int32)
+    mant = ((abs_scaled / (2.0 ** (exp.float() - 7)) - 1.0) * 8).round().clamp(0, 7).to(torch.uint8)
+
+    # Combine
+    fp8_q = torch.where(subnormal_mask, sign_mask | subnormal_mant, fp8_q)
+    fp8_q = torch.where(normal_mask, sign_mask | (exp.to(torch.uint8) << 3) | mant, fp8_q)
+
+    # Reshape back to [e, n, k]
+    fp8_q = fp8_q.permute(0, 1, 3, 2, 4).reshape(e, n, k)
+
+    # Scales shape: [e, n_blocks, k_blocks] -> store as [e, n_blocks, k_blocks]
+    scales_fp32 = scales.to(torch.float32).contiguous()
+
+    return fp8_q.contiguous(), scales_fp32
+
+
+def dequantize_fp8_blockwise(fp8_weights: torch.Tensor, scales: torch.Tensor, group_size: int = 128):
+    """
+    Dequantize FP8 weights back to BF16 for reference computation.
+
+    Args:
+        fp8_weights: [expert_num, n, k] uint8 tensor
+        scales: [expert_num, n // group_size, k // group_size] BF16 tensor
+        group_size: Block size
+
+    Returns:
+        dequantized: [expert_num, n, k] BF16 tensor
+    """
+    e, n, k = fp8_weights.shape
+    n_blocks = n // group_size
+    k_blocks = k // group_size
+
+    # Convert FP8 to float
+    # Build lookup table for FP8 E4M3 -> float
+    fp8_lut = torch.tensor([fp8_e4m3_to_float(i) for i in range(256)], dtype=torch.float32)
+
+    # Use lookup table
+    fp8_float = fp8_lut[fp8_weights.to(torch.int64)]
+
+    # Reshape for block-wise scaling
+    fp8_reshaped = fp8_float.view(e, n_blocks, group_size, k_blocks, group_size)
+    fp8_reshaped = fp8_reshaped.permute(0, 1, 3, 2, 4)  # [e, n_blocks, k_blocks, group_size, group_size]
+
+    # Apply scales
+    scales_f32 = scales.to(torch.float32).unsqueeze(-1).unsqueeze(-1)  # [e, n_blocks, k_blocks, 1, 1]
+    dequantized = fp8_reshaped * scales_f32
+
+    # Reshape back
+    dequantized = dequantized.permute(0, 1, 3, 2, 4).reshape(e, n, k)
+
+    return dequantized.to(torch.bfloat16).contiguous()
+
+
+def build_random_fp8_weights():
+    """
+    Generate random BF16 weights and quantize to FP8.
+
+    Returns:
+        dict with fp8 weights, scales, and original bf16 for reference
+    """
+    torch.manual_seed(42)
+
+    # Generate random BF16 weights with small values
+    gate_proj = (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 100.0).to(
+        torch.bfloat16
+    )
+    up_proj = (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 100.0).to(
+        torch.bfloat16
+    )
+    down_proj = (torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32) / 100.0).to(
+        torch.bfloat16
+    )
+
+    # Quantize to FP8
+    gate_fp8, gate_scales = quantize_to_fp8_blockwise(gate_proj, fp8_group_size)
+    up_fp8, up_scales = quantize_to_fp8_blockwise(up_proj, fp8_group_size)
+    down_fp8, down_scales = quantize_to_fp8_blockwise(down_proj, fp8_group_size)
+
+    # Dequantize for reference computation
+    gate_deq = dequantize_fp8_blockwise(gate_fp8, gate_scales, fp8_group_size)
+    up_deq = dequantize_fp8_blockwise(up_fp8, up_scales, fp8_group_size)
+    down_deq = dequantize_fp8_blockwise(down_fp8, down_scales, fp8_group_size)
+
+    print(f"FP8 weights shape: gate={gate_fp8.shape}, up={up_fp8.shape}, down={down_fp8.shape}")
+    print(f"Scales shape: gate={gate_scales.shape}, up={up_scales.shape}, down={down_scales.shape}")
+
+    # Debug: Print FP8 weight and scale info for expert 0
+    print("\n=== DEBUG: FP8 Weight and Scale Info (Expert 0) ===")
+    print(f"gate_fp8[0] first 8x8 block:")
+    for i in range(8):
+        print(f"  row {i}: {gate_fp8[0, i, :8].numpy().tobytes().hex(' ')}")
+    print(f"gate_fp8[0] stats: min={gate_fp8[0].min()}, max={gate_fp8[0].max()}")
+    print(f"gate_scales[0] first 4x4 block:\n{gate_scales[0, :4, :4]}")
+    print(f"gate_scales[0] stats: min={gate_scales[0].min()}, max={gate_scales[0].max()}")
+
+    print(f"\nup_fp8[0] first 8x8 block:")
+    for i in range(8):
+        print(f"  row {i}: {up_fp8[0, i, :8].numpy().tobytes().hex(' ')}")
+    print(f"up_fp8[0] stats: min={up_fp8[0].min()}, max={up_fp8[0].max()}")
+    print(f"up_scales[0] first 4x4 block:\n{up_scales[0, :4, :4]}")
+    print(f"up_scales[0] stats: min={up_scales[0].min()}, max={up_scales[0].max()}")
+
+    print(f"\ndown_fp8[0] first 8x8 block:")
+    for i in range(8):
+        print(f"  row {i}: {down_fp8[0, i, :8].numpy().tobytes().hex(' ')}")
+    print(f"down_fp8[0] stats: min={down_fp8[0].min()}, max={down_fp8[0].max()}")
+    print(f"down_scales[0] first 4x4 block:\n{down_scales[0, :4, :4]}")
+    print(f"down_scales[0] stats: min={down_scales[0].min()}, max={down_scales[0].max()}")
+
+    return {
+        "gate_fp8": gate_fp8.contiguous(),
+        "up_fp8": up_fp8.contiguous(),
+        "down_fp8": down_fp8.contiguous(),
+        "gate_scales": gate_scales.contiguous(),
+        "up_scales": up_scales.contiguous(),
+        "down_scales": down_scales.contiguous(),
+        "gate_deq": gate_deq.contiguous(),
+        "up_deq": up_deq.contiguous(),
+        "down_deq": down_deq.contiguous(),
+    }
+
+
+def build_moes_from_fp8_data(fp8_data: dict):
+    """
+    Build FP8 MoE modules from quantized data.
+    """
+    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 = 8
+            config.quant_config.group_size = fp8_group_size
+            config.quant_config.zero_point = False
+
+            # Set FP8 weight pointers
+            config.gate_proj = fp8_data["gate_fp8"].data_ptr()
+            config.up_proj = fp8_data["up_fp8"].data_ptr()
+            config.down_proj = fp8_data["down_fp8"].data_ptr()
+
+            # Set scale pointers
+            config.gate_scale = fp8_data["gate_scales"].data_ptr()
+            config.up_scale = fp8_data["up_scales"].data_ptr()
+            config.down_scale = fp8_data["down_scales"].data_ptr()
+            config.pool = CPUInfer.backend_
+
+            moe = kt_kernel_ext.moe.AMXFP8_MOE(config)
+            CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+            CPUInfer.sync()
+            moes.append(moe)
+    return moes
+
+
+def run_fp8_moe_test():
+    """
+    Run FP8 MoE validation test.
+    """
+    print("\n" + "=" * 70)
+    print("FP8 MoE Kernel Validation Test")
+    print("=" * 70)
+
+    # Build FP8 weights
+    print("\nGenerating and quantizing weights...")
+    fp8_data = build_random_fp8_weights()
+
+    # Build MoE modules
+    print("\nBuilding FP8 MoE modules...")
+    moes = build_moes_from_fp8_data(fp8_data)
+
+    # Get dequantized weights for reference
+    gate_deq = fp8_data["gate_deq"]
+    up_deq = fp8_data["up_deq"]
+    down_deq = fp8_data["down_deq"]
+
+    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() / 100
+            input_tensor = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous() * 1.5
+            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()
+
+            assert not torch.isnan(output).any(), "NaN values detected in CPU expert output."
+            assert not torch.isinf(output).any(), "Inf values detected in CPU expert output."
+
+            # Reference computation using dequantized weights
+            t_output = moe_torch(input_tensor, expert_ids, weights, gate_deq, up_deq, down_deq)
+
+            t_output_flat = t_output.flatten()
+            output_flat = output.flatten()
+
+            diff = torch.mean(torch.abs(output_flat - t_output_flat)) / (torch.mean(torch.abs(t_output_flat)) + 1e-12)
+            diffs.append(diff.item())
+            print(f"Iteration {i}: relative L1 diff = {diff:.6f}")
+
+            if i < 3:  # Print detailed output for first few iterations
+                print(f"  kernel output: {output_flat[:debug_print_count]}")
+                print(f"  torch output:  {t_output_flat[:debug_print_count]}")
+
+    mean_diff = float(sum(diffs) / len(diffs))
+    max_diff = float(max(diffs))
+    min_diff = float(min(diffs))
+
+    print("\n" + "=" * 70)
+    print("FP8 MoE Test Results")
+    print("=" * 70)
+    print(f"Mean relative L1 diff: {mean_diff*100:.4f}%")
+    print(f"Max relative L1 diff:  {max_diff*100:.4f}%")
+    print(f"Min relative L1 diff:  {min_diff*100:.4f}%")
+
+    # Pass/Fail criteria
+    threshold = 15.0  # 15% relative error threshold for FP8
+    if mean_diff * 100 < threshold:
+        print(f"\nPASS: Mean error {mean_diff*100:.4f}% < {threshold}% threshold")
+    else:
+        print(f"\nFAIL: Mean error {mean_diff*100:.4f}% >= {threshold}% threshold")
+
+    return {"mean": mean_diff, "max": max_diff, "min": min_diff}
+
+
+if __name__ == "__main__":
+    run_fp8_moe_test()

kt-kernel/examples/test_fp8_perchannel_moe.py

@@ -0,0 +1,408 @@
+"""
+Test script for FP8 Per-Channel MoE kernel validation (GLM-4.7-FP8 style).
+
+This script:
+1. Generates random BF16 weights
+2. Quantizes them to FP8 format with per-channel scales (one scale per output channel)
+3. Runs the FP8 Per-Channel MoE kernel
+4. Compares results with PyTorch reference using dequantized BF16 weights
+
+FP8 Per-Channel format notes:
+- Weight: FP8 (E4M3) stored as uint8, shape [expert_num, n, k]
+- Scale: FP32, shape [expert_num, n] (one scale per output row)
+"""
+
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+
+import torch
+from kt_kernel import kt_kernel_ext
+
+torch.manual_seed(42)
+
+# Model config
+hidden_size = 3072
+intermediate_size = 1536
+max_len = 25600
+
+expert_num = 16
+num_experts_per_tok = 8
+
+qlen = 100
+layer_num = 1
+CPUInfer = kt_kernel_ext.CPUInfer(40)
+validation_iter = 1
+debug_print_count = 16
+
+physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous()
+
+
+def act_fn(x):
+    """SiLU activation function"""
+    return x / (1.0 + torch.exp(-x))
+
+
+def mlp_torch(input, gate_proj, up_proj, down_proj):
+    """Reference MLP computation in PyTorch"""
+    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):
+    """Reference MoE computation in PyTorch"""
+    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
+
+
+# FP8 E4M3 constants
+FP8_E4M3_MAX = 448.0  # Maximum representable value in FP8 E4M3
+
+
+def fp8_e4m3_to_float(fp8_val: int) -> float:
+    """
+    Convert FP8 E4M3 value to float.
+    FP8 E4M3 format: 1 sign bit, 4 exponent bits, 3 mantissa bits
+    """
+    sign = (fp8_val >> 7) & 1
+    exp = (fp8_val >> 3) & 0xF
+    mant = fp8_val & 0x7
+
+    if exp == 0:
+        # Subnormal or zero
+        if mant == 0:
+            return -0.0 if sign else 0.0
+        # Subnormal: value = (-1)^sign * 2^(-6) * (0.mant)
+        return ((-1) ** sign) * (2**-6) * (mant / 8.0)
+    elif exp == 15:
+        # NaN (FP8 E4M3 doesn't have Inf, all exp=15 are NaN)
+        return float("nan")
+    else:
+        # Normal: value = (-1)^sign * 2^(exp-7) * (1.mant)
+        return ((-1) ** sign) * (2 ** (exp - 7)) * (1.0 + mant / 8.0)
+
+
+def float_to_fp8_e4m3(val: float) -> int:
+    """
+    Convert float to FP8 E4M3 value.
+    """
+    if val != val:  # NaN
+        return 0x7F  # NaN representation
+
+    sign = 1 if val < 0 else 0
+    val = abs(val)
+
+    if val == 0:
+        return sign << 7
+
+    # Clamp to max representable value
+    val = min(val, FP8_E4M3_MAX)
+
+    # Find exponent
+    import math
+
+    if val < 2**-9:  # Subnormal threshold
+        # Subnormal
+        mant = int(round(val / (2**-9)))
+        mant = min(mant, 7)
+        return (sign << 7) | mant
+
+    exp = int(math.floor(math.log2(val))) + 7
+    exp = max(1, min(exp, 14))  # Clamp exponent to valid range
+
+    # Calculate mantissa
+    mant = int(round((val / (2 ** (exp - 7)) - 1.0) * 8))
+    mant = max(0, min(mant, 7))
+
+    # Handle overflow to next exponent
+    if mant > 7:
+        mant = 0
+        exp += 1
+        if exp > 14:
+            exp = 14
+            mant = 7
+
+    return (sign << 7) | (exp << 3) | mant
+
+
+def quantize_to_fp8_perchannel(weights: torch.Tensor):
+    """
+    Quantize BF16/FP32 weights to FP8 with per-channel scaling.
+
+    Args:
+        weights: [expert_num, n, k] tensor in BF16/FP32
+
+    Returns:
+        fp8_weights: [expert_num, n, k] uint8 tensor
+        scales: [expert_num, n] FP32 tensor (one scale per output row)
+    """
+    weights_f32 = weights.to(torch.float32)
+    e, n, k = weights_f32.shape
+
+    # Calculate max abs per row (per output channel)
+    max_abs = weights_f32.abs().amax(dim=-1, keepdim=True)  # [e, n, 1]
+    max_abs = torch.clamp(max_abs, min=1e-12)
+
+    # Scale to FP8 range: scale = max_abs / FP8_MAX
+    scales = (max_abs / FP8_E4M3_MAX).squeeze(-1)  # [e, n]
+
+    # Quantize: q = round(val / scale)
+    scaled = weights_f32 / (scales.unsqueeze(-1) + 1e-12)
+
+    # Clamp to FP8 representable range
+    scaled = scaled.clamp(-FP8_E4M3_MAX, FP8_E4M3_MAX)
+
+    # Vectorized FP8 quantization
+    fp8_q = torch.zeros_like(scaled, dtype=torch.uint8)
+
+    sign_mask = (scaled < 0).to(torch.uint8) << 7
+    abs_scaled = scaled.abs()
+
+    # Handle different ranges
+    # Subnormal: 0 < |x| < 2^-6
+    subnormal_mask = (abs_scaled > 0) & (abs_scaled < 2**-6)
+    subnormal_mant = (abs_scaled / (2**-9)).round().clamp(0, 7).to(torch.uint8)
+
+    # Normal values
+    normal_mask = abs_scaled >= 2**-6
+    log2_val = torch.log2(abs_scaled.clamp(min=2**-9))
+    exp = (log2_val.floor() + 7).clamp(1, 14).to(torch.int32)
+    mant = ((abs_scaled / (2.0 ** (exp.float() - 7)) - 1.0) * 8).round().clamp(0, 7).to(torch.uint8)
+
+    # Combine
+    fp8_q = torch.where(subnormal_mask, sign_mask | subnormal_mant, fp8_q)
+    fp8_q = torch.where(normal_mask, sign_mask | (exp.to(torch.uint8) << 3) | mant, fp8_q)
+
+    return fp8_q.contiguous(), scales.to(torch.float32).contiguous()
+
+
+def dequantize_fp8_perchannel(fp8_weights: torch.Tensor, scales: torch.Tensor):
+    """
+    Dequantize FP8 weights back to BF16 for reference computation.
+
+    Args:
+        fp8_weights: [expert_num, n, k] uint8 tensor
+        scales: [expert_num, n] FP32 tensor
+
+    Returns:
+        dequantized: [expert_num, n, k] BF16 tensor
+    """
+    # Build lookup table for FP8 E4M3 -> float
+    fp8_lut = torch.tensor([fp8_e4m3_to_float(i) for i in range(256)], dtype=torch.float32)
+
+    # Use lookup table
+    fp8_float = fp8_lut[fp8_weights.to(torch.int64)]
+
+    # Apply per-channel scales
+    scales_expanded = scales.unsqueeze(-1)  # [e, n, 1]
+    dequantized = fp8_float * scales_expanded
+
+    return dequantized.to(torch.bfloat16).contiguous()
+
+
+def build_random_fp8_perchannel_weights():
+    """
+    Generate random BF16 weights and quantize to FP8 with per-channel scales.
+
+    Returns:
+        dict with fp8 weights, scales, and original bf16 for reference
+    """
+    torch.manual_seed(42)
+
+    # Generate random BF16 weights with small values
+    gate_proj = (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 100.0).to(
+        torch.bfloat16
+    )
+    up_proj = (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 100.0).to(
+        torch.bfloat16
+    )
+    down_proj = (torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32) / 100.0).to(
+        torch.bfloat16
+    )
+
+    # Quantize to FP8 with per-channel scales
+    gate_fp8, gate_scales = quantize_to_fp8_perchannel(gate_proj)
+    up_fp8, up_scales = quantize_to_fp8_perchannel(up_proj)
+    down_fp8, down_scales = quantize_to_fp8_perchannel(down_proj)
+
+    # Dequantize for reference computation
+    gate_deq = dequantize_fp8_perchannel(gate_fp8, gate_scales)
+    up_deq = dequantize_fp8_perchannel(up_fp8, up_scales)
+    down_deq = dequantize_fp8_perchannel(down_fp8, down_scales)
+
+    print(f"FP8 Per-Channel weights shape: gate={gate_fp8.shape}, up={up_fp8.shape}, down={down_fp8.shape}")
+    print(f"Per-Channel scales shape: gate={gate_scales.shape}, up={up_scales.shape}, down={down_scales.shape}")
+
+    # Debug: Print FP8 weight and scale info for expert 0
+    print("\n=== DEBUG: FP8 Per-Channel Weight and Scale Info (Expert 0) ===")
+    print(f"gate_fp8[0] first 8x8 block:")
+    for i in range(8):
+        print(f"  row {i}: {gate_fp8[0, i, :8].numpy().tobytes().hex(' ')}")
+    print(f"gate_fp8[0] stats: min={gate_fp8[0].min()}, max={gate_fp8[0].max()}")
+    print(f"gate_scales[0] first 8 channels: {gate_scales[0, :8]}")
+    print(f"gate_scales[0] stats: min={gate_scales[0].min():.6f}, max={gate_scales[0].max():.6f}")
+
+    return {
+        "gate_fp8": gate_fp8.contiguous(),
+        "up_fp8": up_fp8.contiguous(),
+        "down_fp8": down_fp8.contiguous(),
+        "gate_scales": gate_scales.contiguous(),
+        "up_scales": up_scales.contiguous(),
+        "down_scales": down_scales.contiguous(),
+        "gate_deq": gate_deq.contiguous(),
+        "up_deq": up_deq.contiguous(),
+        "down_deq": down_deq.contiguous(),
+    }
+
+
+def build_moes_from_fp8_perchannel_data(fp8_data: dict):
+    """
+    Build FP8 Per-Channel MoE modules from quantized data.
+    """
+    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 = 8
+            config.quant_config.group_size = 0  # Not used for per-channel
+            config.quant_config.zero_point = False
+            config.quant_config.per_channel = True  # Enable per-channel mode
+
+            # Set FP8 weight pointers
+            config.gate_proj = fp8_data["gate_fp8"].data_ptr()
+            config.up_proj = fp8_data["up_fp8"].data_ptr()
+            config.down_proj = fp8_data["down_fp8"].data_ptr()
+
+            # Set per-channel scale pointers
+            config.gate_scale = fp8_data["gate_scales"].data_ptr()
+            config.up_scale = fp8_data["up_scales"].data_ptr()
+            config.down_scale = fp8_data["down_scales"].data_ptr()
+            config.pool = CPUInfer.backend_
+
+            moe = kt_kernel_ext.moe.AMXFP8PerChannel_MOE(config)
+            CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+            CPUInfer.sync()
+            moes.append(moe)
+    return moes
+
+
+def run_fp8_perchannel_moe_test():
+    """
+    Run FP8 Per-Channel MoE validation test.
+    """
+    print("\n" + "=" * 70)
+    print("FP8 Per-Channel MoE Kernel Validation Test")
+    print("=" * 70)
+
+    # Build FP8 per-channel weights
+    print("\nGenerating and quantizing weights with per-channel scales...")
+    fp8_data = build_random_fp8_perchannel_weights()
+
+    # Build MoE modules
+    print("\nBuilding FP8 Per-Channel MoE modules...")
+    moes = build_moes_from_fp8_perchannel_data(fp8_data)
+
+    # Get dequantized weights for reference
+    gate_deq = fp8_data["gate_deq"]
+    up_deq = fp8_data["up_deq"]
+    down_deq = fp8_data["down_deq"]
+
+    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() / 100
+            input_tensor = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous() * 1.5
+            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()
+
+            assert not torch.isnan(output).any(), "NaN values detected in CPU expert output."
+            assert not torch.isinf(output).any(), "Inf values detected in CPU expert output."
+
+            # Reference computation using dequantized weights
+            t_output = moe_torch(input_tensor, expert_ids, weights, gate_deq, up_deq, down_deq)
+
+            t_output_flat = t_output.flatten()
+            output_flat = output.flatten()
+
+            diff = torch.mean(torch.abs(output_flat - t_output_flat)) / (torch.mean(torch.abs(t_output_flat)) + 1e-12)
+            diffs.append(diff.item())
+            print(f"Iteration {i}: relative L1 diff = {diff:.6f}")
+
+            if i < 3:  # Print detailed output for first few iterations
+                print(f"  kernel output: {output_flat[:debug_print_count]}")
+                print(f"  torch output:  {t_output_flat[:debug_print_count]}")
+
+    mean_diff = float(sum(diffs) / len(diffs))
+    max_diff = float(max(diffs))
+    min_diff = float(min(diffs))
+
+    print("\n" + "=" * 70)
+    print("FP8 Per-Channel MoE Test Results")
+    print("=" * 70)
+    print(f"Mean relative L1 diff: {mean_diff*100:.4f}%")
+    print(f"Max relative L1 diff:  {max_diff*100:.4f}%")
+    print(f"Min relative L1 diff:  {min_diff*100:.4f}%")
+
+    # Pass/Fail criteria
+    threshold = 15.0  # 15% relative error threshold for FP8
+    if mean_diff * 100 < threshold:
+        print(f"\nPASS: Mean error {mean_diff*100:.4f}% < {threshold}% threshold")
+    else:
+        print(f"\nFAIL: Mean error {mean_diff*100:.4f}% >= {threshold}% threshold")
+
+    return {"mean": mean_diff, "max": max_diff, "min": min_diff}
+
+
+if __name__ == "__main__":
+    run_fp8_perchannel_moe_test()

kt-kernel/examples/test_gate.py

@@ -0,0 +1,216 @@
+import math
+import os, sys
+import time
+from typing import Optional
+
+os.environ["BLAS_NUM_THREADS"] = "1"
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+from kt_kernel import kt_kernel_ext
+from kt_kernel_ext.kvcache import ggml_type
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+# from modeling_deepseek_v3 import MoEGate
+from configuration_deepseek_v3 import DeepseekV3Config
+
+seed = 42  # 你可以选择任何整数作为种子
+torch.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+
+seqlen = 64
+
+config = DeepseekV3Config()
+
+hidden_size = config.hidden_size
+num_experts_per_token = config.num_experts_per_tok
+n_routed_experts = config.n_routed_experts
+n_group = config.n_group
+topk_group = config.topk_group
+routed_scaling_factor = config.routed_scaling_factor
+
+weights = torch.randn((n_routed_experts, hidden_size), dtype=torch.float32).to("cpu").contiguous()
+bias = torch.randn((n_routed_experts,), dtype=torch.float32).to("cpu").contiguous()
+
+
+# weights = torch.randn((n_routed_experts, hidden_size), dtype=torch.float16).to('cpu').contiguous  ()
+def load_fp32_tensor(file_path, shape):
+    return torch.zeros(shape, dtype=torch.float32).to("cpu").contiguous()
+    with open(file_path, "rb") as f:
+        raw_data = f.read()
+    tensor = torch.frombuffer(raw_data, dtype=torch.float32)
+    tensor = tensor.view(shape)  # 根据你的 shape reshape
+    return tensor
+
+
+class MoEGate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.n_routed_experts = config.n_routed_experts
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.scoring_func = config.scoring_func
+        self.topk_method = config.topk_method
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+
+        # topk selection algorithm
+        self.norm_topk_prob = config.norm_topk_prob
+        self.gating_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
+        if self.topk_method == "noaux_tc":
+            self.e_score_correction_bias = nn.Parameter(torch.empty((self.n_routed_experts)))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        import torch.nn.init as init
+
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+
+    def forward(self, hidden_states):
+        bsz, seq_len, h = hidden_states.shape
+        ### compute gating score
+        hidden_states = hidden_states.view(-1, h)
+
+        h_to_check = load_fp32_tensor(
+            "/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug/gate_input", (seq_len, h)
+        )
+        diff = (h_to_check - hidden_states).abs().max()
+        # print("hidden_states diff:", diff)
+        # assert diff<0.02
+
+        bias_to_check = load_fp32_tensor(
+            "/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug/bias", (n_routed_experts)
+        )
+        diff = (bias - bias_to_check).abs().max()
+        # print('bias diff:',diff)
+        # assert diff < 0.02
+
+        logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32), None)
+
+        logits_to_check = load_fp32_tensor(
+            "/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug/gate_logits",
+            (seq_len, n_routed_experts),
+        )
+        diff = (logits_to_check - logits).abs().max()
+        # print("logits diff:", diff)
+        # assert diff < 0.02
+
+        if self.scoring_func == "sigmoid":
+            scores = logits.sigmoid()
+        else:
+            raise NotImplementedError(f"insupportable scoring function for MoE gating: {self.scoring_func}")
+
+        ### select top-k experts
+        if self.topk_method == "noaux_tc":
+            # assert not self.training
+            scores_for_choice = scores.view(bsz * seq_len, -1) + self.e_score_correction_bias.unsqueeze(0)
+
+            scores_to_check = load_fp32_tensor(
+                "/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug/scores_to_choice",
+                (seq_len, n_routed_experts),
+            )
+            diff = (scores_for_choice - scores_to_check).abs().max()
+            print(f"score for choice diff = {diff}")
+
+            group_scores = (
+                scores_for_choice.view(bsz * seq_len, self.n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+            )  # [n, n_group]
+
+            group_scores_to_check = load_fp32_tensor(
+                "/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug/group_scores",
+                (seq_len, n_group),
+            )
+            diff = (group_scores - group_scores_to_check).abs().max()
+            print(f"group scores diff = {diff}")
+
+            group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]  # [n, top_k_group]
+            group_mask = torch.zeros_like(group_scores)  # [n, n_group]
+            group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
+            score_mask = (
+                group_mask.unsqueeze(-1)
+                .expand(bsz * seq_len, self.n_group, self.n_routed_experts // self.n_group)
+                .reshape(bsz * seq_len, -1)
+            )  # [n, e]
+            tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))  # [n, e]
+            tmp_scores_to_check = load_fp32_tensor(
+                "/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug/gate_logits_toped",
+                (seq_len, n_routed_experts),
+            )
+            is_close = torch.isclose(tmp_scores, tmp_scores_to_check, rtol=1e-2, atol=1e-2, equal_nan=True)
+            print(f"tmp_score ok {is_close.all()}")
+
+            _, topk_idx = torch.topk(tmp_scores, k=self.top_k, dim=-1, sorted=False)
+            topk_weight = scores.gather(1, topk_idx)
+        else:
+            raise NotImplementedError(f"insupportable TopK function for MoE gating: {self.topk_method}")
+
+        ### norm gate to sum 1
+        if self.top_k > 1 and self.norm_topk_prob:
+            denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
+            topk_weight = topk_weight / denominator
+        topk_weight = topk_weight * self.routed_scaling_factor  # must multiply the scaling factor
+
+        return topk_idx, topk_weight
+
+
+def torch_gate(hidden_states):
+    hidden_states.unsqueeze_(0)
+    gate = MoEGate(config)
+    gate.weight.data = weights
+    gate.e_score_correction_bias.data = bias
+    y = gate(hidden_states)
+    # print(y)
+    return y
+
+
+def cpuinfer_gate(hidden_states):
+    config = kt_kernel_ext.gate.GateConfig(
+        hidden_size,
+        num_experts_per_token,
+        n_routed_experts,
+        n_group,
+        topk_group,
+    )
+
+    CPUInfer = kt_kernel_ext.CPUInfer(64)
+    config.routed_scaling_factor = routed_scaling_factor
+
+    config.pool = CPUInfer.backend_
+    config.weight = weights.data_ptr()
+    config.weight_type = ggml_type.FP32
+    config.e_score_correction_bias = bias.data_ptr()
+    config.e_score_correction_bias_type = ggml_type.FP32
+
+    gate = kt_kernel_ext.gate.MoEGate(config)
+
+    expert_ids = torch.zeros((seqlen, num_experts_per_token), dtype=torch.int64).to("cpu").contiguous()
+    expert_weights = torch.zeros((seqlen, num_experts_per_token), dtype=torch.float32).to("cpu").contiguous()
+
+    gate.forward(seqlen, hidden_states.data_ptr(), expert_ids.data_ptr(), expert_weights.data_ptr())
+
+    # print(expert_ids,expert_weights)
+    return expert_ids, expert_weights
+
+
+input = torch.randn(seqlen, hidden_size, dtype=torch.float32).to("cpu").contiguous()
+# print(input)
+ids, we = cpuinfer_gate(input)
+idx = torch.argsort(ids, dim=-1, descending=True)
+ids = torch.gather(ids, dim=-1, index=idx)
+we = torch.gather(we, dim=-1, index=idx)
+
+
+std_ids, std_we = torch_gate(input)
+idx = torch.argsort(std_ids, dim=-1, descending=True)
+std_we = torch.gather(std_we, dim=-1, index=idx)
+std_ids = torch.gather(std_ids, dim=-1, index=idx)
+
+
+# print("ids diff:", torch.abs(std_ids - ids).max())
+# print("weights diff:", torch.abs(std_we - we).max())
+assert torch.abs(std_ids - ids).max() == 0, "Expert IDs do not match!"
+assert torch.abs(std_we - we).max() < 1e-2, "Expert Weights do not match!"
+print("Expert IDs and Weights match successfully!")

kt-kernel/examples/test_k2_moe_amx.py

@@ -0,0 +1,320 @@
+import math
+import os
+import sys
+from typing import Dict, Literal
+
+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
+
+qlen = 1
+layer_num = 1
+CPUInfer = kt_kernel_ext.CPUInfer(40)
+validation_iter = 10
+k_group_size = 32
+debug_print_count = 16
+
+physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous()
+
+
+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())
+    print(f"gate_buf: {gate_buf}")
+    print(f"up_buf: {up_buf}")
+    intermediate = act_fn(gate_buf) * up_buf
+    ret = torch.mm(intermediate, down_proj.t())
+    print(f"intermediate: {intermediate}")
+    print(f"mlp output: {ret}")
+    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 pack_to_int32(value: torch.Tensor, num_bits: int, packed_dim: Literal[0, 1] = 1) -> torch.Tensor:
+    if value.dtype is not torch.int8:
+        raise ValueError("Tensor must be torch.int8 before packing")
+    if not (1 <= num_bits <= 8):
+        raise ValueError(f"num_bits must be in [1, 8], got {num_bits}")
+
+    offset = 1 << (num_bits - 1)
+    value = (value + offset).to(torch.uint8)
+    device = value.device
+
+    pack_factor = 32 // num_bits
+
+    if packed_dim == 0:
+        value = value.transpose(0, 1)
+
+    rows, cols = value.shape
+    padded_cols = math.ceil(cols / pack_factor) * pack_factor
+    pad_len = padded_cols - cols
+
+    if pad_len > 0:
+        value = torch.nn.functional.pad(value, (0, pad_len))
+
+    num_groups = padded_cols // pack_factor
+
+    # Use int32 here
+    reshaped = value.view(rows, num_groups, pack_factor).to(torch.int32)
+    bit_shifts = torch.arange(pack_factor, device=device, dtype=torch.int32) * num_bits
+    packed = (reshaped << bit_shifts).sum(dim=2, dtype=torch.int32)
+
+    if packed_dim == 0:
+        packed = packed.transpose(0, 1)
+
+    return packed
+
+
+def pack_tensor_per_row(q: torch.Tensor, num_bits: int) -> torch.Tensor:
+    e, rows, cols = q.shape
+    flat = q.view(e * rows, cols)
+    packed = pack_to_int32(flat, num_bits)
+    return packed.view(e, rows, -1).contiguous()
+
+
+def quantize_k2_tensor(weights: torch.Tensor, group_size: int):
+    """
+    Symmetric max-abs/7 quantization per k-group following compressed_tensors packing.
+    Args:
+        weights: [expert_num, rows (N), cols (K)]
+    Returns:
+        packed_q: int32 tensor storing 8 int4s per element with 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 / 7.0).squeeze(-1)
+    q = torch.round(reshaped / scales.unsqueeze(-1)).clamp(-8, 7).to(torch.int8)
+    q = q.view(e, rows, cols)
+    packed = pack_tensor_per_row(q, num_bits=4).view(e, rows, cols // 8).contiguous()
+    scales = scales.to(torch.bfloat16).contiguous().view(e, rows, cols // group_size).contiguous()
+
+    print(f"Quantized weights: {packed.shape}, scales: {scales.shape}")
+    print(f"Quantized tensors: \n{packed},\n {scales}")
+    return packed, scales
+
+
+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_k2_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 = quantize_k2_tensor(gate_proj, k_group_size)
+    up_q, up_scales = quantize_k2_tensor(up_proj, k_group_size)
+    down_q, down_scales = quantize_k2_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(),
+        "original_fp16": {
+            "gate_proj": gate_proj.to(torch.float16).contiguous(),
+            "up_proj": up_proj.to(torch.float16).contiguous(),
+            "down_proj": down_proj.to(torch.float16).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.AMXInt4_KGroup_MOE(config)
+            CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+            CPUInfer.sync()
+            # CPUInfer.submit(moe.warm_up_task())
+            # CPUInfer.sync()
+            moes.append(moe)
+    return moes
+
+
+def run_case(pattern: str) -> Dict[str, float]:
+    print("\n" + "=" * 70)
+    desc = WEIGHT_PATTERNS[pattern][0]
+    print(f"Running case: {pattern} -> {desc}")
+    print("=" * 70)
+
+    quant_data = prepare_k2_quantized_weights(pattern)
+    moes = build_moes_from_quantized_data(quant_data)
+
+    original_weights = quant_data["original_fp16"]
+    gate_fp16 = original_weights["gate_proj"]
+    up_fp16 = original_weights["up_proj"]
+    down_fp16 = original_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()
+
+            input_tensor_fp16 = input_tensor.to(torch.float16)
+            t_output = moe_torch(input_tensor_fp16, expert_ids, weights, gate_fp16, up_fp16, down_fp16).to(
+                torch.bfloat16
+            )
+
+            t_output = t_output.flatten()
+            output = output.flatten()
+
+            diff = torch.mean(torch.abs(output - t_output)) / (torch.mean(torch.abs(t_output)) + 1e-12)
+            diffs.append(diff.item())
+            print(f"[{pattern}] Iteration {i}: relative L1 diff = {diff:.4f}")
+            print(f"           output   {output}")
+            print(f"           t_output {t_output}")
+
+    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_k2_moe_test():
+    summary_rows = []
+    for case_name in WEIGHT_PATTERNS.keys():
+        results = run_case(case_name)
+        summary_rows.append(results)
+        # break
+
+    print("\n=== Case vs. Relative Error Summary ===")
+    print(f"{'Case':<20} {'Mean':>10} {'Max':>10} {'Min':>10}")
+    for row in summary_rows:
+        print(f"{row['case']:<20} {row['mean']*100:9.2f}% {row['max']*100:9.2f}% {row['min']*100:9.2f}%")
+
+
+if __name__ == "__main__":
+    run_k2_moe_test()

kt-kernel/examples/test_k2_write_buffer.py

@@ -0,0 +1,350 @@
+import os
+import sys
+import time
+
+import torch
+import numpy as np
+
+
+from kt_kernel import kt_kernel_ext
+from kt_kernel_ext import CPUInfer
+
+
+def make_cpu_infer(thread_num=80):
+    return CPUInfer(thread_num)
+
+
+def build_config(cpuinfer, expert_num, num_experts_per_tok, hidden_size, intermediate_size, group_size):
+    cfg = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size)
+    cfg.max_len = 1
+    cfg.quant_config.bits = 4
+    cfg.quant_config.group_size = group_size
+    cfg.quant_config.zero_point = False
+    cfg.pool = cpuinfer.backend_
+    return cfg
+
+
+def allocate_weights(expert_num, hidden_size, intermediate_size, group_size):
+    # packed int4 weights: 2 values per byte
+    per_mat_weight_bytes = (hidden_size * intermediate_size) // 2
+    per_mat_scale_elems = (hidden_size * intermediate_size) // group_size
+
+    gate_q = torch.randint(0, 256, (expert_num * per_mat_weight_bytes,), dtype=torch.uint8)
+    up_q = torch.randint(0, 256, (expert_num * per_mat_weight_bytes,), dtype=torch.uint8)
+    down_q = torch.randint(0, 256, (expert_num * per_mat_weight_bytes,), dtype=torch.uint8)
+
+    gate_scale = torch.randn(expert_num * per_mat_scale_elems, dtype=torch.bfloat16)
+    up_scale = torch.randn(expert_num * per_mat_scale_elems, dtype=torch.bfloat16)
+    down_scale = torch.randn(expert_num * per_mat_scale_elems, dtype=torch.bfloat16)
+
+    return (
+        gate_q,
+        up_q,
+        down_q,
+        gate_scale,
+        up_scale,
+        down_scale,
+        per_mat_weight_bytes,
+        per_mat_scale_elems,
+    )
+
+
+def test_with_tp(gpu_tp_count):
+    """Test write_weight_scale_to_buffer with a specific gpu_tp_count"""
+    torch.manual_seed(123)
+
+    expert_num = 8  # Reduced for faster testing
+    gpu_experts = expert_num  # Number of experts on GPU
+
+    num_experts_per_tok = 8
+    hidden_size = 7168
+    intermediate_size = 2048
+    group_size = 32
+
+    cpuinfer = make_cpu_infer()
+    cfg = build_config(cpuinfer, expert_num, num_experts_per_tok, hidden_size, intermediate_size, group_size)
+
+    (
+        gate_q,
+        up_q,
+        down_q,
+        gate_scale,
+        up_scale,
+        down_scale,
+        per_mat_weight_bytes,
+        per_mat_scale_elems,
+    ) = allocate_weights(expert_num, hidden_size, intermediate_size, group_size)
+
+    cfg.gate_proj = gate_q.data_ptr()
+    cfg.up_proj = up_q.data_ptr()
+    cfg.down_proj = down_q.data_ptr()
+    cfg.gate_scale = gate_scale.data_ptr()
+    cfg.up_scale = up_scale.data_ptr()
+    cfg.down_scale = down_scale.data_ptr()
+
+    moe = kt_kernel_ext.moe.AMXInt4_KGroup_MOE(cfg)
+
+    physical_to_logical_map = torch.arange(expert_num, dtype=torch.int64, device="cpu").contiguous()
+    cpuinfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+    cpuinfer.sync()
+
+    # TP configuration
+    # Calculate sizes per TP part (per expert)
+    weight_bytes_per_expert_per_tp = per_mat_weight_bytes // gpu_tp_count
+    scale_elems_per_expert_per_tp = per_mat_scale_elems // gpu_tp_count
+
+    # Total sizes for all gpu_experts
+    total_weight_bytes_per_tp = gpu_experts * weight_bytes_per_expert_per_tp
+    total_scale_elems_per_tp = gpu_experts * scale_elems_per_expert_per_tp
+
+    # Create buffer lists for w13 (gate+up) and w2 (down)
+    # These hold all experts' data for each GPU TP
+    w13_weight_bufs = []
+    w13_scale_bufs = []
+    w2_weight_bufs = []
+    w2_scale_bufs = []
+
+    for tp_idx in range(gpu_tp_count):
+        # w13 combines gate and up, so needs 2x the size per expert
+        w13_weight_bufs.append(torch.empty(2 * total_weight_bytes_per_tp, dtype=torch.uint8))
+        w13_scale_bufs.append(torch.empty(2 * total_scale_elems_per_tp, dtype=torch.bfloat16))
+        w2_weight_bufs.append(torch.empty(total_weight_bytes_per_tp, dtype=torch.uint8))
+        w2_scale_bufs.append(torch.empty(total_scale_elems_per_tp, dtype=torch.bfloat16))
+
+    print(f"Total experts: {expert_num}, GPU experts: {gpu_experts}")
+    print(f"GPU TP count: {gpu_tp_count}")
+    print(f"Original per matrix weight bytes: {per_mat_weight_bytes}")
+    print(f"Original per matrix scale elements: {per_mat_scale_elems}")
+    print(f"Weight bytes per expert per TP: {weight_bytes_per_expert_per_tp}")
+    print(f"Scale elements per expert per TP: {scale_elems_per_expert_per_tp}")
+    print(f"Total weight bytes per TP (w13): {2 * total_weight_bytes_per_tp}")
+    print(f"Total weight bytes per TP (w2): {total_weight_bytes_per_tp}")
+
+    # Helper function to get pointers with expert offset
+    # K2 write_weights_to_buffer writes one expert at a time, so we need to pass
+    # pointers that already point to the correct location for each expert
+    def get_expert_ptrs(expert_id):
+        w13_weight_ptrs = []
+        w13_scale_ptrs = []
+        w2_weight_ptrs = []
+        w2_scale_ptrs = []
+
+        for tp_idx in range(gpu_tp_count):
+            # Calculate byte offsets for this expert
+            # w13: gate_weight + up_weight interleaved by expert
+            # Layout: [expert0_gate, expert0_up, expert1_gate, expert1_up, ...]
+            w13_weight_expert_offset = expert_id * 2 * weight_bytes_per_expert_per_tp
+            w13_scale_expert_offset = expert_id * 2 * scale_elems_per_expert_per_tp
+            w2_weight_expert_offset = expert_id * weight_bytes_per_expert_per_tp
+            w2_scale_expert_offset = expert_id * scale_elems_per_expert_per_tp
+
+            w13_weight_ptrs.append(w13_weight_bufs[tp_idx].data_ptr() + w13_weight_expert_offset)
+            w13_scale_ptrs.append(w13_scale_bufs[tp_idx].data_ptr() + w13_scale_expert_offset * 2)  # bf16 = 2 bytes
+            w2_weight_ptrs.append(w2_weight_bufs[tp_idx].data_ptr() + w2_weight_expert_offset)
+            w2_scale_ptrs.append(w2_scale_bufs[tp_idx].data_ptr() + w2_scale_expert_offset * 2)  # bf16 = 2 bytes
+
+        return w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs
+
+    # Warm up
+    for i in range(2):
+        for expert_id in range(gpu_experts):
+            w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs = get_expert_ptrs(expert_id)
+            cpuinfer.submit(
+                moe.write_weight_scale_to_buffer_task(
+                    gpu_tp_count=gpu_tp_count,
+                    expert_id=expert_id,
+                    w13_weight_ptrs=w13_weight_ptrs,
+                    w13_scale_ptrs=w13_scale_ptrs,
+                    w2_weight_ptrs=w2_weight_ptrs,
+                    w2_scale_ptrs=w2_scale_ptrs,
+                )
+            )
+            cpuinfer.sync()
+
+    # Timing
+    begin_time = time.perf_counter_ns()
+    for expert_id in range(gpu_experts):
+        w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs = get_expert_ptrs(expert_id)
+        cpuinfer.submit(
+            moe.write_weight_scale_to_buffer_task(
+                gpu_tp_count=gpu_tp_count,
+                expert_id=expert_id,
+                w13_weight_ptrs=w13_weight_ptrs,
+                w13_scale_ptrs=w13_scale_ptrs,
+                w2_weight_ptrs=w2_weight_ptrs,
+                w2_scale_ptrs=w2_scale_ptrs,
+            )
+        )
+        cpuinfer.sync()
+    end_time = time.perf_counter_ns()
+    elapsed_ms = (end_time - begin_time) / 1000000
+    total_weights = hidden_size * intermediate_size * gpu_experts * 3
+    total_bytes = total_weights // group_size * 2 + total_weights // 2  # scale (bf16) + weight (int4)
+    print(f"write_weight_scale_to_buffer time: {elapsed_ms:.2f} ms")
+    print(f"Throughput: {total_bytes / (elapsed_ms * 1e6):.2f} GB/s")
+
+    def split_expert_tensor(tensor, chunk):
+        """Split tensor by experts"""
+        return [tensor[i * chunk : (i + 1) * chunk] for i in range(expert_num)]
+
+    # Split by experts first
+    gate_q_experts = split_expert_tensor(gate_q, per_mat_weight_bytes)
+    up_q_experts = split_expert_tensor(up_q, per_mat_weight_bytes)
+    down_q_experts = split_expert_tensor(down_q, per_mat_weight_bytes)
+
+    gate_scale_experts = split_expert_tensor(gate_scale, per_mat_scale_elems)
+    up_scale_experts = split_expert_tensor(up_scale, per_mat_scale_elems)
+    down_scale_experts = split_expert_tensor(down_scale, per_mat_scale_elems)
+
+    # Verify buffers for each TP part
+    for tp_idx in range(gpu_tp_count):
+        expected_w13_weights = []
+        expected_w13_scales = []
+        expected_w2_weights = []
+        expected_w2_scales = []
+
+        weight13_per_tp = per_mat_weight_bytes // gpu_tp_count
+        scale13_per_tp = per_mat_scale_elems // gpu_tp_count
+
+        # Process each GPU expert
+        for expert_id in range(gpu_experts):
+            # For w13 (gate and up), the slicing is straightforward
+            start_weight = tp_idx * weight13_per_tp
+            end_weight = (tp_idx + 1) * weight13_per_tp
+            start_scale = tp_idx * scale13_per_tp
+            end_scale = (tp_idx + 1) * scale13_per_tp
+
+            # Gate
+            gate_weight_tp = gate_q_experts[expert_id][start_weight:end_weight]
+            gate_scale_tp = gate_scale_experts[expert_id][start_scale:end_scale]
+
+            # Up
+            up_weight_tp = up_q_experts[expert_id][start_weight:end_weight]
+            up_scale_tp = up_scale_experts[expert_id][start_scale:end_scale]
+
+            # Down matrix needs special handling because it's sliced column-wise
+            # We need to reconstruct it from column slices
+            down_weight_tp_parts = []
+            down_scale_tp_parts = []
+
+            # Iterate through each column to extract the corresponding parts
+            for col_idx in range(hidden_size):
+                col_weight_start = col_idx * (intermediate_size // 2)
+                col_scale_start = col_idx * (intermediate_size // group_size)
+
+                # Direct mapping: each CPU TP corresponds to a GPU TP
+                tp_slice_weight_size = (intermediate_size // gpu_tp_count) // 2
+                tp_slice_scale_size = (intermediate_size // gpu_tp_count) // group_size
+
+                tp_weight_offset = col_weight_start + tp_idx * tp_slice_weight_size
+                tp_scale_offset = col_scale_start + tp_idx * tp_slice_scale_size
+
+                down_weight_tp_parts.append(
+                    down_q_experts[expert_id][tp_weight_offset : tp_weight_offset + tp_slice_weight_size]
+                )
+                down_scale_tp_parts.append(
+                    down_scale_experts[expert_id][tp_scale_offset : tp_scale_offset + tp_slice_scale_size]
+                )
+
+            # Concatenate all column slices for this TP
+            down_weight_tp = torch.cat(down_weight_tp_parts)
+            down_scale_tp = torch.cat(down_scale_tp_parts)
+
+            # Append to expected lists - interleaved by expert: [gate0, up0, gate1, up1, ...]
+            expected_w13_weights.append(gate_weight_tp)
+            expected_w13_weights.append(up_weight_tp)
+            expected_w13_scales.append(gate_scale_tp)
+            expected_w13_scales.append(up_scale_tp)
+            expected_w2_weights.append(down_weight_tp)
+            expected_w2_scales.append(down_scale_tp)
+
+        # Concatenate all experts for this TP part
+        expected_w13_weight = torch.cat(expected_w13_weights)
+        expected_w13_scale = torch.cat(expected_w13_scales)
+        expected_w2_weight = torch.cat(expected_w2_weights)
+        expected_w2_scale = torch.cat(expected_w2_scales)
+
+        print(f"=== Checking TP part {tp_idx} ===")
+        print(f"  w13 weight shape: actual={w13_weight_bufs[tp_idx].shape}, expected={expected_w13_weight.shape}")
+        print(f"  w13 scale shape: actual={w13_scale_bufs[tp_idx].shape}, expected={expected_w13_scale.shape}")
+        print(f"  w2 weight shape: actual={w2_weight_bufs[tp_idx].shape}, expected={expected_w2_weight.shape}")
+        print(f"  w2 scale shape: actual={w2_scale_bufs[tp_idx].shape}, expected={expected_w2_scale.shape}")
+
+        # Assert all checks pass
+        if not torch.equal(w13_weight_bufs[tp_idx], expected_w13_weight):
+            diff_mask = w13_weight_bufs[tp_idx] != expected_w13_weight
+            first_diff_idx = diff_mask.nonzero()[0].item() if diff_mask.any() else -1
+            print(f"  w13 weight mismatch at index {first_diff_idx}")
+            print(f"    actual: {w13_weight_bufs[tp_idx][first_diff_idx:first_diff_idx+10]}")
+            print(f"    expected: {expected_w13_weight[first_diff_idx:first_diff_idx+10]}")
+            raise AssertionError(f"w13 weight bytes mismatch for TP {tp_idx}")
+
+        if not torch.allclose(w13_scale_bufs[tp_idx], expected_w13_scale):
+            diff = torch.abs(w13_scale_bufs[tp_idx].float() - expected_w13_scale.float())
+            max_diff_idx = diff.argmax().item()
+            print(f"  w13 scale mismatch, max diff at index {max_diff_idx}")
+            print(f"    actual: {w13_scale_bufs[tp_idx][max_diff_idx]}")
+            print(f"    expected: {expected_w13_scale[max_diff_idx]}")
+            raise AssertionError(f"w13 scale values mismatch for TP {tp_idx}")
+
+        if not torch.equal(w2_weight_bufs[tp_idx], expected_w2_weight):
+            diff_mask = w2_weight_bufs[tp_idx] != expected_w2_weight
+            first_diff_idx = diff_mask.nonzero()[0].item() if diff_mask.any() else -1
+            print(f"  w2 weight mismatch at index {first_diff_idx}")
+            print(f"    actual: {w2_weight_bufs[tp_idx][first_diff_idx:first_diff_idx+10]}")
+            print(f"    expected: {expected_w2_weight[first_diff_idx:first_diff_idx+10]}")
+            raise AssertionError(f"w2 weight bytes mismatch for TP {tp_idx}")
+
+        if not torch.allclose(w2_scale_bufs[tp_idx], expected_w2_scale):
+            diff = torch.abs(w2_scale_bufs[tp_idx].float() - expected_w2_scale.float())
+            max_diff_idx = diff.argmax().item()
+            print(f"  w2 scale mismatch, max diff at index {max_diff_idx}")
+            print(f"    actual: {w2_scale_bufs[tp_idx][max_diff_idx]}")
+            print(f"    expected: {expected_w2_scale[max_diff_idx]}")
+            raise AssertionError(f"w2 scale values mismatch for TP {tp_idx}")
+
+    print(
+        f"\n✓ write_weight_scale_to_buffer passed: extracted {gpu_experts} GPU experts across {gpu_tp_count} TP parts"
+    )
+    return True
+
+
+def main():
+    """Run tests for all gpu_tp_count values: 1, 2, 4, 8"""
+    tp_values = [1, 2, 4, 8]
+    all_passed = True
+    results = {}
+
+    print("=" * 60)
+    print("Testing K2 write_weight_scale_to_buffer for TP = 1, 2, 4, 8")
+    print("=" * 60)
+
+    for tp in tp_values:
+        print(f"\n{'='*60}")
+        print(f"Testing with gpu_tp_count = {tp}")
+        print(f"{'='*60}")
+        try:
+            test_with_tp(tp)
+            results[tp] = "PASSED"
+            print(f"✓ TP={tp} PASSED")
+        except Exception as e:
+            results[tp] = f"FAILED: {e}"
+            all_passed = False
+            print(f"✗ TP={tp} FAILED: {e}")
+
+    print("\n" + "=" * 60)
+    print("SUMMARY")
+    print("=" * 60)
+    for tp, result in results.items():
+        status = "✓" if "PASSED" in result else "✗"
+        print(f"  {status} TP={tp}: {result}")
+
+    if all_passed:
+        print("\n✓ ALL TESTS PASSED")
+    else:
+        print("\n✗ SOME TESTS FAILED")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

kt-kernel/examples/test_linear.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""
+Description  :
+Author       : chenht2022
+Date         : 2024-07-25 10:32:05
+Version      : 1.0.0
+LastEditors  : chenht2022
+LastEditTime : 2024-08-06 10:36:59
+Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+"""
+import os, sys
+import time
+
+sys.path.append(os.path.dirname(__file__) + "/../build")
+from kt_kernel import kt_kernel_ext
+import torch
+
+input_size = 16384
+output_size = 5120
+stride = 32
+group_max_len = 1024
+proj_type = 1  # ggml_type::GGML_TYPE_F16
+hidden_type = 1  # ggml_type::GGML_TYPE_F16
+qlen = 30
+layer_num = 10
+CPUInfer = kt_kernel_ext.CPUInfer(48)
+validation_iter = 100
+
+with torch.inference_mode(mode=True):
+    linears = []
+    projs = []
+    for _ in range(layer_num):
+        proj = torch.randn((output_size, input_size), dtype=torch.float16, device="cuda").to("cpu").contiguous()
+        config = kt_kernel_ext.linear.LinearConfig(
+            input_size, output_size, stride, group_max_len, proj.data_ptr(), proj_type, hidden_type
+        )
+        linear = kt_kernel_ext.linear.Linear(config)
+        projs.append(proj)
+        linears.append(linear)
+
+    # validation
+    for i in range(validation_iter):
+        linear = linears[i % layer_num]
+        input = torch.randn((qlen, input_size), dtype=torch.float16).contiguous()
+        output = torch.empty((qlen, output_size), dtype=torch.float16).contiguous()
+        input = input / 100
+
+        CPUInfer.submit(linear.forward(qlen, input.data_ptr(), output.data_ptr()))
+        CPUInfer.sync()
+        # print('cpuinfer output', output)
+
+        proj = projs[i % layer_num]
+        t_output = torch.mm(input, proj.t())
+        # print('torch output', t_output)
+
+        diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
+        print("diff = ", diff)
+        assert diff < 0.001

kt-kernel/examples/test_mla.py

@@ -0,0 +1,724 @@
+import logging
+import os, sys
+import time
+from typing import Optional
+
+os.environ["BLAS_NUM_THREADS"] = "1"
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+from kt_kernel import kt_kernel_ext
+from kt_kernel_ext.kvcache import ggml_type
+import torch
+from torch import inf, nn
+from torch.nn import init
+from torch_attention import apply_rotary_pos_emb, DeepseekV2RMSNorm, KDeepSeekV3Cache, DeepseekV3YarnRotaryEmbedding
+
+logger = logging.getLogger("reader")
+
+from gguf.gguf_reader import GGUFReader
+
+
+def read_gguf_file(gguf_file_path):
+    """
+    Reads and prints key-value pairs and tensor information from a GGUF file in an improved format.
+
+    Parameters:
+    - gguf_file_path: Path to the GGUF file.
+    """
+
+    reader = GGUFReader(gguf_file_path)
+
+    # List all key-value pairs in a columnized format
+    # print("Key-Value Pairs:") # noqa: NP100
+    # max_key_length = max(len(key) for key in reader.fields.keys())
+    for key, field in reader.fields.items():
+        value = field.parts[field.data[0]]
+        # print(f"{key:{max_key_length}} : {value}") # noqa: NP100
+    # print("----") # noqa: NP100
+
+    # List all tensors
+    # print("Tensors:") # noqa: NP100
+    # tensor_info_format = "{:<30} | Shape: {:<15} | Size: {:<12} | Quantization: {}"
+    # print(tensor_info_format.format("Tensor Name", "Shape", "Size", "Quantization")) # noqa: NP100
+    # print("-" * 80) # noqa: NP100
+    re = []
+    for tensor in reader.tensors:
+        shape_str = "x".join(map(str, tensor.shape))
+        size_str = str(tensor.n_elements)
+        quantization_str = tensor.tensor_type.name
+        # print(tensor_info_format.format(tensor.name, shape_str, size_str, quantization_str)) # noqa: NP100
+        re.append(tensor)
+    return re
+
+
+def get_torch_tensor_from_gguf(gguf_weights, name):
+    return torch.from_numpy(gguf_weights[name].data).contiguous()
+
+
+def get_torch_tensor_and_type_from_gguf(gguf_weights, name):
+    return torch.from_numpy(gguf_weights[name].data).contiguous(), gguf_weights[name].tensor_type.name
+
+
+def type_to_ggml_type(type):
+    if type == "F32":
+        return ggml_type.FP32
+    elif type == "F16":
+        return ggml_type.FP16
+    elif type == "BF16":
+        return ggml_type.BF16
+    else:
+        raise ValueError(f"Unsupported data type: {type}")
+
+
+use_real_weights = True
+gguf_path = "/home/bd/models/DeepSeek-R1-BF16"
+
+seed = 42  # 你可以选择任何整数作为种子
+torch.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+
+qlen = 3212
+kvlen = 0
+
+
+page_table = range(20)
+bsz_tensors = torch.tensor([1])
+
+
+page_size = 256
+pages_count = 200
+tp_count = 4
+
+
+hidden_size = 7168
+q_lora_rank = 1536
+kv_lora_rank = 512
+num_heads = 128
+nope_size = 128
+rope_size = 64
+
+rope_theta = 10000
+max_qlen = 4096
+max_kvlen = 4096
+
+max_position_embeddings = 163840
+
+
+rope_scaling = {
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "factor": 40,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0,
+    "original_max_position_embeddings": 4096,
+    "type": "yarn",
+}
+
+
+CPUInfer = kt_kernel_ext.CPUInfer(30)
+validation_iter = 100
+
+
+# data_type = torch.float32
+weight_type = torch.bfloat16
+# weight_type = torch.float16
+
+
+input_type = {
+    torch.float32: torch.float32,
+    torch.float16: torch.float16,
+    torch.bfloat16: torch.float32,
+}[weight_type]
+
+q_a_proj = nn.Linear(hidden_size, q_lora_rank, bias=False, dtype=weight_type)
+q_b_proj = nn.Linear(q_lora_rank, num_heads * (nope_size + rope_size), bias=False, dtype=weight_type)
+kv_a_proj_with_mqa = nn.Linear(hidden_size, kv_lora_rank + rope_size, bias=False, dtype=weight_type)
+kv_b_proj = nn.Linear(num_heads * (nope_size + nope_size), kv_lora_rank, bias=False, dtype=weight_type)
+o_proj = nn.Linear(num_heads * nope_size, hidden_size, bias=False, dtype=weight_type)
+q_a_norm = torch.ones(hidden_size, dtype=torch.float32)
+kv_a_norm = torch.ones(hidden_size, dtype=torch.float32)
+
+
+def read_gguf_directory(directory):
+    """
+    Reads all GGUF files in a directory and prints their contents.
+
+    Parameters:
+    - directory: Path to the directory containing GGUF files.
+    """
+    if not os.path.isdir(directory):
+        logger.error(f"Directory {directory} does not exist.")
+        return
+
+    # List all GGUF files in the directory
+    files = [f for f in os.listdir(directory) if f.endswith(".gguf")]
+    if not files:
+        logger.info(f"No GGUF files found in {directory}.")
+        return
+
+    re = []
+    for file in files:
+        file_path = os.path.join(directory, file)
+        # print(f"Reading {file_path}:") # noqa: NP100
+        # print("\n") # noqa: NP100
+        re.extend(read_gguf_file(file_path))
+    re = {r.name: r for r in re}
+    return re
+
+
+if use_real_weights := True:
+    gguf_weights = read_gguf_directory(gguf_path)
+    layer_idx = 0
+    q_a_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_a.weight")
+    q_a_proj.weight = nn.Parameter(q_a_proj_weight.view(torch.bfloat16), requires_grad=False)
+    q_a_type = type
+
+    q_a_norm_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_a_norm.weight")
+    q_a_norm = q_a_norm_weight.view(torch.float32)
+    # config.q_a_norm = q_a_norm_weight.data_ptr()
+    # config.q_a_norm_type = type_to_ggml_type(type)
+
+    q_b_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_b.weight")
+    q_b_proj.weight = nn.Parameter(q_b_proj_weight.view(torch.bfloat16), requires_grad=False)
+
+    kv_a_proj_with_mqa_weight, type = get_torch_tensor_and_type_from_gguf(
+        gguf_weights, f"blk.{layer_idx}.attn_kv_a_mqa.weight"
+    )
+    kv_a_proj_with_mqa.weight = nn.Parameter(kv_a_proj_with_mqa_weight.view(torch.bfloat16), requires_grad=False)
+
+    kv_a_norm_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_kv_a_norm.weight")
+    kv_a_norm = kv_a_norm_weight.view(torch.float32)
+    # config.kv_a_norm = kv_a_norm_weight.data_ptr()
+    # config.kv_a_norm_type = type_to_ggml_type(type)
+
+    kv_b_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_kv_b.weight")
+    kv_b_proj.weight = nn.Parameter(kv_b_proj_weight.view(torch.bfloat16), requires_grad=False)
+
+    o_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_output.weight")
+    o_proj.weight = nn.Parameter(o_proj_weight.view(torch.bfloat16), requires_grad=False)
+
+else:
+    init.normal_(q_a_proj.weight, mean=0.0, std=0.02)
+    init.normal_(q_b_proj.weight, mean=0.0, std=0.02)
+    init.normal_(kv_a_proj_with_mqa.weight, mean=0.0, std=0.02)
+    init.normal_(kv_b_proj.weight, mean=0.0, std=0.02)
+    init.normal_(o_proj.weight, mean=0.0, std=0.02)
+
+x_reshaped = kv_b_proj.weight.view(num_heads, 2, nope_size, kv_lora_rank)
+q_absorb = x_reshaped[:, 0]
+out_absorb = x_reshaped[:, 1]
+
+
+hidden_states = torch.randn((qlen, hidden_size), dtype=input_type).to("cpu").contiguous()
+
+
+def test_cpu_mla():
+    os.environ["BLAS_NUM_THREADS"] = "1"
+    q_a_proj_weight = q_a_proj.weight.to(weight_type).to("cpu").contiguous()
+    q_b_proj_weight = q_b_proj.weight.to(weight_type).to("cpu").contiguous()
+    kv_a_proj_with_mqa_weight = kv_a_proj_with_mqa.weight.to("cpu").to(weight_type).contiguous()
+    kv_b_proj_weight = kv_b_proj.weight.to(weight_type).to("cpu").contiguous()
+    o_proj_weight = o_proj.weight.to(weight_type).to("cpu").contiguous()
+
+    config = kt_kernel_ext.mla.MLAConfig(
+        hidden_size,
+        q_lora_rank,
+        kv_lora_rank,
+        num_heads,
+        nope_size,
+        rope_size,
+    )
+    config.max_qlen = max_qlen
+    config.max_kvlen = max_kvlen
+    config.max_position_embeddings = max_position_embeddings
+    config.rope_scaling_factor = rope_scaling["factor"]
+    config.rope_theta = rope_theta
+    config.rope_scaling_beta_fast = rope_scaling["beta_fast"]
+    config.rope_scaling_beta_slow = rope_scaling["beta_slow"]
+    config.rope_scaling_mscale = rope_scaling["mscale"]
+    config.rope_scaling_mscale_all_dim = rope_scaling["mscale_all_dim"]
+    config.rope_scaling_original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
+
+    config.q_a_proj = q_a_proj_weight.data_ptr()
+    config.q_b_proj = q_b_proj_weight.data_ptr()
+    config.kv_a_proj_with_mqa = kv_a_proj_with_mqa_weight.data_ptr()
+    config.kv_b_proj = kv_b_proj_weight.data_ptr()
+    config.o_proj = o_proj_weight.data_ptr()
+
+    config.q_a_norm = q_a_norm.data_ptr()
+    config.q_a_norm_type = ggml_type.FP32
+    config.kv_a_norm = kv_a_norm.data_ptr()
+    config.kv_a_norm_type = ggml_type.FP32
+    config.page_count = pages_count
+
+    if weight_type == torch.float32:
+        config.q_a_proj_type = ggml_type.FP32
+        config.q_b_proj_type = ggml_type.FP32
+        config.kv_a_proj_with_mqa_type = ggml_type.FP32
+        config.kv_b_proj_type = ggml_type.FP32
+        config.w_o_type = ggml_type.FP32
+    elif weight_type == torch.float16:
+        config.q_a_proj_type = ggml_type.FP16
+        config.q_b_proj_type = ggml_type.FP16
+        config.kv_a_proj_with_mqa_type = ggml_type.FP16
+        config.kv_b_proj_type = ggml_type.FP16
+        config.w_o_type = ggml_type.FP16
+    elif weight_type == torch.bfloat16:
+        config.q_a_proj_type = ggml_type.BF16
+        config.q_b_proj_type = ggml_type.BF16
+        config.kv_a_proj_with_mqa_type = ggml_type.BF16
+        config.kv_b_proj_type = ggml_type.BF16
+        config.w_o_type = ggml_type.BF16
+    else:
+        raise ValueError(f"Unsupported data type: {weight_type}")
+
+    config.pool = CPUInfer.backend_
+
+    if weight_type == torch.float32:
+        mla = kt_kernel_ext.mla.MLA_F32(config)
+    elif weight_type == torch.float16:
+        mla = kt_kernel_ext.mla.MLA_F16(config)
+    elif weight_type == torch.bfloat16:
+        # mla = kt_kernel_ext.mla.MLA_F32(config)
+        mla = kt_kernel_ext.mla.MLA_QUAN_F32(config)
+    else:
+        raise ValueError(f"Unsupported data type: {weight_type}")
+
+    mla.load_weights()
+    mla.set_local_pages(pages_count)
+
+    output = torch.zeros((qlen, hidden_size), dtype=input_type).to("cpu").contiguous()
+    mla.forward([qlen], [page_table], [kvlen], hidden_states.data_ptr(), output.data_ptr())
+    print("CPU MLA Output: ", output)
+    return output
+
+
+def load_fp16_tensor(file_path, shape):
+    # return load_fp32_tensor(file_path, shape)
+    return torch.zeros(shape)
+    with open(file_path, "rb") as f:
+        raw_data = f.read()
+    tensor = torch.frombuffer(raw_data, dtype=weight_type)
+    tensor = tensor.view(shape)  # 根据你的 shape reshape
+    return tensor
+
+
+def load_fp32_tensor(file_path, shape):
+    return torch.zeros(shape)
+    with open(file_path, "rb") as f:
+        raw_data = f.read()
+    tensor = torch.frombuffer(raw_data, dtype=torch.float32)
+    tensor = tensor.view(shape)  # 根据你的 shape reshape
+    return tensor
+
+
+def test_torch():
+    torch.set_grad_enabled(False)
+
+    softmax_scale = (nope_size + rope_size) ** -0.5
+    # 1代表的是压缩的kv的头数
+    k_caches = torch.randn(1, pages_count, page_size, 1, kv_lora_rank + rope_size).to(weight_type)
+    kv_cache = KDeepSeekV3Cache(page_size=page_size, kv_lora_rank=kv_lora_rank, k_caches=k_caches)
+
+    q_a_layernorm = DeepseekV2RMSNorm(q_lora_rank)
+    q_a_layernorm.weight = nn.Parameter(q_a_norm, requires_grad=False)
+
+    x = torch.randn(q_lora_rank, dtype=weight_type) * 100
+    print(x)
+    print(q_a_layernorm(x))
+
+    kv_a_layernorm = DeepseekV2RMSNorm(kv_lora_rank)
+    kv_a_layernorm.weight = nn.Parameter(kv_a_norm, requires_grad=False)
+
+    # 第三步：拆分成两个 tensor
+    # q_absorb, out_absorb = x_permuted[:, 0], x_permuted[:, 1]  # 都是 (num_heads, nope_size, kv_lora_rank
+    # q_absorb = kv_b_proj[:, ] # torch.randn(num_heads, nope_size, kv_lora_rank, dtype=data_type)
+    # out_absorb = kv_b_proj # torch.randn(num_heads, nope_size, kv_lora_rank, dtype=data_type)
+
+    rotary_emb = DeepseekV3YarnRotaryEmbedding(
+        rope_size,
+        max_position_embeddings=max_position_embeddings,
+        scaling_factor=rope_scaling["factor"],
+        base=rope_theta,
+        beta_fast=rope_scaling["beta_fast"],
+        beta_slow=rope_scaling["beta_slow"],
+        mscale=rope_scaling["mscale"],
+        mscale_all_dim=rope_scaling["mscale_all_dim"],
+        original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+    )
+    # 构造一个qlen 长度的输入 hidden_states, 对应的历史 kv_indptr 是[0:bsz]
+    # kv_indices 是[0:bsz]，page_idx=[0:bsz], page_offset=[kvlen:qlen+kvlen]
+    # last_page_len = [qlen+kvlen,...] layer_idx = 1
+    # position_ids = [kvlen:qlen+kvlen]
+    q_indptr = torch.tensor([0, qlen]).to(torch.int32)
+
+    kv_indptr = torch.tensor([0, (qlen + kvlen + page_size - 1) // page_size]).to(torch.int32)
+    kv_indices = torch.tensor(range(pages_count)).to(torch.int32)
+
+    page_idx = torch.tensor([i // page_size for i in range(kvlen, kvlen + qlen)]).to(torch.int32)
+    page_offset = torch.tensor([i % page_size for i in range(kvlen, kvlen + qlen)]).to(torch.int32)
+
+    last_page_len = torch.tensor([256], device=hidden_states.device)
+    position_ids = torch.tensor(range(kvlen, kvlen + qlen)).to(torch.int32)
+
+    # 按照行创建 mask [qlen,kvlen+qlen]
+    attention_masks = torch.zeros((max_qlen, max_kvlen), dtype=weight_type)
+    for i in range(max_qlen):
+        attention_masks[i, i + kvlen + 1 :] = -inf
+
+    def torch_attn(
+        hidden_states_i: torch.Tensor,
+        kv_cache: KDeepSeekV3Cache,
+        position_ids: torch.Tensor,
+        page_idx: torch.Tensor,
+        page_offset: torch.Tensor,
+        attention_masks: Optional[list[torch.Tensor]] = None,
+        q_indptr: Optional[torch.Tensor] = None,
+        kv_indices: Optional[torch.Tensor] = None,
+        kv_indptr: Optional[torch.Tensor] = None,
+        bsz_tensors: Optional[torch.Tensor] = None,
+        last_page_len: Optional[torch.Tensor] = None,
+        layer_idx: Optional[int] = None,
+    ):
+        global out_absorb
+        global q_absorb
+        hidden_states = hidden_states_i.to(weight_type)
+        # range bsz_tensors
+        final_attention_output = torch.tensor([], device=hidden_states.device)
+        for i in range(bsz_tensors[0]):
+            batch_num_tokens_tensors = q_indptr[i + 1] - q_indptr[i]
+            batch_last_page_len = last_page_len[i]
+            # kv_total_len is kv_len, batch_compressed_kv is compressed_kv, batch_k_pe is k_pe
+            batch_page_idx = page_idx[q_indptr[i] : q_indptr[i + 1]]
+            batch_page_offset = page_offset[q_indptr[i] : q_indptr[i + 1]]
+            # kv_page_nums is the number of pages for the current batch
+            kv_page_nums = kv_indptr[i + 1] - kv_indptr[i]
+            # kv_total_len is the total length of the kv cache for the current batch (kv_len for algorithm)
+            kv_total_len = kv_page_nums * page_size
+            if batch_last_page_len is not None:
+                kv_total_len = kv_total_len - (page_size - batch_last_page_len)
+            # print(f"kv_total_len's shape {kv_total_len.shape}")
+            # kv_index is the index of the kv cache pages for the current batch
+            kv_index = kv_indices[kv_indptr[i] : kv_indptr[i + 1]]
+            # we can index [kv_index, page_offset_indices] to get the kv cache for the current batch
+            # from q_indptr[i] to q_indptr[i+1] is the range of the current batch
+            batch_hidden_states = hidden_states[q_indptr[i] : q_indptr[i + 1]]
+            batch_position_ids = position_ids[q_indptr[i] : q_indptr[i + 1]]
+            qlen, _ = batch_hidden_states.size()
+            # print("qlen -> ", qlen)
+
+            hidden_states_to_check = load_fp16_tensor("./debug/query_0_tp_0_input.bin", batch_hidden_states.shape)
+            diff = torch.abs(batch_hidden_states - hidden_states_to_check).max()
+            print("hidden_states diff -> ", diff)
+
+            q_lora = q_a_proj(batch_hidden_states)
+            # q_lora_to_check = load_fp16_tensor('./debug/query_0_tp_0_qlora.bin', q_lora.shape)
+            # q_lora_to_check_test = load_fp16_tensor('./debug/query_0_tp_0_qlora_test.bin', q_lora.shape)
+            # diff = torch.abs(q_lora - q_lora_to_check).max()
+            # diff_test = torch.abs(q_lora - q_lora_to_check_test).max()
+            # print("q_lora max diff -> ", diff)
+            # print("q_lora max diff test -> ", diff_test)
+            # mae =  torch.mean(torch.abs(q_lora - q_lora_to_check))
+            # mae_test =  torch.mean(torch.abs(q_lora - q_lora_to_check_test))
+            # print("q_lora mae -> ", mae)
+            # print("q_lora mae test -> ", mae_test)
+
+            q_lora_norm = q_a_layernorm(q_lora)
+            # q_lora_norm_to_check = load_fp16_tensor('./debug/query_0_tp_0_qlora_norm.bin', q_lora_norm.shape)
+            # q_lora_norm_to_check_test = load_fp16_tensor('./debug/query_0_tp_0_qlora_norm_test.bin', q_lora_norm.shape)
+            # diff = torch.abs(q_lora_norm - q_lora_norm_to_check).max()
+            # mae =  torch.mean(torch.abs(q_lora_norm - q_lora_norm_to_check))
+            # diff_test = torch.abs(q_lora_norm - q_lora_norm_to_check_test).max()
+            # mae_test =  torch.mean(torch.abs(q_lora_norm - q_lora_norm_to_check_test))
+            # print("q_lora_norm diff -> ", diff)
+            # print("q_lora_norm mae -> ", mae)
+            # print("q_lora_norm diff test -> ", diff_test)
+            # print("q_lora_norm mae test -> ", mae_test)
+
+            q = q_b_proj(q_lora_norm)
+            # for v3, bsz, qlen, num_heads(128), qk_head_dim(192=128(nope)+64(rope))
+            q = q.view(qlen, num_heads, nope_size + rope_size)
+            # q_nope is [qlen, num_heads(128), qk_nope_head_dim(128)]
+            # q_pe is [qlen, num_heads(128), qk_rope_head_dim(64)]
+            q_nope, q_pe = torch.split(q, [nope_size, rope_size], dim=-1)
+
+            # compressed_kv is [qlen, kv_lora_rank(512) + rope(64)]
+            compressed_kv = kv_a_proj_with_mqa(batch_hidden_states)
+            # compressed_kv is [qlen, kv_lora_rank(512)], k_pe is [qlen, rope(64)]
+            compressed_kv, k_pe = torch.split(compressed_kv, [kv_lora_rank, rope_size], dim=-1)
+            compressed_kv = compressed_kv.contiguous()
+
+            # compressed_kv_page_0 = compressed_kv[0:page_size, :]
+            # compressed_kv_to_check = load_fp16_tensor('./debug/query_0_tp_0_page_0_kv_lora_rank',
+            #                                           compressed_kv_page_0.shape)
+            # diff = torch.abs(compressed_kv_page_0 - compressed_kv_to_check).max()
+            # mae =  torch.mean(torch.abs(compressed_kv_page_0 - compressed_kv_to_check))
+            # print("compressed_kv diff -> ", diff)
+            # print("compressed_kv mae -> ", mae)
+
+            compressed_kv = kv_a_layernorm(compressed_kv)
+            # k_pe is [qlen, 1, qk_rope_head_dim(64)]
+
+            # compressed_kv_page_0 = compressed_kv[0:page_size, :]
+            # compressed_kv_to_check = load_fp16_tensor('./debug/query_0_tp_0_page_0_kv_lora_rank_norm',
+            #                                           compressed_kv_page_0.shape)
+            # diff = torch.abs(compressed_kv_page_0 - compressed_kv_to_check).max()
+            # mae =  torch.mean(torch.abs(compressed_kv_page_0 - compressed_kv_to_check))
+            # print("compressed_kv diff norm -> ", diff)
+            # print("compressed_kv mae norm -> ", mae)
+
+            k_pe = k_pe.view(qlen, 1, rope_size)
+            # compressed_kv is [qlen, 1, kv_lora_rank(512)]
+            compressed_kv = compressed_kv.view(qlen, 1, kv_lora_rank)
+
+            cos, sin = rotary_emb(q_pe, batch_position_ids)
+
+            # q_nope_check = q_nope.transpose(0, 1) # qlen is 1, no GPU overhead, same below
+
+            # q_nope_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_nope', q_nope_check[0].shape)
+            # q_nope_0_to_check_test = load_fp16_tensor('./debug/query_0_tp_0_q_nope_test', q_nope_check[0].shape)
+            # diff = torch.abs(q_nope_check[0] - q_nope_0_to_check).max()
+            # mae =  torch.mean(torch.abs(q_nope_check[0] - q_nope_0_to_check))
+            # diff_test = torch.abs(q_nope_check[0] - q_nope_0_to_check_test).max()
+            # mae_test =  torch.mean(torch.abs(q_nope_check[0] - q_nope_0_to_check_test))
+            # print("q_nope[0] diff -> ", diff)
+            # print("q_nope[0] mae -> ", mae)
+            # print("q_nope[0] diff test -> ", diff_test)
+            # print("q_nope[0] mae test -> ", mae_test)
+
+            q_pe_nope = q_pe.transpose(0, 1)
+            # q_pe_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_rope', q_pe_nope[0].shape)
+            # q_pe_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_rope_no_rope', q_pe_nope[0].shape)
+            # q_pe_0_to_check_test = load_fp16_tensor('./debug/query_0_tp_0_q_rope_no_rope_test', q_pe_nope[0].shape)
+            # diff = torch.abs(q_pe_nope[0] - q_pe_0_to_check).max()
+            # mae =  torch.mean(torch.abs(q_pe_nope[0] - q_pe_0_to_check))
+            # diff_test = torch.abs(q_pe_nope[0] - q_pe_0_to_check_test).max()
+            # mae_test =  torch.mean(torch.abs(q_pe_nope[0] - q_pe_0_to_check_test))
+            # print("q_pe nope[0] diff -> ", diff)
+            # print("q_pe nope[0] mae -> ", mae)
+            # print("q_pe nope[0] diff test -> ", diff_test)
+            # print("q_pe nope[0] mae test -> ", mae_test)
+
+            # cos_to_check = load_fp32_tensor('./debug/query_0_tp_0_rope_cos', (qlen,32))
+            # diff = torch.abs(cos[:,:32]-cos_to_check).max()
+            # mae =  torch.mean(torch.abs(cos[:,:32]-cos_to_check))
+            # print("cos diff -> ", diff)
+            # print("cos mae -> ", mae)
+            # sin_to_check = load_fp32_tensor('./debug/query_0_tp_0_rope_sin', (qlen,32))
+            # diff = torch.abs(sin[:,:32]-sin_to_check).max()
+            # mae =  torch.mean(torch.abs(sin[:,:32]-sin_to_check))
+            # print("sin diff -> ", diff)
+            # print("sin mae -> ", mae)
+
+            # new_q_pe = q_pe.transpose(0, 1)
+            # qa = new_q_pe[:,:,range(0,64,2)]
+            # qb = new_q_pe[:,:,range(1,65,2)]
+            # # q1 = (qa * cos[:,:32] - qb * sin[:,:32])
+            # # q2 = (qb*cos[:,:32] + qa*sin[:,:32])
+            # q1 = (qa * cos_to_check - qb * sin_to_check)
+            # q2 = (qb*cos_to_check + qa*sin_to_check)
+            # q_new = torch.cat((q1,q2), dim=-1)
+            # print(f"q_pe shape{q_pe.shape}, k_pe shape {k_pe.shape}")
+            # new_q_pe = torch.zeros_like(q_pe)
+            # new_q_pe[:,:,range(0,64,2)] = 1
+            # new_q_pe[:,:,range(1,65,2)] = 10
+            q_pe, k_pe = apply_rotary_pos_emb(q_pe.unsqueeze(0), k_pe.unsqueeze(0), cos, sin, unsqueeze_dim=1)
+            q_pe = q_pe.squeeze(0)
+            # q_pe is [num_heads(128), qlen, qk_rope_head_dim(64)]
+            q_pe.transpose_(0, 1)
+
+            # diff = torch.abs(q_pe - q_new).max()
+            # print("q_pe diff -> ", diff)
+
+            # q_pe_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_rope', q_pe[0].shape)
+            # diff = torch.abs(q_pe[0] - q_pe_0_to_check).max()
+            # mae =  torch.mean(torch.abs(q_pe[0] - q_pe_0_to_check))
+            # print("q_pe[0] diff -> ", diff)
+            # print("q_pe[0] mae -> ", mae)
+
+            # diff = torch.abs(q_pe_0_to_check - q_new[0]).max()
+            # mae =  torch.mean(torch.abs(q_pe_0_to_check - q_new[0]))
+            # print("q_pe[0] 2  diff -> ", diff)
+            # print("q_pe[0] 2 mae -> ", mae)
+
+            if kv_cache is not None:
+                cache_kwargs = {
+                    "sin": sin,
+                    "cos": cos,
+                    "page_idx": batch_page_idx,
+                    "page_offset": batch_page_offset,
+                }  # Specific to RoPE models
+                compressed_kv_with_k_pe = kv_cache.update(
+                    compressed_kv.unsqueeze(0), k_pe, layer_idx, batch_page_idx, batch_page_offset, cache_kwargs
+                )
+                compressed_kv = compressed_kv_with_k_pe[:, :, :, :kv_lora_rank].view(-1, page_size, kv_lora_rank)
+                k_pe = compressed_kv_with_k_pe[:, :, :, kv_lora_rank:].view(-1, page_size, rope_size)
+            # q_absorb is [num_heads(128), qk_nope_head_dim(128), kv_lora_rank(512)]
+            # out_absorb is [num_heads(128), kv_lora_rank(512), v_head_dim(128)] v_head_dim is also the nope dim
+            # q_absorb, out_absorb = get_absorbed()
+            # q_nope is [num_heads(128), qlen, qk_nope_head_dim(128)]
+            q_nope = q_nope.transpose(0, 1)  # qlen is 1, no GPU overhead, same below
+
+            # q_nope_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_nope', q_nope[0].shape)
+            # diff = torch.abs(q_nope[0] - q_nope_0_to_check).max()
+            # mae =  torch.mean(torch.abs(q_nope[0] - q_nope_0_to_check))
+            # print("q_nope[0] diff -> ", diff)
+
+            # q_nope is [num_heads(128), qlen, kv_lora_rank(512)]
+            q_nope = torch.matmul(q_nope, q_absorb)  # batched MM
+
+            # k_b_proj_check = load_fp16_tensor('./debug/query_0_tp_0_k_b_lora', (nope_size,kv_lora_rank))
+            # diff = torch.abs(q_absorb[0] - k_b_proj_check).max()
+            # print("kv b lora weight[0] diff -> ", diff)
+
+            # q_absorb_check = load_fp16_tensor('./debug/query_0_tp_0_q_absorb', (kv_lora_rank,1024))
+            # q_absorb_check = q_absorb_check[:,0:qlen].transpose(0,1)
+            # diff = torch.abs(q_nope[0] - q_absorb_check).max()
+            # mae =  torch.mean(torch.abs(q_nope[0] - q_absorb_check))
+            # print("q_nope absorb diff -> ", diff)
+            # print("q_nope absorb mae -> ", mae)
+
+            # # q_nope is [qlen, num_heads(128), kv_lora_rank(512)]
+            # q_nope = q_nope.transpose(0, 1)
+
+            # we need to index out the compressed_kv and k_pe for the current batch
+            batch_compressed_kv = None
+            batch_k_pe = None
+            for page_index in kv_index:
+                if kv_total_len > page_size:
+                    tmp_compressed_kv = compressed_kv[page_index, 0:page_size, :]
+                    tmp_k_pe = k_pe[page_index, 0:page_size, :]
+                    if batch_compressed_kv is None or batch_k_pe is None:
+                        batch_compressed_kv = tmp_compressed_kv
+                        batch_k_pe = tmp_k_pe
+                    else:
+                        batch_compressed_kv = torch.cat((batch_compressed_kv, tmp_compressed_kv), dim=0)
+                        batch_k_pe = torch.cat((batch_k_pe, tmp_k_pe), dim=0)
+                    kv_total_len -= page_size
+                else:
+                    tmp_compressed_kv = compressed_kv[page_index, 0:kv_total_len, :]
+                    tmp_k_pe = k_pe[page_index, 0:kv_total_len, :]
+                    if batch_compressed_kv is None or batch_k_pe is None:
+                        batch_compressed_kv = tmp_compressed_kv
+                        batch_k_pe = tmp_k_pe
+                    else:
+                        batch_compressed_kv = torch.cat((batch_compressed_kv, tmp_compressed_kv), dim=0)
+                        batch_k_pe = torch.cat((batch_k_pe, tmp_k_pe), dim=0)
+                    break
+            # batch_compressed_kv is [kv_total_len(k_len), kv_lora_rank(512)]
+            # batch_k_pe is [kv_total_len(k_len), qk_rope_head_dim(64)]
+
+            # k_pe_to_check = load_fp16_tensor('./debug/query_0_tp_0_page_0_k_rope', (256,64))
+            # diff = torch.abs(batch_k_pe[:256] - k_pe_to_check).max()
+            # mae =  torch.mean(torch.abs(batch_k_pe[:256] - k_pe_to_check))
+            # print("k_pe diff -> ", diff)
+            # print("k_pe mae -> ", mae)
+
+            pe_weights = torch.matmul(q_pe, batch_k_pe.mT)
+            kv_total_len = kv_page_nums * page_size
+            # pe_weights_0 = load_fp16_tensor('./debug/query_0_tp_0_pe_attention_weights', (1024,4096))
+            # pe_weights_0 = pe_weights_0[0:qlen, 0:kv_total_len]
+            # diff = torch.abs(pe_weights[0] - pe_weights_0).max()
+            # print("pe_weights[0] diff -> ", diff)
+
+            attention_weights = pe_weights + torch.matmul(q_nope, batch_compressed_kv.mT)
+
+            # raw_weights = load_fp16_tensor('./debug/query_0_tp_0_raw_attention_weights', (1024, 4096))
+            # raw_weights = raw_weights[0:qlen, 0:kv_total_len]
+            # diff = torch.abs(attention_weights[0] - raw_weights).max()
+            # print("raw attention_weights[0] diff -> ", diff)
+
+            attention_weights = attention_weights * softmax_scale
+            # attention_weights is [num_heads(128), qlen, k_len]
+
+            # attention_weights = attention_weights.transpose(0,1).unsqueeze(0).squeeze(-1).expand(qlen,-1,-1).transpose(0,1)
+
+            # attention_masks[i] is [qlen, k_len]
+
+            print(attention_weights.shape)
+            print(attention_masks.shape)
+            attention_weights = (
+                attention_weights + attention_masks[: attention_weights.shape[1], : attention_weights.shape[2]]
+            )
+            # attention_weights shape is [num_heads(128), qlen, k_len]
+
+            attention_weights = nn.functional.softmax(attention_weights, dim=-1, dtype=weight_type).to(q_pe.dtype)
+
+            # attention_weights_0 = load_fp16_tensor('./debug/query_0_tp_0_attention_weights', (1024, 4096))
+            # attention_weights_0 = attention_weights_0[0:qlen, 0:kv_total_len]
+            # diff = torch.abs(attention_weights[0] - attention_weights_0).max()
+            # print("attention_weights[0] diff -> ", diff)
+
+            attn_output = torch.matmul(attention_weights, batch_compressed_kv)  # [num_heads(128),qlen, lora_rank(512)]
+            # out_absorb shape is [num_heads(128), kv_lora_rank(512), v_head_dim(128)]
+
+            # o_absorb_check = load_fp16_tensor('./debug/query_0_tp_0_o_absorb', (qlen,kv_lora_rank))
+            # diff = torch.abs(attn_output[0] - o_absorb_check).max()
+            # print("o absorb[0] diff -> ", diff)
+
+            out_absorb = out_absorb.transpose(1, 2)  # [qlen, num_heads(128), v_head_dim(128)]
+            # q for qlen, n for num_heads, h for v_head_dim, v for kv_lora_rank
+            attn_output = torch.matmul(attn_output, out_absorb)  # [num_heads(128), qlen, v_head_dim(128)]
+
+            # attn_output_check_0 = load_fp16_tensor('./debug/query_0_tp_0_attention_output', (qlen, nope_size))
+            # diff = torch.abs(attn_output[0] - attn_output_check_0).max()
+            # print("attn_output[0] diff -> ", diff)
+
+            attn_output = attn_output.transpose(0, 1)  # [qlen, num_heads(128), v_head_dim(128)]
+            attn_output = attn_output.reshape(qlen, num_heads * nope_size)
+
+            w_o = o_proj.weight.view([hidden_size, num_heads * nope_size])
+            output = torch.matmul(attn_output, w_o.transpose(0, 1))
+            output = output.view(qlen, hidden_size)
+
+            # output_0_check = load_fp16_tensor('./debug/query_0_tp_0_qlen_output', (qlen, hidden_size))
+            # h1_o = w_o[:,:128]
+            # local_o_check = load_fp16_tensor('./debug/query_0_tp_0_local_w_o', (hidden_size, 128))
+            # diff = torch.abs(local_o_check - h1_o).max()
+            # print("local w_o diff -> ", diff)
+
+            # h1_output = torch.matmul(attn_output[:,:128],h1_o.transpose(0,1))
+            # diff = torch.abs(h1_output - output_0_check).max()
+            # print("h1_output diff -> ", diff)
+
+            # output_check = load_fp16_tensor('./debug/output.bin', output.shape)
+            # diff = torch.abs(output - output_check).max()
+            # mae =   torch.mean(torch.abs(output - output_check))
+            # print("output diff -> ", diff)
+
+            final_attention_output = torch.cat((final_attention_output, output), dim=0)
+        return final_attention_output
+
+    torch_output = torch_attn(
+        hidden_states,
+        kv_cache,
+        position_ids,
+        page_idx,
+        page_offset,
+        attention_masks=attention_masks,
+        q_indptr=q_indptr,
+        kv_indices=kv_indices,
+        kv_indptr=kv_indptr,
+        bsz_tensors=bsz_tensors,
+        last_page_len=last_page_len,
+        layer_idx=0,
+    )
+    print("Torch Output: ", torch_output)
+    return torch_output
+
+
+torch.set_printoptions(sci_mode=False, precision=5)
+output_cpu = test_cpu_mla()
+output_torch = test_torch()
+print("Output CPU: ", output_cpu)
+print("Output Torch: ", output_torch)
+diff = (output_cpu - output_torch).abs()
+# 计算相对误差
+diff_relative = diff / (output_cpu.abs())
+# 把 diff_relative 中的 NaN 替换为 0
+diff_relative = torch.where(torch.isnan(diff_relative), torch.zeros_like(diff_relative), diff_relative)
+diff_relative_mean = torch.mean(torch.abs(output_cpu - output_torch)) / torch.mean(torch.abs(output_torch))
+
+print(
+    f"Diff: ave:{diff.mean()}, max:{diff.max()}, min:{diff.min()},  relative_mean:{diff_relative_mean}, relative_max:{diff_relative.max()}, relative_min:{diff_relative.min()}"
+)
+assert diff_relative_mean < 2e-1, "CPU and Torch outputs are not close enough!"

kt-kernel/examples/test_mla_qlen.py

@@ -0,0 +1,345 @@
+import logging
+import os, sys
+import time
+from typing import Optional
+
+os.environ["BLAS_NUM_THREADS"] = "1"
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+from kt_kernel import kt_kernel_ext
+from kt_kernel_ext.kvcache import ggml_type
+import torch
+from torch import inf, nn
+from torch.nn import init
+from torch_attention import apply_rotary_pos_emb, DeepseekV2RMSNorm, KDeepSeekV3Cache, DeepseekV3YarnRotaryEmbedding
+
+logger = logging.getLogger("reader")
+
+from gguf.gguf_reader import GGUFReader
+
+
+def read_gguf_file(gguf_file_path):
+    """
+    Reads and prints key-value pairs and tensor information from a GGUF file in an improved format.
+
+    Parameters:
+    - gguf_file_path: Path to the GGUF file.
+    """
+
+    reader = GGUFReader(gguf_file_path)
+
+    # List all key-value pairs in a columnized format
+    # print("Key-Value Pairs:") # noqa: NP100
+    # max_key_length = max(len(key) for key in reader.fields.keys())
+    for key, field in reader.fields.items():
+        value = field.parts[field.data[0]]
+        # print(f"{key:{max_key_length}} : {value}") # noqa: NP100
+    # print("----") # noqa: NP100
+
+    # List all tensors
+    # print("Tensors:") # noqa: NP100
+    # tensor_info_format = "{:<30} | Shape: {:<15} | Size: {:<12} | Quantization: {}"
+    # print(tensor_info_format.format("Tensor Name", "Shape", "Size", "Quantization")) # noqa: NP100
+    # print("-" * 80) # noqa: NP100
+    re = []
+    for tensor in reader.tensors:
+        shape_str = "x".join(map(str, tensor.shape))
+        size_str = str(tensor.n_elements)
+        quantization_str = tensor.tensor_type.name
+        # print(tensor_info_format.format(tensor.name, shape_str, size_str, quantization_str)) # noqa: NP100
+        re.append(tensor)
+    return re
+
+
+def get_torch_tensor_from_gguf(gguf_weights, name):
+    return torch.from_numpy(gguf_weights[name].data).contiguous()
+
+
+def get_torch_tensor_and_type_from_gguf(gguf_weights, name):
+    return torch.from_numpy(gguf_weights[name].data).contiguous(), gguf_weights[name].tensor_type.name
+
+
+def type_to_ggml_type(type):
+    if type == "F32":
+        return ggml_type.FP32
+    elif type == "F16":
+        return ggml_type.FP16
+    elif type == "BF16":
+        return ggml_type.BF16
+    else:
+        raise ValueError(f"Unsupported data type: {type}")
+
+
+use_real_weights = True
+gguf_path = "/home/bd/models/DeepSeek-R1-BF16"
+
+seed = 42  # 你可以选择任何整数作为种子
+torch.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+
+qlen = 1024
+kvlen = 0
+
+
+page_table = range(20)
+bsz_tensors = torch.tensor([1])
+
+
+page_size = 256
+pages_count = 200
+tp_count = 4
+
+
+hidden_size = 7168
+q_lora_rank = 1536
+kv_lora_rank = 512
+num_heads = 128
+nope_size = 128
+rope_size = 64
+
+rope_theta = 10000
+max_qlen = 1024
+max_kvlen = 4096
+
+max_position_embeddings = 163840
+
+
+rope_scaling = {
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "factor": 40,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0,
+    "original_max_position_embeddings": 4096,
+    "type": "yarn",
+}
+
+
+CPUInfer = kt_kernel_ext.CPUInfer(64)
+validation_iter = 100
+
+
+# data_type = torch.float32
+weight_type = torch.bfloat16
+# weight_type = torch.float16
+
+
+input_type = {
+    torch.float32: torch.float32,
+    torch.float16: torch.float16,
+    torch.bfloat16: torch.float32,
+}[weight_type]
+
+q_a_proj = nn.Linear(hidden_size, q_lora_rank, bias=False, dtype=weight_type)
+q_b_proj = nn.Linear(q_lora_rank, num_heads * (nope_size + rope_size), bias=False, dtype=weight_type)
+kv_a_proj_with_mqa = nn.Linear(hidden_size, kv_lora_rank + rope_size, bias=False, dtype=weight_type)
+kv_b_proj = nn.Linear(num_heads * (nope_size + nope_size), kv_lora_rank, bias=False, dtype=weight_type)
+o_proj = nn.Linear(num_heads * nope_size, hidden_size, bias=False, dtype=weight_type)
+q_a_norm = torch.ones(hidden_size, dtype=torch.float32)
+kv_a_norm = torch.ones(hidden_size, dtype=torch.float32)
+
+
+def read_gguf_directory(directory):
+    """
+    Reads all GGUF files in a directory and prints their contents.
+
+    Parameters:
+    - directory: Path to the directory containing GGUF files.
+    """
+    if not os.path.isdir(directory):
+        logger.error(f"Directory {directory} does not exist.")
+        return
+
+    # List all GGUF files in the directory
+    files = [f for f in os.listdir(directory) if f.endswith(".gguf")]
+    if not files:
+        logger.info(f"No GGUF files found in {directory}.")
+        return
+
+    re = []
+    for file in files:
+        file_path = os.path.join(directory, file)
+        # print(f"Reading {file_path}:") # noqa: NP100
+        # print("\n") # noqa: NP100
+        re.extend(read_gguf_file(file_path))
+    re = {r.name: r for r in re}
+    return re
+
+
+if use_real_weights := True:
+    gguf_weights = read_gguf_directory(gguf_path)
+    layer_idx = 0
+    q_a_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_a.weight")
+    q_a_proj.weight = nn.Parameter(q_a_proj_weight.view(torch.bfloat16), requires_grad=False)
+    q_a_type = type
+
+    q_a_norm_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_a_norm.weight")
+    q_a_norm = q_a_norm_weight.view(torch.float32)
+    # config.q_a_norm = q_a_norm_weight.data_ptr()
+    # config.q_a_norm_type = type_to_ggml_type(type)
+
+    q_b_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_q_b.weight")
+    q_b_proj.weight = nn.Parameter(q_b_proj_weight.view(torch.bfloat16), requires_grad=False)
+
+    kv_a_proj_with_mqa_weight, type = get_torch_tensor_and_type_from_gguf(
+        gguf_weights, f"blk.{layer_idx}.attn_kv_a_mqa.weight"
+    )
+    kv_a_proj_with_mqa.weight = nn.Parameter(kv_a_proj_with_mqa_weight.view(torch.bfloat16), requires_grad=False)
+
+    kv_a_norm_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_kv_a_norm.weight")
+    kv_a_norm = kv_a_norm_weight.view(torch.float32)
+    # config.kv_a_norm = kv_a_norm_weight.data_ptr()
+    # config.kv_a_norm_type = type_to_ggml_type(type)
+
+    kv_b_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_kv_b.weight")
+    kv_b_proj.weight = nn.Parameter(kv_b_proj_weight.view(torch.bfloat16), requires_grad=False)
+
+    o_proj_weight, type = get_torch_tensor_and_type_from_gguf(gguf_weights, f"blk.{layer_idx}.attn_output.weight")
+    o_proj.weight = nn.Parameter(o_proj_weight.view(torch.bfloat16), requires_grad=False)
+
+else:
+    init.normal_(q_a_proj.weight, mean=0.0, std=0.02)
+    init.normal_(q_b_proj.weight, mean=0.0, std=0.02)
+    init.normal_(kv_a_proj_with_mqa.weight, mean=0.0, std=0.02)
+    init.normal_(kv_b_proj.weight, mean=0.0, std=0.02)
+    init.normal_(o_proj.weight, mean=0.0, std=0.02)
+
+x_reshaped = kv_b_proj.weight.view(num_heads, 2, nope_size, kv_lora_rank)
+q_absorb = x_reshaped[:, 0]
+out_absorb = x_reshaped[:, 1]
+
+
+hidden_states = torch.randn((qlen, hidden_size), dtype=input_type).to("cpu").contiguous()
+
+
+def build_mla():
+    os.environ["BLAS_NUM_THREADS"] = "1"
+    q_a_proj_weight = q_a_proj.weight.to(weight_type).to("cpu").contiguous()
+    q_b_proj_weight = q_b_proj.weight.to(weight_type).to("cpu").contiguous()
+    kv_a_proj_with_mqa_weight = kv_a_proj_with_mqa.weight.to("cpu").to(weight_type).contiguous()
+    kv_b_proj_weight = kv_b_proj.weight.to(weight_type).to("cpu").contiguous()
+    o_proj_weight = o_proj.weight.to(weight_type).to("cpu").contiguous()
+
+    config = kt_kernel_ext.mla.MLAConfig(
+        hidden_size,
+        q_lora_rank,
+        kv_lora_rank,
+        num_heads,
+        nope_size,
+        rope_size,
+    )
+    config.max_qlen = max_qlen
+    config.max_kvlen = max_kvlen
+    config.max_position_embeddings = max_position_embeddings
+    config.rope_scaling_factor = rope_scaling["factor"]
+    config.rope_theta = rope_theta
+    config.rope_scaling_beta_fast = rope_scaling["beta_fast"]
+    config.rope_scaling_beta_slow = rope_scaling["beta_slow"]
+    config.rope_scaling_mscale = rope_scaling["mscale"]
+    config.rope_scaling_mscale_all_dim = rope_scaling["mscale_all_dim"]
+    config.rope_scaling_original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
+
+    config.q_a_proj = q_a_proj_weight.data_ptr()
+    config.q_b_proj = q_b_proj_weight.data_ptr()
+    config.kv_a_proj_with_mqa = kv_a_proj_with_mqa_weight.data_ptr()
+    config.kv_b_proj = kv_b_proj_weight.data_ptr()
+    config.o_proj = o_proj_weight.data_ptr()
+
+    config.q_a_norm = q_a_norm.data_ptr()
+    config.q_a_norm_type = ggml_type.FP32
+    config.kv_a_norm = kv_a_norm.data_ptr()
+    config.kv_a_norm_type = ggml_type.FP32
+
+    if weight_type == torch.float32:
+        config.q_a_proj_type = ggml_type.FP32
+        config.q_b_proj_type = ggml_type.FP32
+        config.kv_a_proj_with_mqa_type = ggml_type.FP32
+        config.kv_b_proj_type = ggml_type.FP32
+        config.w_o_type = ggml_type.FP32
+    elif weight_type == torch.float16:
+        config.q_a_proj_type = ggml_type.FP16
+        config.q_b_proj_type = ggml_type.FP16
+        config.kv_a_proj_with_mqa_type = ggml_type.FP16
+        config.kv_b_proj_type = ggml_type.FP16
+        config.w_o_type = ggml_type.FP16
+    elif weight_type == torch.bfloat16:
+        config.q_a_proj_type = ggml_type.BF16
+        config.q_b_proj_type = ggml_type.BF16
+        config.kv_a_proj_with_mqa_type = ggml_type.BF16
+        config.kv_b_proj_type = ggml_type.BF16
+        config.w_o_type = ggml_type.BF16
+    else:
+        raise ValueError(f"Unsupported data type: {weight_type}")
+
+    config.pool = CPUInfer.backend_
+
+    if weight_type == torch.float32:
+        mla = kt_kernel_ext.mla.MLA_F32(config)
+    elif weight_type == torch.float16:
+        mla = kt_kernel_ext.mla.MLA_F16(config)
+    elif weight_type == torch.bfloat16:
+        mla = kt_kernel_ext.mla.MLA_F32(config)
+    else:
+        raise ValueError(f"Unsupported data type: {weight_type}")
+
+    mla.load_weights()
+    mla.set_local_pages(pages_count)
+    return mla
+
+
+def load_fp32_tensor(file_path, shape):
+    with open(file_path, "rb") as f:
+        raw_data = f.read()
+    tensor = torch.frombuffer(raw_data, dtype=torch.float32)
+    tensor = tensor.view(shape)  # 根据你的 shape reshape
+    return tensor
+
+
+# page3 = load_fp32_tensor('/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug1/query_0_tp_0_page_3_kv_lora_rank_norm.f32',(page_size,kv_lora_rank))
+# page3_2 = load_fp32_tensor('/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug2/query_0_tp_0_page_3_kv_lora_rank_norm.f32',(page_size,kv_lora_rank))
+
+# diff = torch.abs(page3 - page3_2)
+# print(f'Diff: ave:{diff.mean()}, max:{diff.max()}')
+
+# q_pe_1 = load_fp32_tensor('/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug1/query_0_tp_0_q_rope.f32',(1, rope_size))
+# q_pe_2 = load_fp32_tensor('/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug2/query_0_tp_0_q_rope.f32',(qlen, rope_size))
+# diff = torch.abs(q_pe_1 - q_pe_2[-1])
+# print(f'Q PE Diff: ave:{diff.mean()}, max:{diff.max()}')
+
+# q_nope_1 = load_fp32_tensor('/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug1/query_0_tp_0_q_nope.f32',(1, nope_size))
+# q_nope_2 = load_fp32_tensor('/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug2/query_0_tp_0_q_nope.f32',(qlen, nope_size))
+# diff = torch.abs(q_nope_1 - q_nope_2[-1])
+# print(f'Q Nope Diff: ave:{diff.mean()}, max:{diff.max()}')
+
+
+# pe_attn_w_1 = load_fp32_tensor('/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug1/query_0_tp_0_pe_attention_weights.f32',(1,max_kvlen))
+# pe_attn_w_2 = load_fp32_tensor('/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug2/query_0_tp_0_pe_attention_weights.f32',(qlen,max_kvlen))
+# diff = torch.abs(pe_attn_w_1 - pe_attn_w_2[-1])
+# print(f'PE Attention Weights Diff: ave:{diff.mean()}, max:{diff.max()}')
+
+
+# raw_attn_w_1 = load_fp32_tensor('/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug1/query_0_tp_0_raw_attention_weights.f32',(1,max_kvlen))
+# raw_attn_w_2 = load_fp32_tensor('/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug2/query_0_tp_0_raw_attention_weights.f32',(qlen,max_kvlen))
+# diff = torch.abs(raw_attn_w_1 - raw_attn_w_2[-1])
+# print(f'Raw Attention Weights Diff: ave:{diff.mean()}, max:{diff.max()}')
+
+
+# output_1 = load_fp32_tensor('/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug1/output.bin.f32',shape=(1, hidden_size))
+# output_2 = load_fp32_tensor('/home/yzw/xwy/Projects/ktransformers-dev/csrc/ktransformers_ext/examples/debug2/output.bin.f32',shape=(qlen, hidden_size))
+
+# diff = torch.abs(output_1 - output_2[-1])
+# print(f'Output Diff: ave:{diff.mean()}, max:{diff.max()}')
+
+
+mla = build_mla()
+output = torch.zeros((qlen, hidden_size), dtype=input_type).to("cpu").contiguous()
+mla.forward([qlen], [page_table], [kvlen], hidden_states.data_ptr(), output.data_ptr())
+print("CPU MLA Output: ", output[-1])
+
+
+output_2 = torch.zeros((1, hidden_size), dtype=input_type).to("cpu").contiguous()
+mla.forward([1], [page_table], [qlen - 1], hidden_states[-1].data_ptr(), output_2.data_ptr())
+print("CPU MLA Output 2: ", output_2[-1])
+
+diff = torch.abs(output[-1] - output_2[-1])
+print(f"Diff: ave:{diff.mean()}, max:{diff.max()}")
+assert diff.max() < 1e-1, "CPU and Torch outputs are not close enough!"

kt-kernel/examples/test_mla_quant.py

@@ -0,0 +1,353 @@
+import logging
+import os, sys
+import time
+from typing import Optional
+
+os.environ["BLAS_NUM_THREADS"] = "1"
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+from kt_kernel import kt_kernel_ext
+from kt_kernel_ext.kvcache import ggml_type
+import torch
+from torch import inf, nn
+from torch.nn import init
+from torch_attention import apply_rotary_pos_emb, DeepseekV2RMSNorm, KDeepSeekV3Cache, DeepseekV3YarnRotaryEmbedding
+
+logger = logging.getLogger("reader")
+
+from gguf.gguf_reader import GGUFReader
+
+
+def load_fp32_tensor_raw(file_path):
+    # return torch.zeros(shape)
+    with open(file_path, "rb") as f:
+        raw_data = f.read()
+    tensor = torch.frombuffer(raw_data, dtype=torch.float32)
+    return tensor
+
+
+def load_fp16_tensor(file_path, shape=None):
+    # return load_fp32_tensor(file_path, shape)
+    return load_fp32_tensor_raw(file_path)
+    # return torch.zeros(shape)
+    with open(file_path, "rb") as f:
+        raw_data = f.read()
+    tensor = torch.frombuffer(raw_data, dtype=weight_type)
+    tensor = tensor.view(shape)  # 根据你的 shape reshape
+    return tensor
+
+
+def load_fp32_tensor(file_path, shape):
+    # return torch.zeros(shape)
+    with open(file_path, "rb") as f:
+        raw_data = f.read()
+    tensor = torch.frombuffer(raw_data, dtype=torch.float32)
+    tensor = tensor.view(shape)  # 根据你的 shape reshape
+    return tensor
+
+
+def test_torch():
+    torch.set_grad_enabled(False)
+
+    hidden_states_to_check_decode = load_fp16_tensor("./debug_decode/query_0_tp_0_input.bin")
+    hidden_states_to_check_prefill = load_fp16_tensor("./debug_prefill/query_0_tp_0_input.bin")
+    # diff = torch.abs(hidden_states_to_check_prefill - hidden_states_to_check_decode).max()
+    # print("hidden_states diff -> ", diff)
+
+    q_lora_to_check_decode = load_fp16_tensor("./debug_decode/query_0_tp_0_qlora.bin")
+    q_lora_to_check_test_decode = load_fp16_tensor("./debug_decode/query_0_tp_0_qlora_test.bin")
+    q_lora_to_check_prefill = load_fp16_tensor("./debug_prefill/query_0_tp_0_qlora.bin")
+    q_lora_to_check_test_prefill = load_fp16_tensor("./debug_prefill/query_0_tp_0_qlora_test.bin")
+    # diff = torch.abs(q_lora_to_check_prefill - q_lora_to_check_decode).max()
+    # diff_test = torch.abs(q_lora_to_check_prefill - q_lora_to_check_decode).max()
+    # print("q_lora max diff -> ", diff)
+    # print("q_lora max diff test -> ", diff_test)
+    # mae =  torch.mean(torch.abs(q_lora_to_check_prefill - q_lora_to_check_decode))
+    # mae_test =  torch.mean(torch.abs(q_lora_to_check_prefill - q_lora_to_check_decode))
+    # print("q_lora mae -> ", mae)
+    # print("q_lora mae test -> ", mae_test)
+
+    # q_lora_norm = q_a_layernorm(q_lora)
+    # q_lora_norm_to_check = load_fp16_tensor('./debug/query_0_tp_0_qlora_norm.bin', q_lora_norm.shape)
+    # q_lora_norm_to_check_test = load_fp16_tensor('./debug/query_0_tp_0_qlora_norm_test.bin', q_lora_norm.shape)
+    # diff = torch.abs(q_lora_norm - q_lora_norm_to_check).max()
+    # mae =  torch.mean(torch.abs(q_lora_norm - q_lora_norm_to_check))
+    # diff_test = torch.abs(q_lora_norm - q_lora_norm_to_check_test).max()
+    # mae_test =  torch.mean(torch.abs(q_lora_norm - q_lora_norm_to_check_test))
+    # print("q_lora_norm diff -> ", diff)
+    # print("q_lora_norm mae -> ", mae)
+    # print("q_lora_norm diff test -> ", diff_test)
+    # print("q_lora_norm mae test -> ", mae_test)
+
+    # q = q_b_proj(q_lora_norm)
+    # for v3, bsz, qlen, num_heads(128), qk_head_dim(192=128(nope)+64(rope))
+    # q = q.view(qlen, num_heads, nope_size+rope_size)
+    # q_nope is [qlen, num_heads(128), qk_nope_head_dim(128)]
+    # q_pe is [qlen, num_heads(128), qk_rope_head_dim(64)]
+    # q_nope, q_pe = torch.split(
+    #     q, [nope_size, rope_size], dim=-1
+    # )
+
+    # compressed_kv is [qlen, kv_lora_rank(512) + rope(64)]
+    # compressed_kv = kv_a_proj_with_mqa(batch_hidden_states)
+    # compressed_kv is [qlen, kv_lora_rank(512)], k_pe is [qlen, rope(64)]
+    # compressed_kv, k_pe = torch.split(
+    #     compressed_kv, [kv_lora_rank, rope_size], dim=-1
+    # )
+    # compressed_kv = compressed_kv.contiguous()
+
+    # compressed_kv_page_0 = compressed_kv[0:page_size, :]
+    compressed_kv_to_check_decode = load_fp16_tensor("./debug_decode/query_0_tp_0_page_0_kv_lora_rank")
+    compressed_kv_to_check_prefill = load_fp16_tensor("./debug_prefill/query_0_tp_0_page_0_kv_lora_rank")
+    # diff = torch.abs(compressed_kv_to_check_prefill - compressed_kv_to_check_decode).max()
+    # mae =  torch.mean(torch.abs(compressed_kv_to_check_prefill - compressed_kv_to_check_decode))
+    # print("compressed_kv diff -> ", diff)
+    # print("compressed_kv mae -> ", mae)
+
+    # compressed_kv = kv_a_layernorm(compressed_kv)
+    # k_pe is [qlen, 1, qk_rope_head_dim(64)]
+
+    # compressed_kv_page_0 = compressed_kv[0:page_size, :]
+    compressed_kv_to_check_decode = load_fp16_tensor("./debug_decode/query_0_tp_0_page_0_kv_lora_rank_norm")
+    compressed_kv_to_check_prefill = load_fp16_tensor("./debug_prefill/query_0_tp_0_page_0_kv_lora_rank_norm")
+    # diff = torch.abs(compressed_kv_page_0 - compressed_kv_to_check).max()
+    # mae =  torch.mean(torch.abs(compressed_kv_page_0 - compressed_kv_to_check))
+    # print("compressed_kv diff norm -> ", diff)
+    # print("compressed_kv mae norm -> ", mae)
+
+    # k_pe = k_pe.view(qlen, 1, rope_size)
+    # compressed_kv is [qlen, 1, kv_lora_rank(512)]
+    # compressed_kv = compressed_kv.view(qlen, 1, kv_lora_rank)
+
+    # cos, sin = rotary_emb(q_pe, batch_position_ids)
+
+    # q_nope_check = q_nope.transpose(0, 1) # qlen is 1, no GPU overhead, same below
+
+    # q_nope_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_nope', q_nope_check[0].shape)
+    # q_nope_0_to_check_test = load_fp16_tensor('./debug/query_0_tp_0_q_nope_test', q_nope_check[0].shape)
+    # diff = torch.abs(q_nope_check[0] - q_nope_0_to_check).max()
+    # mae =  torch.mean(torch.abs(q_nope_check[0] - q_nope_0_to_check))
+    # diff_test = torch.abs(q_nope_check[0] - q_nope_0_to_check_test).max()
+    # mae_test =  torch.mean(torch.abs(q_nope_check[0] - q_nope_0_to_check_test))
+    # print("q_nope[0] diff -> ", diff)
+    # print("q_nope[0] mae -> ", mae)
+    # print("q_nope[0] diff test -> ", diff_test)
+    # print("q_nope[0] mae test -> ", mae_test)
+
+    # q_pe_nope = q_pe.transpose(0,1)
+    q_pe_0_to_check_decode = load_fp16_tensor("./debug_decode/query_0_tp_0_q_rope")
+    q_pe_0_to_check_prefill = load_fp16_tensor("./debug_prefill/query_0_tp_0_q_rope")
+
+    # q_pe_0_to_check_decode_test = load_fp16_tensor('./debug_decode/query_0_tp_0_q_rope_test')
+    # q_pe_0_to_check_prefill_test = load_fp16_tensor('./debug_prefill/query_0_tp_0_q_rope_test')
+
+    # q_pe_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_rope_no_rope', q_pe_nope[0].shape)
+    # q_pe_0_to_check_test = load_fp16_tensor('./debug/query_0_tp_0_q_rope_no_rope_test', q_pe_nope[0].shape)
+    # diff = torch.abs(q_pe_nope[0] - q_pe_0_to_check).max()
+    # mae =  torch.mean(torch.abs(q_pe_nope[0] - q_pe_0_to_check))
+    # diff_test = torch.abs(q_pe_nope[0] - q_pe_0_to_check_test).max()
+    # mae_test =  torch.mean(torch.abs(q_pe_nope[0] - q_pe_0_to_check_test))
+    # print("q_pe nope[0] diff -> ", diff)
+    # print("q_pe nope[0] mae -> ", mae)
+    # print("q_pe nope[0] diff test -> ", diff_test)
+    # print("q_pe nope[0] mae test -> ", mae_test)
+
+    # cos_to_check = load_fp32_tensor('./debug/query_0_tp_0_rope_cos', (qlen,32))
+    # diff = torch.abs(cos[:,:32]-cos_to_check).max()
+    # mae =  torch.mean(torch.abs(cos[:,:32]-cos_to_check))
+    # print("cos diff -> ", diff)
+    # print("cos mae -> ", mae)
+    # sin_to_check = load_fp32_tensor('./debug/query_0_tp_0_rope_sin', (qlen,32))
+    # diff = torch.abs(sin[:,:32]-sin_to_check).max()
+    # mae =  torch.mean(torch.abs(sin[:,:32]-sin_to_check))
+    # print("sin diff -> ", diff)
+    # print("sin mae -> ", mae)
+
+    # new_q_pe = q_pe.transpose(0, 1)
+    # qa = new_q_pe[:,:,range(0,64,2)]
+    # qb = new_q_pe[:,:,range(1,65,2)]
+    # q1 = (qa * cos[:,:32] - qb * sin[:,:32])
+    # q2 = (qb*cos[:,:32] + qa*sin[:,:32])
+    # q1 = (qa * cos_to_check - qb * sin_to_check)
+    # q2 = (qb*cos_to_check + qa*sin_to_check)
+    # q_new = torch.cat((q1,q2), dim=-1)
+    # print(f"q_pe shape{q_pe.shape}, k_pe shape {k_pe.shape}")
+    # new_q_pe = torch.zeros_like(q_pe)
+    # new_q_pe[:,:,range(0,64,2)] = 1
+    # new_q_pe[:,:,range(1,65,2)] = 10
+    # q_pe, k_pe = apply_rotary_pos_emb(q_pe.unsqueeze(0), k_pe.unsqueeze(0), cos, sin, unsqueeze_dim=1)
+    # q_pe = q_pe.squeeze(0)
+    # q_pe is [num_heads(128), qlen, qk_rope_head_dim(64)]
+    # q_pe.transpose_(0, 1)
+
+    # diff = torch.abs(q_pe - q_new).max()
+    # print("q_pe diff -> ", diff)
+
+    # q_pe_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_rope', q_pe[0].shape)
+    # diff = torch.abs(q_pe[0] - q_pe_0_to_check).max()
+    # mae =  torch.mean(torch.abs(q_pe[0] - q_pe_0_to_check))
+    # print("q_pe[0] diff -> ", diff)
+    # print("q_pe[0] mae -> ", mae)
+
+    # diff = torch.abs(q_pe_0_to_check - q_new[0]).max()
+    # mae =  torch.mean(torch.abs(q_pe_0_to_check - q_new[0]))
+    # print("q_pe[0] 2  diff -> ", diff)
+    # print("q_pe[0] 2 mae -> ", mae)
+
+    # if kv_cache is not None:
+    #     cache_kwargs = {"sin": sin, "cos": cos, "page_idx": batch_page_idx, "page_offset": batch_page_offset}  # Specific to RoPE models
+    #     compressed_kv_with_k_pe = kv_cache.update(compressed_kv.unsqueeze(0), k_pe, layer_idx, batch_page_idx, batch_page_offset, cache_kwargs)
+    #     compressed_kv = compressed_kv_with_k_pe [:, :, :, :kv_lora_rank].view(-1, page_size, kv_lora_rank)
+    #     k_pe = compressed_kv_with_k_pe [:, :, :, kv_lora_rank:].view(-1, page_size, rope_size)
+    # # q_absorb is [num_heads(128), qk_nope_head_dim(128), kv_lora_rank(512)]
+    # # out_absorb is [num_heads(128), kv_lora_rank(512), v_head_dim(128)] v_head_dim is also the nope dim
+    # # q_absorb, out_absorb = get_absorbed()
+    # # q_nope is [num_heads(128), qlen, qk_nope_head_dim(128)]
+    # q_nope = q_nope.transpose(0, 1) # qlen is 1, no GPU overhead, same below
+
+    # q_nope_0_to_check = load_fp16_tensor('./debug/query_0_tp_0_q_nope', q_nope[0].shape)
+    # diff = torch.abs(q_nope[0] - q_nope_0_to_check).max()
+    # mae =  torch.mean(torch.abs(q_nope[0] - q_nope_0_to_check))
+    # print("q_nope[0] diff -> ", diff)
+
+    # # q_nope is [num_heads(128), qlen, kv_lora_rank(512)]
+    # q_nope = torch.matmul(q_nope, q_absorb) # batched MM
+
+    # k_b_proj_check = load_fp16_tensor('./debug/query_0_tp_0_k_b_lora', (nope_size,kv_lora_rank))
+    # diff = torch.abs(q_absorb[0] - k_b_proj_check).max()
+    # print("kv b lora weight[0] diff -> ", diff)
+
+    # q_absorb_check = load_fp16_tensor('./debug/query_0_tp_0_q_absorb', (kv_lora_rank,1024))
+    # q_absorb_check = q_absorb_check[:,0:qlen].transpose(0,1)
+    # diff = torch.abs(q_nope[0] - q_absorb_check).max()
+    # mae =  torch.mean(torch.abs(q_nope[0] - q_absorb_check))
+    # print("q_nope absorb diff -> ", diff)
+    # print("q_nope absorb mae -> ", mae)
+
+    # # q_nope is [qlen, num_heads(128), kv_lora_rank(512)]
+    # q_nope = q_nope.transpose(0, 1)
+
+    # we need to index out the compressed_kv and k_pe for the current batch
+    # batch_compressed_kv = None
+    # batch_k_pe = None
+    # for page_index in kv_index:
+    #     if kv_total_len > page_size:
+    #         tmp_compressed_kv = compressed_kv[page_index, 0:page_size, :]
+    #         tmp_k_pe = k_pe[page_index, 0:page_size, :]
+    #         if batch_compressed_kv is None or batch_k_pe is None:
+    #             batch_compressed_kv = tmp_compressed_kv
+    #             batch_k_pe = tmp_k_pe
+    #         else:
+    #             batch_compressed_kv = torch.cat((batch_compressed_kv, tmp_compressed_kv), dim=0)
+    #             batch_k_pe = torch.cat((batch_k_pe, tmp_k_pe), dim=0)
+    #         kv_total_len -= page_size
+    #     else:
+    #         tmp_compressed_kv = compressed_kv[page_index, 0:kv_total_len, :]
+    #         tmp_k_pe = k_pe[page_index, 0:kv_total_len, :]
+    #         if batch_compressed_kv is None or batch_k_pe is None:
+    #             batch_compressed_kv = tmp_compressed_kv
+    #             batch_k_pe = tmp_k_pe
+    #         else:
+    #             batch_compressed_kv = torch.cat((batch_compressed_kv, tmp_compressed_kv), dim=0)
+    #             batch_k_pe = torch.cat((batch_k_pe, tmp_k_pe), dim=0)
+    #         break
+    # batch_compressed_kv is [kv_total_len(k_len), kv_lora_rank(512)]
+    # batch_k_pe is [kv_total_len(k_len), qk_rope_head_dim(64)]
+
+    k_pe_to_check_decode = load_fp16_tensor("./debug_decode/query_0_tp_0_page_0_k_rope", (256, 64))
+    k_pe_to_check_prefill = load_fp16_tensor("./debug_prefill/query_0_tp_0_page_0_k_rope", (256, 64))
+    # diff = torch.abs(k_pe_to_check_prefill - k_pe_to_check_decode).max()
+    # mae =  torch.mean(k_pe_to_check_prefill - k_pe_to_check_decode)
+    # print("k_pe diff -> ", diff)
+    # print("k_pe mae -> ", mae)
+
+    # pe_weights = torch.matmul(q_pe,batch_k_pe.mT)
+    # kv_total_len = kv_page_nums * page_size
+    pe_weights_0_decode = load_fp16_tensor("./debug_decode/query_0_tp_0_pe_attention_weights", (1024, 4096))
+    pe_weights_0_prefill = load_fp16_tensor("./debug_prefill/query_0_tp_0_pe_attention_weights", (1024, 4096))
+
+    # diff = torch.abs(pe_weights[0] - pe_weights_0).max()
+    # print("pe_weights[0] diff -> ", diff)
+
+    # attention_weights = (pe_weights + torch.matmul(q_nope, batch_compressed_kv.mT))
+
+    # raw_weights = load_fp16_tensor('./debug/query_0_tp_0_raw_attention_weights', (1024, 4096))
+    # raw_weights = raw_weights[0:qlen, 0:kv_total_len]
+    # diff = torch.abs(attention_weights[0] - raw_weights).max()
+    # print("raw attention_weigh/ts[0] diff -> ", diff)
+
+    # attention_weights = attention_weights * softmax_scale
+    # attention_weights is [num_heads(128), qlen, k_len]
+
+    # attention_weights = attention_weights.transpose(0,1).unsqueeze(0).squeeze(-1).expand(qlen,-1,-1).transpose(0,1)
+
+    # attention_masks[i] is [qlen, k_len]
+
+    # attention_weights = (attention_weights + attention_masks)
+    # attention_weights shape is [num_heads(128), qlen, k_len]
+
+    # attention_weights = nn.functional.softmax(attention_weights,dim=-1,dtype=weight_type).to(q_pe.dtype)
+
+    attention_weights_0_decode = load_fp16_tensor("./debug_decode/query_0_tp_0_attention_weights", (1024, 4096))
+    attention_weights_0_prefill = load_fp16_tensor("./debug_prefill/query_0_tp_0_attention_weights", (1024, 4096))
+
+    # attention_weights_0 = attention_weights_0[0:qlen, 0:kv_total_len]
+    # diff = torch.abs(attention_weights[0] - attention_weights_0).max()
+    # print("attention_weights[0] diff -> ", diff)
+
+    # attn_output = torch.matmul(attention_weights, batch_compressed_kv) # [num_heads(128),qlen, lora_rank(512)]
+    # out_absorb shape is [num_heads(128), kv_lora_rank(512), v_head_dim(128)]
+
+    # o_absorb_check = load_fp16_tensor('./debug/query_0_tp_0_o_absorb', (qlen,kv_lora_rank))
+    # diff = torch.abs(attn_output[0] - o_absorb_check).max()
+    # print("o absorb[0] diff -> ", diff)
+
+    # out_absorb = out_absorb.transpose(1, 2) # [qlen, num_heads(128), v_head_dim(128)]
+    # # q for qlen, n for num_heads, h for v_head_dim, v for kv_lora_rank
+    # attn_output = torch.matmul(attn_output, out_absorb) # [num_heads(128), qlen, v_head_dim(128)]
+
+    # attn_output_check_0 = load_fp16_tensor('./debug/query_0_tp_0_attention_output', (qlen, nope_size))
+    # diff = torch.abs(attn_output[0] - attn_output_check_0).max()
+    # print("attn_output[0] diff -> ", diff)
+
+    # attn_output = attn_output.transpose(0, 1) # [qlen, num_heads(128), v_head_dim(128)]
+    # attn_output = attn_output.reshape(qlen, num_heads * nope_size)
+
+    # w_o = o_proj.weight.view([hidden_size,num_heads * nope_size])
+    # output = torch.matmul(attn_output,w_o.transpose(0,1))
+    # output = output.view(qlen, hidden_size)
+
+    # output_0_check = load_fp16_tensor('./debug/query_0_tp_0_qlen_output', (qlen, hidden_size))
+    # h1_o = w_o[:,:128]
+    # local_o_check = load_fp16_tensor('./debug/query_0_tp_0_local_w_o', (hidden_size, 128))
+    # diff = torch.abs(local_o_check - h1_o).max()
+    # print("local w_o diff -> ", diff)
+
+    # h1_output = torch.matmul(attn_output[:,:128],h1_o.transpose(0,1))
+    # diff = torch.abs(h1_output - output_0_check).max()
+    # print("h1_output diff -> ", diff)
+
+    output_check_decode = load_fp16_tensor("./debug_decode/output.bin")
+    output_check_prefill = load_fp16_tensor("./debug_prefill/output.bin")
+    # diff = torch.abs(output - output_check).max()
+    # mae =   torch.mean(torch.abs(output - output_check))
+    # print("output diff -> ", diff)
+
+    return None
+
+
+torch.set_printoptions(sci_mode=False, precision=5)
+# output_cpu = test_cpu_mla()
+# output_cpu_quant = test_cpu_mla_quant()
+output_torch = test_torch()
+# print("Output CPU: ", output_cpu)
+# print("Output CPU: ", output_cpu_quant)
+# print("Output Torch: ", output_torch)
+# diff = (output_cpu - output_torch).abs()
+# # 计算相对误差
+# diff_relative = diff / (output_cpu.abs())
+# # 把 diff_relative 中的 NaN 替换为 0
+# diff_relative = torch.where(torch.isnan(diff_relative), torch.zeros_like(diff_relative), diff_relative)
+# diff_relative_mean = torch.mean(torch.abs(output_cpu-output_torch)) / torch.mean(torch.abs(output_torch))
+
+# print(f'Diff: ave:{diff.mean()}, max:{diff.max()}, min:{diff.min()},  relative_mean:{diff_relative_mean}, relative_max:{diff_relative.max()}, relative_min:{diff_relative.min()}')
+# assert diff_relative_mean < 2e-1, "CPU and Torch outputs are not close enough!"

kt-kernel/examples/test_mla_simple.py

@@ -0,0 +1,305 @@
+import math
+import random
+import os, sys
+import time
+import subprocess
+import platform
+import json
+from typing import Any, Dict, Optional, Tuple
+import numpy as np
+import torch.nn.init as init
+from torch_attention import apply_rotary_pos_emb,DeepseekV2RMSNorm,KDeepSeekV3Cache,DeepseekV3YarnRotaryEmbedding
+
+import torch
+from tqdm import tqdm
+from torch import nn
+"""
+"rope_scaling": {
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "factor": 40,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0,
+    "original_max_position_embeddings": 4096,
+    "type": "yarn"
+  },
+""" 
+
+rope_scaling = {
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "factor": 40,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0,
+    "original_max_position_embeddings": 4096,
+    "type": "yarn"
+}
+seed = 42  # 你可以选择任何整数作为种子
+torch.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+np.random.seed(seed)
+random.seed(seed)
+
+# "rope_theta": 10000
+rope_theta = 10000
+
+
+hidden_size = 7168
+num_heads = 128
+kv_lora_rank = 512
+q_lora_rank = 512
+nope_size = 128
+rope_size = 64
+
+# page 的个数
+page_nums = 10
+page_size = 512
+layer_num = 10
+max_position_embeddings =  163840
+
+
+warm_up_iter = 1000
+test_iter = 1000
+
+q_len = 200
+his_kv_len = 128
+
+bsz_tensors=torch.tensor([1])
+
+softmax_scale = (nope_size + rope_size) ** -0.5
+# 1代表的是压缩的kv的头数
+k_caches = torch.randn(layer_num,page_nums, page_size,1, kv_lora_rank + rope_size).to(torch.float16)
+kv_cache = KDeepSeekV3Cache(page_size=page_size, kv_lora_rank=kv_lora_rank, k_caches=k_caches)
+
+q_a_layernorm = DeepseekV2RMSNorm(q_lora_rank)
+
+x = torch.randn(q_lora_rank, dtype=torch.float16)*100
+print(x)
+print(q_a_layernorm(x))
+
+kv_a_layernorm = DeepseekV2RMSNorm(kv_lora_rank)
+
+q_a_proj = nn.Linear(hidden_size, q_lora_rank, bias=False, dtype=torch.float16)
+q_b_proj = nn.Linear(q_lora_rank, num_heads * (nope_size+rope_size) , bias=False, dtype=torch.float16)
+kv_a_proj_with_mqa = nn.Linear(hidden_size, kv_lora_rank + rope_size, bias=False, dtype=torch.float16)
+kv_b_proj = nn.Linear(kv_lora_rank, num_heads * (nope_size + nope_size), bias=False, dtype=torch.float16)
+o_proj = nn.Linear(num_heads * nope_size, hidden_size, bias=False, dtype=torch.float16)
+
+init.normal_(q_a_proj.weight, mean=0.0, std=0.02)
+init.normal_(q_b_proj.weight, mean=0.0, std=0.02)
+init.normal_(kv_a_proj_with_mqa.weight, mean=0.0, std=0.02)
+init.normal_(kv_b_proj.weight, mean=0.0, std=0.02)
+init.normal_(o_proj.weight, mean=0.0, std=0.02)
+# # 这里的权重初始化是为了测试
+# # 将权重设置为全 1
+# with torch.no_grad():
+#     q_a_proj.weight.fill_(1.0)
+#     q_b_proj.weight.fill_(1.0)
+#     kv_a_proj_with_mqa.weight.fill_(1.0)
+#     kv_b_proj.weight.fill_(1.0)
+#     o_proj.weight.fill_(1.0)
+
+q_absorb = torch.randn(num_heads, nope_size, kv_lora_rank, dtype=torch.float16)
+out_absorb = torch.randn(num_heads, nope_size, kv_lora_rank, dtype=torch.float16)
+
+rotary_emb = DeepseekV3YarnRotaryEmbedding(
+    rope_size,
+    max_position_embeddings=max_position_embeddings,
+    scaling_factor=rope_scaling["factor"],
+    base=rope_theta,
+    beta_fast=rope_scaling["beta_fast"],
+    beta_slow=rope_scaling["beta_slow"],
+    mscale=rope_scaling["mscale"],
+    mscale_all_dim=rope_scaling["mscale_all_dim"],
+    original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+)
+# 构造一个q_len 长度的输入 hidden_states, 对应的历史 kv_indptr 是[0:bsz]
+# kv_indices 是[0:bsz]，page_idx=[0:bsz], page_offset=[his_kv_len:q_len+his_kv_len]
+# last_page_len = [q_len+his_kv_len,...] layer_idx = 1
+# position_ids = [his_kv_len:q_len+his_kv_len]
+hidden_states = torch.randn(q_len, hidden_size, dtype=torch.float16)
+q_indptr = torch.tensor([0,q_len]).to(torch.int32)
+kv_indptr = torch.tensor(range(0, bsz_tensors[0] + 1)).to(torch.int32)
+kv_indices = torch.tensor(range(0, bsz_tensors[0])).to(torch.int32)
+page_idx = torch.tensor(range(0, bsz_tensors[0])).to(torch.int32)
+page_offset = torch.tensor(range(his_kv_len, his_kv_len + q_len)).to(torch.int32)
+last_page_len = torch.tensor([q_len+his_kv_len]*bsz_tensors[0], device=hidden_states.device)
+position_ids = torch.tensor(range(his_kv_len, his_kv_len + q_len)).to(torch.int32)
+
+
+# 按照行创建 mask [q_len,his_kv_len+q_len]
+attention_masks = torch.zeros((q_len, his_kv_len + q_len), dtype=torch.float16)
+for i in range(q_len):
+    attention_masks[i, i + his_kv_len + 1: i + his_kv_len + q_len] = -65504.0
+
+
+def torch_attn(hidden_states: torch.Tensor,
+                kv_cache: KDeepSeekV3Cache,
+                position_ids: torch.Tensor,
+                page_idx: torch.Tensor,
+                page_offset: torch.Tensor,
+                attention_masks: Optional[list[torch.Tensor]] = None,
+                q_indptr: Optional[torch.Tensor] = None,
+                kv_indices: Optional[torch.Tensor] = None,
+                kv_indptr: Optional[torch.Tensor] = None,
+                bsz_tensors: Optional[torch.Tensor] = None,
+                last_page_len: Optional[torch.Tensor] = None,
+                layer_idx: Optional[int] = None,
+                ):
+    global out_absorb
+    global q_absorb
+    # range bsz_tensors
+    final_attention_output = torch.tensor([], device=hidden_states.device)
+    for i in range(bsz_tensors[0]):
+        print("page_idx", page_idx)
+        print("page_offset", page_offset)
+        print("q_indptr", q_indptr)
+        print("kv_indices", kv_indices)
+        print("kv_indptr", kv_indptr)
+
+        batch_num_tokens_tensors = q_indptr[i+1] - q_indptr[i]
+        batch_last_page_len = last_page_len[i]
+        # kv_total_len is kv_len, batch_compressed_kv is compressed_kv, batch_k_pe is k_pe
+        batch_page_idx = page_idx[q_indptr[i]:q_indptr[i+1]]
+        print('batch_page_idx',batch_page_idx)
+        batch_page_offset = page_offset[q_indptr[i]:q_indptr[i+1]]
+        # kv_page_nums is the number of pages for the current batch
+        kv_page_nums = kv_indptr[i+1] - kv_indptr[i]
+        # kv_total_len is the total length of the kv cache for the current batch (kv_len for algorithm)
+        kv_total_len = kv_page_nums * page_size
+        if batch_last_page_len is not None:
+            kv_total_len = kv_total_len - (page_size - batch_last_page_len)
+        # print(f"kv_total_len's shape {kv_total_len.shape}")
+        # kv_index is the index of the kv cache pages for the current batch
+        kv_index = kv_indices[kv_indptr[i]:kv_indptr[i+1]]
+        # we can index [kv_index, page_offset_indices] to get the kv cache for the current batch
+        # from q_indptr[i] to q_indptr[i+1] is the range of the current batch
+        batch_hidden_states = hidden_states[q_indptr[i]:q_indptr[i+1]]
+        batch_position_ids = position_ids[q_indptr[i]:q_indptr[i+1]]
+        q_len, _ = batch_hidden_states.size()
+        # print("q_len -> ", q_len)
+        q_lora = q_a_proj(batch_hidden_states)
+        print('q_a_proj',q_a_proj.weight)
+        print('q_lora',q_lora)
+        
+        q = q_b_proj(q_a_layernorm(q_lora))
+        print('q_b_proj',q_b_proj.weight)
+        # for v3, bsz, q_len, num_heads(128), qk_head_dim(192=128(nope)+64(rope))
+        q = q.view(q_len, num_heads, nope_size+rope_size)
+        # q_nope is [q_len, num_heads(128), qk_nope_head_dim(128)]
+        # q_pe is [q_len, num_heads(128), qk_rope_head_dim(64)]
+        q_nope, q_pe = torch.split(
+            q, [nope_size, rope_size], dim=-1
+        )
+        print('q_nope',q_nope)
+        print('q_pe',q_pe)
+        # compressed_kv is [q_len, kv_lora_rank(512) + rope(64)]
+        compressed_kv = kv_a_proj_with_mqa(batch_hidden_states)
+        # compressed_kv is [q_len, kv_lora_rank(512)], k_pe is [q_len, rope(64)]
+        compressed_kv, k_pe = torch.split(
+            compressed_kv, [kv_lora_rank, rope_size], dim=-1
+        )
+        compressed_kv = compressed_kv.contiguous()
+        compressed_kv = kv_a_layernorm(compressed_kv)
+        # k_pe is [q_len, 1, qk_rope_head_dim(64)]
+        print('compressed_kv ',compressed_kv)
+        print('k_pe ',k_pe)
+        k_pe = k_pe.view(q_len, 1, rope_size)
+        # compressed_kv is [q_len, 1, kv_lora_rank(512)]
+        compressed_kv = compressed_kv.view(q_len, 1, kv_lora_rank)
+        
+        cos, sin = rotary_emb(q_pe, batch_position_ids)
+        # print(f"q_pe shape{q_pe.shape}, k_pe shape {k_pe.shape}")
+        q_pe, k_pe = apply_rotary_pos_emb(q_pe.unsqueeze(0), k_pe.unsqueeze(0), cos, sin, unsqueeze_dim=1)
+        q_pe = q_pe.squeeze(0)
+        # q_pe is [num_heads(128), q_len, qk_rope_head_dim(64)]
+        q_pe.transpose_(0, 1)            
+        if kv_cache is not None:
+            cache_kwargs = {"sin": sin, "cos": cos, "page_idx": batch_page_idx, "page_offset": batch_page_offset}  # Specific to RoPE models
+            compressed_kv_with_k_pe = kv_cache.update(compressed_kv.unsqueeze(0), k_pe, layer_idx, batch_page_idx, batch_page_offset, cache_kwargs)
+            compressed_kv = compressed_kv_with_k_pe [:, :, :, :kv_lora_rank].view(-1, page_size, kv_lora_rank)
+            k_pe = compressed_kv_with_k_pe [:, :, :, kv_lora_rank:].view(-1, page_size, rope_size)
+        # q_absorb is [num_heads(128), qk_nope_head_dim(128), kv_lora_rank(512)]
+        # out_absorb is [num_heads(128), kv_lora_rank(512), v_head_dim(128)] v_head_dim is also the nope dim
+        # q_absorb, out_absorb = get_absorbed()
+        # q_nope is [num_heads(128), q_len, qk_nope_head_dim(128)]
+        q_nope = q_nope.transpose(0, 1) # q_len is 1, no GPU overhead, same below
+        # q_nope is [num_heads(128), q_len, kv_lora_rank(512)]
+        q_nope = torch.matmul(q_nope, q_absorb) # batched MM
+
+        # # q_nope is [q_len, num_heads(128), kv_lora_rank(512)]
+        # q_nope = q_nope.transpose(0, 1)
+
+        # we need to index out the compressed_kv and k_pe for the current batch
+        batch_compressed_kv = None
+        batch_k_pe = None
+        for page_index in kv_index:
+            if kv_total_len > page_size:
+                tmp_compressed_kv = compressed_kv[page_index, 0:page_size, :]
+                tmp_k_pe = k_pe[page_index, 0:page_size, :]
+                if batch_compressed_kv is None or batch_k_pe is None:
+                    batch_compressed_kv = tmp_compressed_kv
+                    batch_k_pe = tmp_k_pe
+                else: 
+                    batch_compressed_kv = torch.cat((batch_compressed_kv, tmp_compressed_kv), dim=0)
+                    batch_k_pe = torch.cat((batch_k_pe, tmp_k_pe), dim=0)
+                kv_total_len -= page_size
+            else:
+                tmp_compressed_kv = compressed_kv[page_index, 0:kv_total_len, :]
+                tmp_k_pe = k_pe[page_index, 0:kv_total_len, :]
+                if batch_compressed_kv is None or batch_k_pe is None:
+                    batch_compressed_kv = tmp_compressed_kv
+                    batch_k_pe = tmp_k_pe
+                else: 
+                    batch_compressed_kv = torch.cat((batch_compressed_kv, tmp_compressed_kv), dim=0)
+                    batch_k_pe = torch.cat((batch_k_pe, tmp_k_pe), dim=0)
+                break
+        # batch_compressed_kv is [kv_total_len(k_len), kv_lora_rank(512)]
+        # batch_k_pe is [kv_total_len(k_len), qk_rope_head_dim(64)]
+        pe_weights = torch.matmul(q_pe,batch_k_pe.mT)
+        print('pe_weights',pe_weights)
+        attention_weights = (pe_weights + torch.matmul(q_nope, batch_compressed_kv.mT)) * softmax_scale
+        # attention_weights is [num_heads(128), q_len, k_len]
+        
+        # attention_weights = attention_weights.transpose(0,1).unsqueeze(0).squeeze(-1).expand(q_len,-1,-1).transpose(0,1)
+        
+        # attention_masks[i] is [q_len, k_len]
+        
+        attention_weights = (attention_weights + attention_masks[i])
+        # attention_weights shape is [num_heads(128), q_len, k_len]
+        attention_weights = nn.functional.softmax(attention_weights,dim=-1,dtype=torch.float16).to(q_pe.dtype)
+        attn_output = torch.matmul(attention_weights, batch_compressed_kv) # [num_heads(128),q_len, lora_rank(512)]
+        # out_absorb shape is [num_heads(128), kv_lora_rank(512), v_head_dim(128)]
+        out_absorb = out_absorb.transpose(1,2)
+        # q for q_len, n for num_heads, h for v_head_dim, v for kv_lora_rank
+        attn_output = torch.matmul(attn_output, out_absorb) # [num_heads(128), q_len, v_head_dim(128)]
+        attn_output = attn_output.transpose(0, 1) # [q_len, num_heads(128), v_head_dim(128)]
+        attn_output = attn_output.reshape(q_len, num_heads * nope_size)
+        attn_output = o_proj(attn_output)
+        final_attention_output = torch.cat((final_attention_output, attn_output), dim=0)
+    return final_attention_output
+
+
+
+def torch_attn_for_test(hidden_states,kv_cache,):
+    pass
+
+def test_mla_simple():
+    result = torch_attn(
+        hidden_states,
+        kv_cache,
+        position_ids,
+        page_idx,
+        page_offset,
+        attention_masks=attention_masks,
+        q_indptr=q_indptr,
+        kv_indices=kv_indices,
+        kv_indptr=kv_indptr,
+        bsz_tensors=bsz_tensors,
+        last_page_len=last_page_len,
+        layer_idx=1
+    )
+    print(result.shape)
+    print(result)
+    
+test_mla_simple()

kt-kernel/examples/test_mla_torch.py

@@ -0,0 +1,336 @@
+import os, sys
+import time
+from typing import Optional
+
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+from kt_kernel import kt_kernel_ext
+from kt_kernel_ext.kvcache import ggml_type
+import torch
+from torch import nn
+from torch.nn import init
+from torch_attention import apply_rotary_pos_emb, DeepseekV2RMSNorm, KDeepSeekV3Cache, DeepseekV3YarnRotaryEmbedding
+
+
+seed = 42  # 你可以选择任何整数作为种子
+torch.manual_seed(seed)
+torch.cuda.manual_seed_all(seed)
+
+qlen = 1024
+kvlen = 0
+
+
+page_table = range(20)
+bsz_tensors = torch.tensor([1])
+
+
+page_size = 256
+pages_count = 200
+tp_count = 4
+
+
+hidden_size = 7168
+q_lora_rank = 1536
+kv_lora_rank = 512
+num_heads = 128
+nope_size = 128
+rope_size = 64
+
+rope_theta = 10000
+max_qlen = 1024
+max_kvlen = 4096
+
+max_position_embeddings = 163840
+
+
+rope_scaling = {
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "factor": 40,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0,
+    "original_max_position_embeddings": 4096,
+    "type": "yarn",
+}
+
+
+CPUInfer = kt_kernel_ext.CPUInfer(64)
+validation_iter = 100
+
+
+q_a_proj = nn.Linear(hidden_size, q_lora_rank, bias=False, dtype=torch.float16)
+q_b_proj = nn.Linear(q_lora_rank, num_heads * (nope_size + rope_size), bias=False, dtype=torch.float16)
+kv_a_proj_with_mqa = nn.Linear(hidden_size, kv_lora_rank + rope_size, bias=False, dtype=torch.float16)
+kv_b_proj = nn.Linear(kv_lora_rank, num_heads * (nope_size + nope_size), bias=False, dtype=torch.float16)
+o_proj = nn.Linear(num_heads * nope_size, hidden_size, bias=False, dtype=torch.float16)
+
+init.normal_(q_a_proj.weight, mean=0.0, std=0.02)
+init.normal_(q_b_proj.weight, mean=0.0, std=0.02)
+init.normal_(kv_a_proj_with_mqa.weight, mean=0.0, std=0.02)
+init.normal_(kv_b_proj.weight, mean=0.0, std=0.02)
+init.normal_(o_proj.weight, mean=0.0, std=0.02)
+
+q_a_proj_weight = q_a_proj.weight.to(torch.float16).to("cpu").contiguous()
+q_b_proj_weight = q_b_proj.weight.to(torch.float16).to("cpu").contiguous()
+kv_a_proj_with_mqa_weight = kv_a_proj_with_mqa.weight.to("cpu").to(torch.float16).contiguous()
+kv_b_proj_weight = kv_b_proj.weight.to(torch.float16).to("cpu").contiguous()
+o_proj_weight = o_proj.weight.to(torch.float16).to("cpu").contiguous()
+
+
+config = kt_kernel_ext.mla.MLAConfig(
+    hidden_size,
+    q_lora_rank,
+    kv_lora_rank,
+    num_heads,
+    nope_size,
+    rope_size,
+)
+config.max_qlen = max_qlen
+config.max_kvlen = max_kvlen
+config.max_position_embeddings = max_position_embeddings
+config.rope_scaling_factor = rope_scaling["factor"]
+config.rope_theta = rope_theta
+config.rope_scaling_beta_fast = rope_scaling["beta_fast"]
+config.rope_scaling_beta_slow = rope_scaling["beta_slow"]
+config.rope_scaling_mscale = rope_scaling["mscale"]
+config.rope_scaling_mscale_all_dim = rope_scaling["mscale_all_dim"]
+config.rope_scaling_original_max_position_embeddings = rope_scaling["original_max_position_embeddings"]
+
+config.q_a_proj = q_a_proj_weight.data_ptr()
+config.q_b_proj = q_b_proj_weight.data_ptr()
+config.kv_a_proj_with_mqa = kv_a_proj_with_mqa_weight.data_ptr()
+config.kv_b_proj = kv_b_proj_weight.data_ptr()
+config.o_proj = o_proj_weight.data_ptr()
+
+config.q_a_proj_type = ggml_type.FP16
+config.q_b_proj_type = ggml_type.FP16
+config.kv_a_proj_with_mqa_type = ggml_type.FP16
+config.kv_b_proj_type = ggml_type.FP16
+config.w_o_type = ggml_type.FP16
+
+
+config.pool = CPUInfer.backend_
+
+
+mla = kt_kernel_ext.mla.MLA(config)
+mla.load_weights()
+mla.set_local_pages(pages_count)
+
+
+input = torch.randn((qlen, hidden_size), dtype=torch.float16).to("cpu").contiguous()
+
+
+output = torch.zeros((qlen, hidden_size), dtype=torch.float16).to("cpu").contiguous()
+mla.forward([qlen], [page_table], [kvlen], input.data_ptr(), output.data_ptr())
+print("CPU MLA Output: ", output)
+
+
+softmax_scale = (nope_size + rope_size) ** -0.5
+# 1代表的是压缩的kv的头数
+k_caches = torch.randn(1, pages_count, page_size, 1, kv_lora_rank + rope_size).to(torch.float16)
+kv_cache = KDeepSeekV3Cache(page_size=page_size, kv_lora_rank=kv_lora_rank, k_caches=k_caches)
+
+q_a_layernorm = DeepseekV2RMSNorm(q_lora_rank)
+
+x = torch.randn(q_lora_rank, dtype=torch.float16) * 100
+print(x)
+print(q_a_layernorm(x))
+
+kv_a_layernorm = DeepseekV2RMSNorm(kv_lora_rank)
+
+
+q_absorb = torch.randn(num_heads, nope_size, kv_lora_rank, dtype=torch.float16)
+out_absorb = torch.randn(num_heads, nope_size, kv_lora_rank, dtype=torch.float16)
+
+rotary_emb = DeepseekV3YarnRotaryEmbedding(
+    rope_size,
+    max_position_embeddings=max_position_embeddings,
+    scaling_factor=rope_scaling["factor"],
+    base=rope_theta,
+    beta_fast=rope_scaling["beta_fast"],
+    beta_slow=rope_scaling["beta_slow"],
+    mscale=rope_scaling["mscale"],
+    mscale_all_dim=rope_scaling["mscale_all_dim"],
+    original_max_position_embeddings=rope_scaling["original_max_position_embeddings"],
+)
+# 构造一个qlen 长度的输入 hidden_states, 对应的历史 kv_indptr 是[0:bsz]
+# kv_indices 是[0:bsz]，page_idx=[0:bsz], page_offset=[kvlen:qlen+kvlen]
+# last_page_len = [qlen+kvlen,...] layer_idx = 1
+# position_ids = [kvlen:qlen+kvlen]
+hidden_states = torch.randn(qlen, hidden_size, dtype=torch.float16)
+q_indptr = torch.tensor([0, qlen]).to(torch.int32)
+
+kv_indptr = torch.tensor([0, (qlen + kvlen + page_size - 1) // page_size]).to(torch.int32)
+kv_indices = torch.tensor(range(pages_count)).to(torch.int32)
+
+page_idx = torch.tensor([i // page_size for i in range(kvlen, kvlen + qlen)]).to(torch.int32)
+page_offset = torch.tensor([i % page_size for i in range(kvlen, kvlen + qlen)]).to(torch.int32)
+
+last_page_len = torch.tensor([(qlen + kvlen) % page_size], device=hidden_states.device)
+position_ids = torch.tensor(range(kvlen, kvlen + qlen)).to(torch.int32)
+
+
+# 按照行创建 mask [qlen,kvlen+qlen]
+attention_masks = torch.zeros((qlen, kvlen + qlen), dtype=torch.float16)
+for i in range(qlen):
+    attention_masks[i, i + kvlen + 1 : i + kvlen + qlen] = -65504.0
+
+
+def torch_attn(
+    hidden_states: torch.Tensor,
+    kv_cache: KDeepSeekV3Cache,
+    position_ids: torch.Tensor,
+    page_idx: torch.Tensor,
+    page_offset: torch.Tensor,
+    attention_masks: Optional[list[torch.Tensor]] = None,
+    q_indptr: Optional[torch.Tensor] = None,
+    kv_indices: Optional[torch.Tensor] = None,
+    kv_indptr: Optional[torch.Tensor] = None,
+    bsz_tensors: Optional[torch.Tensor] = None,
+    last_page_len: Optional[torch.Tensor] = None,
+    layer_idx: Optional[int] = None,
+):
+    global out_absorb
+    global q_absorb
+    # range bsz_tensors
+    final_attention_output = torch.tensor([], device=hidden_states.device)
+    for i in range(bsz_tensors[0]):
+        batch_num_tokens_tensors = q_indptr[i + 1] - q_indptr[i]
+        batch_last_page_len = last_page_len[i]
+        # kv_total_len is kv_len, batch_compressed_kv is compressed_kv, batch_k_pe is k_pe
+        batch_page_idx = page_idx[q_indptr[i] : q_indptr[i + 1]]
+        batch_page_offset = page_offset[q_indptr[i] : q_indptr[i + 1]]
+        # kv_page_nums is the number of pages for the current batch
+        kv_page_nums = kv_indptr[i + 1] - kv_indptr[i]
+        # kv_total_len is the total length of the kv cache for the current batch (kv_len for algorithm)
+        kv_total_len = kv_page_nums * page_size
+        if batch_last_page_len is not None:
+            kv_total_len = kv_total_len - (page_size - batch_last_page_len)
+        # print(f"kv_total_len's shape {kv_total_len.shape}")
+        # kv_index is the index of the kv cache pages for the current batch
+        kv_index = kv_indices[kv_indptr[i] : kv_indptr[i + 1]]
+        # we can index [kv_index, page_offset_indices] to get the kv cache for the current batch
+        # from q_indptr[i] to q_indptr[i+1] is the range of the current batch
+        batch_hidden_states = hidden_states[q_indptr[i] : q_indptr[i + 1]]
+        batch_position_ids = position_ids[q_indptr[i] : q_indptr[i + 1]]
+        qlen, _ = batch_hidden_states.size()
+        # print("qlen -> ", qlen)
+        q_lora = q_a_proj(batch_hidden_states)
+        print("q_a_proj", q_a_proj.weight)
+        print("q_lora", q_lora)
+
+        q = q_b_proj(q_a_layernorm(q_lora))
+        print("q_b_proj", q_b_proj.weight)
+        # for v3, bsz, qlen, num_heads(128), qk_head_dim(192=128(nope)+64(rope))
+        q = q.view(qlen, num_heads, nope_size + rope_size)
+        # q_nope is [qlen, num_heads(128), qk_nope_head_dim(128)]
+        # q_pe is [qlen, num_heads(128), qk_rope_head_dim(64)]
+        q_nope, q_pe = torch.split(q, [nope_size, rope_size], dim=-1)
+        print("q_nope", q_nope)
+        print("q_pe", q_pe)
+        # compressed_kv is [qlen, kv_lora_rank(512) + rope(64)]
+        compressed_kv = kv_a_proj_with_mqa(batch_hidden_states)
+        # compressed_kv is [qlen, kv_lora_rank(512)], k_pe is [qlen, rope(64)]
+        compressed_kv, k_pe = torch.split(compressed_kv, [kv_lora_rank, rope_size], dim=-1)
+        compressed_kv = compressed_kv.contiguous()
+        compressed_kv = kv_a_layernorm(compressed_kv)
+        # k_pe is [qlen, 1, qk_rope_head_dim(64)]
+        print("compressed_kv ", compressed_kv)
+        print("k_pe ", k_pe)
+        k_pe = k_pe.view(qlen, 1, rope_size)
+        # compressed_kv is [qlen, 1, kv_lora_rank(512)]
+        compressed_kv = compressed_kv.view(qlen, 1, kv_lora_rank)
+
+        cos, sin = rotary_emb(q_pe, batch_position_ids)
+        # print(f"q_pe shape{q_pe.shape}, k_pe shape {k_pe.shape}")
+        q_pe, k_pe = apply_rotary_pos_emb(q_pe.unsqueeze(0), k_pe.unsqueeze(0), cos, sin, unsqueeze_dim=1)
+        q_pe = q_pe.squeeze(0)
+        # q_pe is [num_heads(128), qlen, qk_rope_head_dim(64)]
+        q_pe.transpose_(0, 1)
+        if kv_cache is not None:
+            cache_kwargs = {
+                "sin": sin,
+                "cos": cos,
+                "page_idx": batch_page_idx,
+                "page_offset": batch_page_offset,
+            }  # Specific to RoPE models
+            compressed_kv_with_k_pe = kv_cache.update(
+                compressed_kv.unsqueeze(0), k_pe, layer_idx, batch_page_idx, batch_page_offset, cache_kwargs
+            )
+            compressed_kv = compressed_kv_with_k_pe[:, :, :, :kv_lora_rank].view(-1, page_size, kv_lora_rank)
+            k_pe = compressed_kv_with_k_pe[:, :, :, kv_lora_rank:].view(-1, page_size, rope_size)
+        # q_absorb is [num_heads(128), qk_nope_head_dim(128), kv_lora_rank(512)]
+        # out_absorb is [num_heads(128), kv_lora_rank(512), v_head_dim(128)] v_head_dim is also the nope dim
+        # q_absorb, out_absorb = get_absorbed()
+        # q_nope is [num_heads(128), qlen, qk_nope_head_dim(128)]
+        q_nope = q_nope.transpose(0, 1)  # qlen is 1, no GPU overhead, same below
+        # q_nope is [num_heads(128), qlen, kv_lora_rank(512)]
+        q_nope = torch.matmul(q_nope, q_absorb)  # batched MM
+
+        # # q_nope is [qlen, num_heads(128), kv_lora_rank(512)]
+        # q_nope = q_nope.transpose(0, 1)
+
+        # we need to index out the compressed_kv and k_pe for the current batch
+        batch_compressed_kv = None
+        batch_k_pe = None
+        for page_index in kv_index:
+            if kv_total_len > page_size:
+                tmp_compressed_kv = compressed_kv[page_index, 0:page_size, :]
+                tmp_k_pe = k_pe[page_index, 0:page_size, :]
+                if batch_compressed_kv is None or batch_k_pe is None:
+                    batch_compressed_kv = tmp_compressed_kv
+                    batch_k_pe = tmp_k_pe
+                else:
+                    batch_compressed_kv = torch.cat((batch_compressed_kv, tmp_compressed_kv), dim=0)
+                    batch_k_pe = torch.cat((batch_k_pe, tmp_k_pe), dim=0)
+                kv_total_len -= page_size
+            else:
+                tmp_compressed_kv = compressed_kv[page_index, 0:kv_total_len, :]
+                tmp_k_pe = k_pe[page_index, 0:kv_total_len, :]
+                if batch_compressed_kv is None or batch_k_pe is None:
+                    batch_compressed_kv = tmp_compressed_kv
+                    batch_k_pe = tmp_k_pe
+                else:
+                    batch_compressed_kv = torch.cat((batch_compressed_kv, tmp_compressed_kv), dim=0)
+                    batch_k_pe = torch.cat((batch_k_pe, tmp_k_pe), dim=0)
+                break
+        # batch_compressed_kv is [kv_total_len(k_len), kv_lora_rank(512)]
+        # batch_k_pe is [kv_total_len(k_len), qk_rope_head_dim(64)]
+        pe_weights = torch.matmul(q_pe, batch_k_pe.mT)
+        print("pe_weights", pe_weights)
+        attention_weights = (pe_weights + torch.matmul(q_nope, batch_compressed_kv.mT)) * softmax_scale
+        # attention_weights is [num_heads(128), qlen, k_len]
+
+        # attention_weights = attention_weights.transpose(0,1).unsqueeze(0).squeeze(-1).expand(qlen,-1,-1).transpose(0,1)
+
+        # attention_masks[i] is [qlen, k_len]
+
+        attention_weights = attention_weights + attention_masks[i]
+        # attention_weights shape is [num_heads(128), qlen, k_len]
+        attention_weights = nn.functional.softmax(attention_weights, dim=-1, dtype=torch.float16).to(q_pe.dtype)
+        attn_output = torch.matmul(attention_weights, batch_compressed_kv)  # [num_heads(128),qlen, lora_rank(512)]
+        # out_absorb shape is [num_heads(128), kv_lora_rank(512), v_head_dim(128)]
+        out_absorb = out_absorb.transpose(1, 2)
+        # q for qlen, n for num_heads, h for v_head_dim, v for kv_lora_rank
+        attn_output = torch.matmul(attn_output, out_absorb)  # [num_heads(128), qlen, v_head_dim(128)]
+        attn_output = attn_output.transpose(0, 1)  # [qlen, num_heads(128), v_head_dim(128)]
+        attn_output = attn_output.reshape(qlen, num_heads * nope_size)
+        attn_output = o_proj(attn_output)
+        final_attention_output = torch.cat((final_attention_output, attn_output), dim=0)
+    return final_attention_output
+
+
+torch_output = torch_attn(
+    input,
+    kv_cache,
+    position_ids,
+    page_idx,
+    page_offset,
+    attention_masks=attention_masks,
+    q_indptr=q_indptr,
+    kv_indices=kv_indices,
+    kv_indptr=kv_indptr,
+    bsz_tensors=bsz_tensors,
+    last_page_len=last_page_len,
+    layer_idx=0,
+)
+print("Torch Output: ", torch_output)

kt-kernel/examples/test_mlp.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""
+Description  :
+Author       : chenht2022
+Date         : 2024-07-25 10:32:05
+Version      : 1.0.0
+LastEditors  : chenht2022
+LastEditTime : 2024-08-06 10:37:28
+Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+"""
+import os, sys
+import time
+
+sys.path.append(os.path.dirname(__file__) + "/../build")
+from kt_kernel import kt_kernel_ext
+import torch
+
+hidden_size = 5120
+intermediate_size = 3072
+stride = 32
+group_max_len = 1024
+gate_type = 1  # ggml_type::GGML_TYPE_F16
+up_type = 1  # ggml_type::GGML_TYPE_F16
+down_type = 1  # ggml_type::GGML_TYPE_F16
+hidden_type = 1  # ggml_type::GGML_TYPE_F16
+qlen = 30
+layer_num = 10
+CPUInfer = kt_kernel_ext.CPUInfer(48)
+validation_iter = 100
+
+
+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
+
+
+with torch.inference_mode(mode=True):
+    mlps = []
+    gate_projs = []
+    up_projs = []
+    down_projs = []
+    for _ in range(layer_num):
+        gate_proj = (
+            torch.randn((intermediate_size, hidden_size), dtype=torch.float16, device="cuda").to("cpu").contiguous()
+        )
+        up_proj = (
+            torch.randn((intermediate_size, hidden_size), dtype=torch.float16, device="cuda").to("cpu").contiguous()
+        )
+        down_proj = (
+            torch.randn((hidden_size, intermediate_size), dtype=torch.float16, device="cuda").to("cpu").contiguous()
+        )
+        config = kt_kernel_ext.mlp.MLPConfig(
+            hidden_size,
+            intermediate_size,
+            stride,
+            group_max_len,
+            gate_proj.data_ptr(),
+            up_proj.data_ptr(),
+            down_proj.data_ptr(),
+            gate_type,
+            up_type,
+            down_type,
+            hidden_type,
+        )
+        mlp = kt_kernel_ext.mlp.MLP(config)
+        gate_projs.append(gate_proj)
+        up_projs.append(up_proj)
+        down_projs.append(down_proj)
+        mlps.append(mlp)
+
+    # validation
+    for i in range(validation_iter):
+        mlp = mlps[i % layer_num]
+        input = torch.randn((qlen, hidden_size), dtype=torch.float16).contiguous()
+        output = torch.empty((qlen, hidden_size), dtype=torch.float16).contiguous()
+        input = input / 100
+
+        CPUInfer.submit(mlp.forward(qlen, input.data_ptr(), output.data_ptr()))
+        CPUInfer.sync()
+        # print('cpuinfer output', output)
+
+        gate_proj = gate_projs[i % layer_num]
+        up_proj = up_projs[i % layer_num]
+        down_proj = down_projs[i % layer_num]
+        t_output = mlp_torch(input, gate_proj, up_proj, down_proj)
+        # print('torch output', t_output)
+
+        diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
+        print("diff = ", diff)
+        assert diff < 0.001

kt-kernel/examples/test_moe.py

@@ -0,0 +1,210 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""
+Description  :
+Author       : chenht2022
+Date         : 2024-07-25 10:32:05
+Version      : 1.0.0
+LastEditors  : SkqLiao
+LastEditTime : 2025-03-13 11:38:05
+Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+"""
+import os, sys
+import time
+
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+from kt_kernel import kt_kernel_ext
+import torch
+from tqdm import tqdm
+from kt_kernel_ext.kvcache import ggml_type
+
+torch.manual_seed(0)
+
+expert_num = 8
+hidden_size = 2048  # 7168
+intermediate_size = 2048
+stride = 32
+group_min_len = 10
+group_max_len = 2560
+num_experts_per_tok = 8
+layer_num = 1
+# expert_num = 8
+# hidden_size = 7168
+# intermediate_size = 2048
+# stride = 32
+# group_min_len = 10
+# group_max_len = 10240
+# num_experts_per_tok = 8
+# qlen = 1024
+# layer_num = 1
+CPUInfer = kt_kernel_ext.CPUInfer(64)
+validation_iter = 10
+
+
+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 to_cpuinfer_tensor(tensor, type):
+    size = torch.prod(torch.tensor(tensor.shape, dtype=torch.int32)).item()
+    return kt_kernel_ext.utils.from_float(tensor.data_ptr(), size, type)
+
+
+def from_cpuinfer_tensor(tensor, size, type):
+    return kt_kernel_ext.utils.to_float(tensor.data_ptr(), size, type)
+
+
+qlens = [1, 64]  # [64, 512, 2048, 8192, 16384]
+# gate_types = [ggml_type.FP32, ggml_type.FP16, ggml_type.Q8_0, ggml_type.Q6_K, ggml_type.Q5_K, ggml_type.Q4_K, ggml_type.Q3_K]
+# up_types = [ggml_type.FP32, ggml_type.FP16, ggml_type.Q8_0, ggml_type.Q6_K, ggml_type.Q5_K, ggml_type.Q4_K, ggml_type.Q3_K]
+# down_types = [ggml_type.FP32, ggml_type.FP16, ggml_type.Q8_0, ggml_type.Q6_K, ggml_type.Q6_K, ggml_type.Q6_K, ggml_type.Q5_K]
+gate_types = [ggml_type.Q4_K]
+up_types = [ggml_type.Q4_K]
+down_types = [ggml_type.Q6_K]
+hidden_type = ggml_type.BF16
+print(f"Parameters: expert_num: {expert_num} hidden_size: {hidden_size} intermediate_size: {intermediate_size}")
+print(f"group_max_len: ", group_max_len)
+
+for qlen in qlens:
+    for gate_type, up_type, down_type in zip(gate_types, up_types, down_types):
+        with torch.inference_mode(mode=True):
+            moes = []
+            gate_projs = []
+            up_projs = []
+            down_projs = []
+            print("Preparing data...")
+            converted_tensors = []
+            for _ in range(layer_num):
+                size = expert_num * intermediate_size * hidden_size
+                gate_proj = (
+                    torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device="cuda")
+                    .to("cpu")
+                    .contiguous()
+                )
+                up_proj = (
+                    torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device="cuda")
+                    .to("cpu")
+                    .contiguous()
+                )
+                down_proj = (
+                    torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device="cuda")
+                    .to("cpu")
+                    .contiguous()
+                )
+
+                gate_tensor = to_cpuinfer_tensor(gate_proj, gate_type)
+                up_tensor = to_cpuinfer_tensor(up_proj, up_type)
+                down_tensor = to_cpuinfer_tensor(down_proj, down_type)
+
+                config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size)
+                config.pool = CPUInfer.backend_
+                config.stride = stride
+                config.group_min_len = group_min_len
+                config.group_max_len = group_max_len
+                config.gate_proj = gate_tensor.data_ptr()
+                config.up_proj = up_tensor.data_ptr()
+                config.down_proj = down_tensor.data_ptr()
+                config.gate_type = gate_type
+                config.up_type = up_type
+                config.down_type = down_type
+                config.hidden_type = hidden_type
+
+                moe = kt_kernel_ext.moe.MOE(config)
+                gate_projs.append(gate_proj)
+                up_projs.append(up_proj)
+                down_projs.append(down_proj)
+                CPUInfer.submit(moe.load_weights_task())
+                CPUInfer.sync()
+                moes.append(moe)
+                converted_tensors.append((gate_tensor, up_tensor, down_tensor))
+            print("Finished initialization!")
+
+            CPUInfer.submit(moes[0].warm_up_task())
+            CPUInfer.sync()
+            print("Warm up finished!")
+
+            # validation
+            progress_bar = tqdm(range(validation_iter), desc="Starting")
+            total_diff = 0
+
+            for i in tqdm(progress_bar):
+                progress_bar.set_description("Round: {}/{}".format(i + 1, validation_iter))
+                expert_ids = torch.stack(
+                    [torch.randperm(expert_num)[:num_experts_per_tok] for _ in range(qlen)]
+                ).contiguous()
+                weights = torch.rand((qlen, num_experts_per_tok), dtype=torch.float32).contiguous()
+                input_proj = torch.randn((qlen, hidden_size), dtype=torch.float32).contiguous() / 100
+                output_proj = torch.empty((qlen, hidden_size), dtype=torch.float32).contiguous()
+
+                input_tensor = to_cpuinfer_tensor(input_proj, hidden_type)
+                output_tensor = to_cpuinfer_tensor(output_proj, hidden_type)
+
+                qlen_tensor = torch.tensor([qlen], dtype=torch.int32)
+                moe = moes[i % layer_num]
+                CPUInfer.submit(
+                    moe.forward_task(
+                        qlen_tensor.data_ptr(),
+                        num_experts_per_tok,
+                        expert_ids.data_ptr(),
+                        weights.data_ptr(),
+                        input_tensor.data_ptr(),
+                        output_tensor.data_ptr(),
+                    )
+                )
+                CPUInfer.sync()
+                cpu_output = from_cpuinfer_tensor(output_tensor, qlen * hidden_size, hidden_type)
+
+                gate_proj = gate_projs[i % layer_num]
+                up_proj = up_projs[i % layer_num]
+                down_proj = down_projs[i % layer_num]
+                t_output = moe_torch(input_proj, expert_ids, weights, gate_proj, up_proj, down_proj)
+                print("cpuinfer output", cpu_output)
+                print("torch output", t_output)
+                diff = torch.mean(torch.abs(cpu_output.flatten() - t_output.flatten())) / torch.mean(
+                    torch.abs(t_output.flatten())
+                )
+                assert diff < 0.5
+                total_diff += diff
+
+            print(f"gate_type: {gate_type}, up_type: {up_type}, down_type: {down_type}")
+            print(f"Average diff: {total_diff / validation_iter:.4f}")

kt-kernel/examples/test_moe_amx.py

@@ -0,0 +1,562 @@
+import os, sys
+
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+print("sys.path:", sys.path)
+
+import torch
+from kt_kernel import kt_kernel_ext
+
+# Model configuration
+expert_num = 256
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 25600
+num_experts_per_tok = 8
+qlen = 1
+# qlen = 640
+layer_num = 1
+
+# Test configuration
+num_threads = 90
+CPUInfer = kt_kernel_ext.CPUInfer(num_threads)
+# validation_iter = 10000
+validation_iter = 2
+k_group_size = 64
+debug_print_count = 16  # Number of values to print in debug output
+physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous()
+
+# Performance test configuration
+perf_warmup_iter = 5  # Number of warmup iterations for performance test
+perf_test_iter = 20  # Number of iterations for performance measurement
+perf_qlen = 128  # Sequence length for performance testing
+
+
+def act_fn(x):
+    return x / (1.0 + torch.exp(-x))
+
+
+def mlp_torch(input, gate_proj, up_proj, down_proj, debug_expert_id=None, debug_print=False):
+    gate_buf = torch.mm(input, gate_proj.t())
+    up_buf = torch.mm(input, up_proj.t())
+
+    if debug_print and debug_expert_id is not None:
+        print(f"[TORCH DEBUG] Expert {debug_expert_id}:")
+        print(f"  gate_buf[:{debug_print_count}] = {gate_buf.flatten()[:debug_print_count]}")
+        print(f"  up_buf[:{debug_print_count}] = {up_buf.flatten()[:debug_print_count]}")
+
+    intermediate = act_fn(gate_buf) * up_buf
+
+    if debug_print and debug_expert_id is not None:
+        print(f"  intermediate[:{debug_print_count}] = {intermediate.flatten()[:debug_print_count]}")
+
+    ret = torch.mm(intermediate, down_proj.t())
+
+    if debug_print and debug_expert_id is not None:
+        print(f"  down_output[:{debug_print_count}] = {ret.flatten()[:debug_print_count]}")
+
+    return ret
+
+
+def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj, debug_print=False):
+    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]]
+
+    # Get the first expert from expert_ids array to match AWQ-MoE behavior
+    target_debug_expert = expert_ids[0, 0].item()  # First expert in expert_ids array
+
+    outputs = []
+    start_idx = 0
+    activated_experts = []
+
+    for i, num_tokens in enumerate(tokens_per_expert):
+        end_idx = start_idx + num_tokens
+        if num_tokens == 0:
+            continue
+        activated_experts.append(i)
+        tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
+        # Only debug the target expert that matches AWQ-MoE's first expert
+        should_debug = debug_print and i == target_debug_expert
+        expert_out = mlp_torch(
+            tokens_for_this_expert, gate_proj[i], up_proj[i], down_proj[i], debug_expert_id=i, debug_print=should_debug
+        )
+        outputs.append(expert_out)
+        start_idx = end_idx
+
+    if debug_print:
+        print(f"[TORCH DEBUG] Processing activated experts: {activated_experts}")
+        print(f"[TORCH DEBUG] Target debug expert (matches AWQ): {target_debug_expert}")
+
+    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)
+    )
+
+    if debug_print:
+        print(f"[TORCH DEBUG] Final MoE output[:{debug_print_count}] = {t_output.flatten()[:debug_print_count]}")
+
+    return t_output
+
+
+def test_moe(quant_mode: str):
+    assert (
+        quant_mode == "bf16"
+        or quant_mode == "int8"
+        or quant_mode == "int4"
+        or quant_mode == "int4_1"
+        or quant_mode == "int4_1k"
+    )
+    with torch.inference_mode(mode=True):
+        moes = []
+        gate_projs = []
+        up_projs = []
+        down_projs = []
+        for _ in range(layer_num):
+            gate_proj = (
+                torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cuda")
+                .to("cpu")
+                .contiguous()
+            )
+            up_proj = (
+                torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cuda")
+                .to("cpu")
+                .contiguous()
+            )
+            down_proj = (
+                torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.bfloat16, device="cuda")
+                .to("cpu")
+                .contiguous()
+            )
+            config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size, 0)
+            config.max_len = max_len
+            config.gate_proj = gate_proj.data_ptr()
+            config.up_proj = up_proj.data_ptr()
+            config.down_proj = down_proj.data_ptr()
+            config.gate_scale = 0
+            config.pool = CPUInfer.backend_
+            if quant_mode == "bf16":
+                moe = kt_kernel_ext.moe.AMXBF16_MOE(config)
+                CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+                CPUInfer.sync()
+                CPUInfer.submit(moe.warm_up_task())
+                CPUInfer.sync()
+            elif quant_mode == "int8":
+                moe = kt_kernel_ext.moe.AMXInt8_MOE(config)
+                CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+                CPUInfer.sync()
+                # CPUInfer.submit(moe.warm_up_task())
+                # CPUInfer.sync()
+            elif quant_mode == "int4":
+                moe = kt_kernel_ext.moe.AMXInt4_MOE(config)
+                CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+                CPUInfer.sync()
+                CPUInfer.submit(moe.warm_up_task())
+                CPUInfer.sync()
+            elif quant_mode == "int4_1":
+                moe = kt_kernel_ext.moe.AMXInt4_1_MOE(config)
+                CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+                CPUInfer.sync()
+                CPUInfer.submit(moe.warm_up_task())
+                CPUInfer.sync()
+            elif quant_mode == "int4_1k":
+                config.quant_config.bits = 4
+                config.quant_config.group_size = k_group_size
+                config.quant_config.zero_point = True
+                moe = kt_kernel_ext.moe.AMXInt4_1KGroup_MOE(config)
+                # import debugpy
+                # debugpy.listen(("127.0.0.1", 5678))
+                # debugpy.wait_for_client()
+                # debugpy.breakpoint()
+                print(f"the physical_logical map:{physical_to_logical_map.data_ptr()}")
+                CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+                CPUInfer.sync()
+                # CPUInfer.submit(moe.warm_up_task())
+                # CPUInfer.sync()
+            gate_projs.append(gate_proj)
+            up_projs.append(up_proj)
+            down_projs.append(down_proj)
+            moes.append(moe)
+
+        # validation
+        for i in range(validation_iter):
+            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.rand((qlen, num_experts_per_tok), dtype=torch.float32).contiguous()
+            input = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            input = input / 100
+            moe = moes[i % layer_num]
+
+            # Enable debug for first few iterations
+            enable_debug = i < 2
+            enable_debug = False
+            if enable_debug:
+                print(f"\n=== Iteration {i} Debug Info ===")
+                print(f"input[:{debug_print_count}] = {input.flatten()[:debug_print_count]}")
+                print(f"expert_ids = {expert_ids}")
+                print(f"weights = {weights}")
+                # Print which experts will be activated for comparison
+                activated_experts = []
+                for token in range(expert_ids.shape[0]):
+                    for expert_idx in range(expert_ids.shape[1]):
+                        expert_id = expert_ids[token][expert_idx].item()
+                        if expert_id not in activated_experts:
+                            activated_experts.append(expert_id)
+                print(f"[TORCH DEBUG] Activated experts: {sorted(activated_experts)}")
+                print(f"[TORCH DEBUG] First expert from expert_ids array: {expert_ids[0, 0].item()}")
+            print(f"expert_ids = {expert_ids}")
+            # print('expert ids:',expert_ids)
+            CPUInfer.submit(
+                moe.forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids.data_ptr(),
+                    weights.data_ptr(),
+                    input.data_ptr(),
+                    output.data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+
+            if enable_debug:
+                print(f"[AWQ-MOE DEBUG] AMX output[:{debug_print_count}] = {output.flatten()[:debug_print_count]}")
+
+            gate_proj = gate_projs[i % layer_num]
+            up_proj = up_projs[i % layer_num]
+            down_proj = down_projs[i % layer_num]
+            t_output = moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj, debug_print=enable_debug)
+            print("torch output", t_output)
+            print("amx output", output)
+
+            # print(output - t_output)
+            # print(torch.abs(output - t_output))
+            diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
+            # print(f'output_shape:{output.shape}, t_output_shape:{t_output.shape}\n')
+            print(f"Iteration {i}, diff = {diff:.6f}")
+
+            if enable_debug:
+                abs_diff = torch.abs(output - t_output)
+                print(f"[COMPARE] Max abs diff = {torch.max(abs_diff):.6f}")
+                print(f"[COMPARE] Mean abs diff = {torch.mean(abs_diff):.6f}")
+                print(f"[COMPARE] Relative diff = {diff:.6f}")
+                print("=" * 50)
+
+            if quant_mode == "int4" or quant_mode == "int4_1" or quant_mode == "int4_1k":
+                assert diff < 0.35
+            else:
+                assert diff < 0.05
+
+
+def test_moe_performance(quant_mode: str):
+    """
+    Test MOE inference performance (forward latency and throughput).
+
+    Measures:
+    - Forward pass latency (ms)
+    - Throughput (tokens/second)
+
+    Args:
+        quant_mode: Quantization mode, "bf16" or "int8"
+    """
+    import time
+
+    assert quant_mode in ("bf16", "int8"), f"Performance test only supports bf16 and int8, got {quant_mode}"
+
+    print(f"\n{'='*60}")
+    print(f"Performance Test - {quant_mode.upper()} mode (Inference)")
+    print(f"{'='*60}")
+    print(f"Configuration:")
+    print(f"  qlen (batch size): {perf_qlen}")
+    print(f"  warmup iterations: {perf_warmup_iter}")
+    print(f"  test iterations: {perf_test_iter}")
+    print(f"  num_threads: {num_threads}")
+    print(f"{'='*60}")
+
+    with torch.inference_mode(mode=True):
+        # Initialize weights
+        gate_proj = (
+            torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cuda")
+            .to("cpu")
+            .contiguous()
+        )
+        up_proj = (
+            torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cuda")
+            .to("cpu")
+            .contiguous()
+        )
+        down_proj = (
+            torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.bfloat16, device="cuda")
+            .to("cpu")
+            .contiguous()
+        )
+
+        # Create MOE config
+        config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size, 0)
+        config.max_len = max_len
+        config.gate_proj = gate_proj.data_ptr()
+        config.up_proj = up_proj.data_ptr()
+        config.down_proj = down_proj.data_ptr()
+        config.gate_scale = 0
+        config.pool = CPUInfer.backend_
+
+        # Create MOE instance based on quant_mode
+        if quant_mode == "bf16":
+            moe = kt_kernel_ext.moe.AMXBF16_MOE(config)
+        elif quant_mode == "int8":
+            moe = kt_kernel_ext.moe.AMXInt8_MOE(config)
+        else:
+            raise ValueError(f"Unsupported quant_mode for performance test: {quant_mode}")
+
+        print(f"[INFO] Using {quant_mode.upper()} MOE class")
+
+        # Load weights
+        CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+        CPUInfer.sync()
+
+        # Warm up task
+        if quant_mode == "bf16":
+            CPUInfer.submit(moe.warm_up_task())
+            CPUInfer.sync()
+
+        # Prepare test data
+        bsz_tensor = torch.tensor([perf_qlen], device="cpu")
+        expert_ids = torch.stack(
+            [torch.randperm(expert_num)[:num_experts_per_tok] for _ in range(perf_qlen)]
+        ).contiguous()
+        weights = torch.rand((perf_qlen, num_experts_per_tok), dtype=torch.float32).contiguous()
+        input_data = torch.randn((perf_qlen, hidden_size), dtype=torch.bfloat16).contiguous() / 100
+        output = torch.empty((perf_qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+
+        # =========================================================================
+        # Warmup Phase
+        # =========================================================================
+        print(f"\n[INFO] Warmup phase ({perf_warmup_iter} iterations)...")
+        for _ in range(perf_warmup_iter):
+            CPUInfer.submit(
+                moe.forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids.data_ptr(),
+                    weights.data_ptr(),
+                    input_data.data_ptr(),
+                    output.data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+
+        # =========================================================================
+        # Forward Performance Test
+        # =========================================================================
+        print(f"[INFO] Testing forward pass performance ({perf_test_iter} iterations)...")
+        forward_times = []
+        for _ in range(perf_test_iter):
+            start_time = time.perf_counter()
+            CPUInfer.submit(
+                moe.forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids.data_ptr(),
+                    weights.data_ptr(),
+                    input_data.data_ptr(),
+                    output.data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+            end_time = time.perf_counter()
+            forward_times.append((end_time - start_time) * 1000)  # Convert to ms
+
+        # =========================================================================
+        # Results Summary
+        # =========================================================================
+        import statistics
+
+        avg_forward = statistics.mean(forward_times)
+        std_forward = statistics.stdev(forward_times) if len(forward_times) > 1 else 0
+        min_forward = min(forward_times)
+        max_forward = max(forward_times)
+
+        # Calculate throughput (tokens per second)
+        forward_throughput = perf_qlen / (avg_forward / 1000)  # tokens/second
+
+        print(f"\n{'='*60}")
+        print(f"Performance Results - {quant_mode.upper()} mode (Inference)")
+        print(f"{'='*60}")
+        print(f"\nForward Pass:")
+        print(f"  Average latency: {avg_forward:.3f} ms (±{std_forward:.3f})")
+        print(f"  Min latency:     {min_forward:.3f} ms")
+        print(f"  Max latency:     {max_forward:.3f} ms")
+        print(f"  Throughput:      {forward_throughput:.1f} tokens/s")
+
+        print(f"\n[OK] Performance Test - {quant_mode.upper()} mode completed")
+
+        return {
+            "quant_mode": quant_mode,
+            "forward_avg_ms": avg_forward,
+            "forward_std_ms": std_forward,
+            "forward_throughput": forward_throughput,
+        }
+
+
+def run_performance_tests():
+    """Run performance tests for AMXBF16 and AMXINT8 modes (Inference)."""
+    print("\n" + "=" * 70)
+    print(" MOE AMX Inference Performance Test Suite")
+    print("=" * 70)
+    print(f"Configuration:")
+    print(f"  expert_num: {expert_num}")
+    print(f"  hidden_size: {hidden_size}")
+    print(f"  intermediate_size: {intermediate_size}")
+    print(f"  num_experts_per_tok: {num_experts_per_tok}")
+    print(f"  perf_qlen: {perf_qlen}")
+    print(f"  num_threads: {num_threads}")
+    print("=" * 70)
+
+    # Only test BF16 and INT8 as requested
+    quant_modes = ["bf16", "int8"]
+
+    results = []
+    try:
+        for quant_mode in quant_modes:
+            result = test_moe_performance(quant_mode)
+            results.append(result)
+
+        # Print comparison table
+        print("\n" + "=" * 70)
+        print(" Performance Comparison Summary (Inference)")
+        print("=" * 70)
+        print(f"\n{'Mode':<10} {'Forward(ms)':<15} {'Throughput(tok/s)':<20}")
+        print("-" * 45)
+        for r in results:
+            print(
+                f"{r['quant_mode'].upper():<10} " f"{r['forward_avg_ms']:<15.3f} " f"{r['forward_throughput']:<20.1f}"
+            )
+        print("-" * 45)
+
+        # Calculate speedup if we have both results
+        if len(results) == 2:
+            bf16_forward = results[0]["forward_avg_ms"]
+            int8_forward = results[1]["forward_avg_ms"]
+            speedup = bf16_forward / int8_forward
+            print(f"\nINT8 vs BF16 speedup: {speedup:.2f}x")
+
+        print("\n" + "=" * 70)
+        print(" PERFORMANCE TESTS COMPLETED!")
+        print("=" * 70)
+
+    except Exception as e:
+        print(f"\n[FAILED] Performance test failed with error: {e}")
+        import traceback
+
+        traceback.print_exc()
+        import sys
+
+        sys.exit(1)
+
+    return results
+
+
+def run_all_tests():
+    """Run all MOE accuracy tests for bf16 and int8 modes."""
+    print("\n" + "=" * 70)
+    print(" MOE AMX Inference Accuracy Test Suite")
+    print("=" * 70)
+    print(f"Configuration:")
+    print(f"  expert_num: {expert_num}")
+    print(f"  hidden_size: {hidden_size}")
+    print(f"  intermediate_size: {intermediate_size}")
+    print(f"  num_experts_per_tok: {num_experts_per_tok}")
+    print(f"  qlen: {qlen}")
+    print(f"  num_threads: {num_threads}")
+    print("=" * 70)
+
+    # Only test BF16 and INT8 as requested
+    quant_modes = ["bf16", "int8"]
+
+    try:
+        for quant_mode in quant_modes:
+            print(f"\n{'='*70}")
+            print(f" Testing MOE AMX - {quant_mode.upper()} Mode")
+            print(f"{'='*70}")
+            test_moe(quant_mode)
+
+        print("\n" + "=" * 70)
+        print(" ALL ACCURACY TESTS PASSED!")
+        print(f" Tested quantization modes: {', '.join(m.upper() for m in quant_modes)}")
+        print("=" * 70)
+
+    except Exception as e:
+        print(f"\n[FAILED] Test failed with error: {e}")
+        import traceback
+
+        traceback.print_exc()
+        import sys
+
+        sys.exit(1)
+
+
+# =============================================================================
+# Main Entry Point
+# =============================================================================
+
+if __name__ == "__main__":
+    import argparse
+    import sys
+
+    parser = argparse.ArgumentParser(description="MOE AMX Inference Test Suite")
+    parser.add_argument(
+        "--mode",
+        choices=["all", "accuracy", "perf"],
+        default="perf",
+        help="Test mode: 'all' runs both, 'accuracy' runs correctness tests, 'perf' runs performance tests",
+    )
+    parser.add_argument(
+        "--qlen",
+        type=int,
+        default=None,
+        help=f"Override perf_qlen for performance tests (default: {perf_qlen})",
+    )
+    parser.add_argument(
+        "--warmup",
+        type=int,
+        default=None,
+        help=f"Override warmup iterations for performance tests (default: {perf_warmup_iter})",
+    )
+    parser.add_argument(
+        "--iter",
+        type=int,
+        default=None,
+        help=f"Override test iterations for performance tests (default: {perf_test_iter})",
+    )
+    args = parser.parse_args()
+
+    # Override performance test parameters if specified
+    if args.qlen is not None or args.warmup is not None or args.iter is not None:
+        # Need to use global to modify module-level variables
+        if args.qlen is not None:
+            globals()["perf_qlen"] = args.qlen
+        if args.warmup is not None:
+            globals()["perf_warmup_iter"] = args.warmup
+        if args.iter is not None:
+            globals()["perf_test_iter"] = args.iter
+
+    if args.mode == "all":
+        run_all_tests()
+        run_performance_tests()
+    elif args.mode == "accuracy":
+        run_all_tests()
+    elif args.mode == "perf":
+        run_performance_tests()

kt-kernel/examples/test_moe_kernel.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""
+Description  :
+Author       : chenht2022
+Date         : 2024-07-25 10:32:05
+Version      : 1.0.0
+LastEditors  : chenht2022
+LastEditTime : 2024-08-06 10:38:05
+Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+"""
+import os, sys
+import time
+
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+os.environ["BLAS_NUM_THREADS"] = "1"
+import torch
+from kt_kernel import kt_kernel_ext
+
+
+expert_num = 16
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 4096
+num_experts_per_tok = 8
+m_block = 320
+n_block_up_gate = 32
+n_block_down = 64
+n_block_up_gate_prefi = 32
+n_block_down_prefi = 64
+# qlen = 1
+qlen = 1024
+layer_num = 1
+CPUInfer = kt_kernel_ext.CPUInfer(160)
+# validation_iter = 10000
+validation_iter = 1
+
+
+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 test_moe(quant_mode: str):
+    assert quant_mode == "int8" or quant_mode == "int4" or quant_mode == "int4_1"
+    with torch.inference_mode(mode=True):
+        moes = []
+        gate_projs = []
+        up_projs = []
+        down_projs = []
+        for _ in range(layer_num):
+            gate_proj = (
+                torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cpu")
+                .to("cpu")
+                .contiguous()
+            )
+            up_proj = (
+                torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cpu")
+                .to("cpu")
+                .contiguous()
+            )
+            down_proj = (
+                torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.bfloat16, device="cpu")
+                .to("cpu")
+                .contiguous()
+            )
+            config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size)
+            config.max_len = max_len
+            config.gate_proj = gate_proj.data_ptr()
+            config.up_proj = up_proj.data_ptr()
+            config.down_proj = down_proj.data_ptr()
+            config.pool = CPUInfer.backend_
+            if quant_mode == "int8":
+                d = kt_kernel_ext.moe.tiling.get_int8()
+                nbug_prefi = n_block_up_gate_prefi
+                nbd_prefi = n_block_down_prefi
+                kb = d["k_block"]
+                nb = d["n_block"]
+                mb = m_block
+                nbug = n_block_up_gate
+                nbd = n_block_down
+                print(
+                    f"Int8 Tiling: nbug {nbug}, nbd {nbd}, nb {nb}, mb {mb}, kb {kb}, nbug_prefi {nbug_prefi}, nbd_prefi {nbd_prefi}"
+                )
+                kt_kernel_ext.moe.tiling.set_int8(nbug, nbd, nb, mb, kb, nbug_prefi, nbd_prefi)
+                moe = kt_kernel_ext.moe.Int8_KERNEL_MOE(config)
+                CPUInfer.submit(moe.load_weights_task())
+                CPUInfer.sync()
+                # CPUInfer.submit(moe.warm_up_task())
+                # CPUInfer.sync()
+            elif quant_mode == "int4":
+                moe = kt_kernel_ext.moe.Int4_KERNEL_MOE(config)
+                CPUInfer.submit(moe.load_weights_task())
+                CPUInfer.sync()
+                CPUInfer.submit(moe.warm_up_task())
+                CPUInfer.sync()
+            else:
+                raise ValueError(f"Unsupported quantization mode: {quant_mode}")
+            gate_projs.append(gate_proj)
+            up_projs.append(up_proj)
+            down_projs.append(down_proj)
+            moes.append(moe)
+
+        # validation
+        for i in range(validation_iter):
+            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.rand((qlen, num_experts_per_tok), dtype=torch.float32).contiguous()
+            input = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            input = input / 100
+            # 打印 input 的内容
+            print("input:", input)
+            moe = moes[i % layer_num]
+            # print('expert ids:',expert_ids)
+            CPUInfer.submit(
+                moe.forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids.data_ptr(),
+                    weights.data_ptr(),
+                    input.data_ptr(),
+                    output.data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+            print("cpuinfer output", output)
+
+            gate_proj = gate_projs[i % layer_num]
+            up_proj = up_projs[i % layer_num]
+            down_proj = down_projs[i % layer_num]
+            t_output = moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj)
+            print("torch output", t_output)
+
+            # print(output - t_output)
+            diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
+            print("diff = ", diff)
+            if quant_mode == "int4":
+                assert diff < 0.35
+            else:
+                assert diff < 0.05
+
+
+test_moe("int8")
+# test_moe("int4")

kt-kernel/examples/test_moe_kml.py

@@ -0,0 +1,176 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""
+Description  :
+Author       : chenht2022
+Date         : 2024-07-25 10:32:05
+Version      : 1.0.0
+LastEditors  : chenht2022
+LastEditTime : 2024-08-06 10:38:05
+Copyright (c) 2024 by KVCache.AI, All Rights Reserved.
+"""
+import os, sys
+import time
+
+sys.path.insert(0, os.path.dirname(__file__) + "/../build")
+os.environ["BLAS_NUM_THREADS"] = "1"
+from kt_kernel import kt_kernel_ext
+import torch
+
+expert_num = 16
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 4096
+num_experts_per_tok = 8
+qlen = 512
+# qlen = 640
+layer_num = 1
+CPUInfer = kt_kernel_ext.CPUInfer(112)
+# validation_iter = 10000
+validation_iter = 1
+
+
+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 test_moe(quant_mode: str):
+    assert quant_mode == "bf16" or quant_mode == "int8" or quant_mode == "int4" or quant_mode == "int4_1"
+    with torch.inference_mode(mode=True):
+        moes = []
+        gate_projs = []
+        up_projs = []
+        down_projs = []
+        for _ in range(layer_num):
+            gate_proj = (
+                torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cpu")
+                .to("cpu")
+                .contiguous()
+            )
+            up_proj = (
+                torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.bfloat16, device="cpu")
+                .to("cpu")
+                .contiguous()
+            )
+            down_proj = (
+                torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.bfloat16, device="cpu")
+                .to("cpu")
+                .contiguous()
+            )
+            config = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size)
+            config.max_len = max_len
+            config.gate_proj = gate_proj.data_ptr()
+            config.up_proj = up_proj.data_ptr()
+            config.down_proj = down_proj.data_ptr()
+            config.pool = CPUInfer.backend_
+            if quant_mode == "bf16":
+                moe = kt_kernel_ext.moe.AMXBF16_MOE(config)
+                CPUInfer.submit(moe.load_weights_task())
+                CPUInfer.sync()
+                CPUInfer.submit(moe.warm_up_task())
+                CPUInfer.sync()
+            elif quant_mode == "int8":
+                moe = kt_kernel_ext.moe.KMLInt8_MOE(config)
+                CPUInfer.submit(moe.load_weights_task())
+                CPUInfer.sync()
+                # CPUInfer.submit(moe.warm_up_task())
+                # CPUInfer.sync()
+            elif quant_mode == "int4":
+                moe = kt_kernel_ext.moe.KMLInt4_MOE(config)
+                CPUInfer.submit(moe.load_weights_task())
+                CPUInfer.sync()
+                CPUInfer.submit(moe.warm_up_task())
+                CPUInfer.sync()
+            else:
+                raise ValueError(f"Unsupported quantization mode: {quant_mode}")
+            gate_projs.append(gate_proj)
+            up_projs.append(up_proj)
+            down_projs.append(down_proj)
+            moes.append(moe)
+
+        # validation
+        for i in range(validation_iter):
+            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.rand((qlen, num_experts_per_tok), dtype=torch.float32).contiguous()
+            input = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            input = input / 100
+            # 打印 input 的内容
+            print("input:", input)
+            moe = moes[i % layer_num]
+            # print('expert ids:',expert_ids)
+            CPUInfer.submit(
+                moe.forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids.data_ptr(),
+                    weights.data_ptr(),
+                    input.data_ptr(),
+                    output.data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+            print("cpuinfer output", output)
+
+            gate_proj = gate_projs[i % layer_num]
+            up_proj = up_projs[i % layer_num]
+            down_proj = down_projs[i % layer_num]
+            t_output = moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj)
+            print("torch output", t_output)
+
+            # print(output - t_output)
+            diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
+            print("diff = ", diff)
+            if quant_mode == "int4":
+                assert diff < 0.35
+            else:
+                assert diff < 0.05
+
+
+# test_moe("bf16")
+# test_moe("int8")
+test_moe("int4")

kt-kernel/examples/test_mxfp8_moe_avx2.py

@@ -0,0 +1,204 @@
+"""AVX2 MXFP8 MoE validation for MiniMax-M3-Preview.
+
+Forces the AVX2 backend (`kt_kernel_ext.moe.AVX2MXFP8_MOE`) for layer-`LAYER`
+experts, compares output against a torch dequant+matmul reference.
+
+Optional `--compare-amx` flag runs both AMX and AVX2 backends on identical
+inputs and asserts numerical equivalence. The two paths share buffer layout
+and do the same FMA arithmetic; observed differences should be at BF16-noise
+level (typical mean abs ~ 1e-4).
+
+Usage:
+    python test_mxfp8_moe_avx2.py --weight-path /mnt/data/models/Minimax-M3-preview
+    python test_mxfp8_moe_avx2.py --weight-path ... --compare-amx
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/build")
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/python")
+
+from kt_kernel import kt_kernel_ext  # noqa: E402
+from kt_kernel.utils.loader import MXFP8SafeTensorLoader  # noqa: E402
+
+
+# ---- Reference implementation (shared with test_mxfp8_moe_m3.py) ----
+
+def dequantize_mxfp8(weight_u8: torch.Tensor, scale_u8: torch.Tensor, group_size: int = 32) -> torch.Tensor:
+    """Dequantize [N, K] FP8 E4M3fn (as uint8) with [N, K/gs] ue8m0 scales -> [N, K] BF16."""
+    n, k = weight_u8.shape
+    assert k % group_size == 0
+    assert scale_u8.shape == (n, k // group_size)
+    w_fp32 = weight_u8.view(torch.float8_e4m3fn).float()
+    s_fp32 = (2.0 ** (scale_u8.to(torch.int32) - 127)).float()
+    s_full = s_fp32.repeat_interleave(group_size, dim=-1)
+    return (w_fp32 * s_full).to(torch.bfloat16).contiguous()
+
+
+def reference_mlp_m3(x, gate, up, down, alpha=1.702, limit=7.0):
+    g = torch.mm(x.float(), gate.float().t()).clamp(-limit, limit)
+    u = torch.mm(x.float(), up.float().t()).clamp(-limit, limit)
+    act = g * torch.sigmoid(g * alpha) * (u + 1.0)
+    return torch.mm(act, down.float().t())
+
+
+def reference_moe_m3(hidden, expert_ids, weights, gate_w, up_w, down_w, alpha=1.702, limit=7.0):
+    out = torch.zeros(hidden.shape[0], down_w.shape[1], dtype=torch.float32)
+    for tok in range(hidden.shape[0]):
+        for slot in range(expert_ids.shape[1]):
+            eid = int(expert_ids[tok, slot])
+            if eid < 0:
+                continue
+            w = float(weights[tok, slot])
+            y = reference_mlp_m3(hidden[tok:tok+1], gate_w[eid], up_w[eid], down_w[eid], alpha, limit)
+            out[tok] += w * y[0]
+    return out.to(hidden.dtype)
+
+
+# ---- Backend dispatch ----
+
+def _backend_cls(backend_name: str):
+    if backend_name == "amx":
+        cls = getattr(kt_kernel_ext.moe, "AMXMXFP8_KGroup_MOE", None)
+        if cls is None:
+            raise RuntimeError("AMXMXFP8_KGroup_MOE not in .so — rebuild with AVX-512 + VBMI + AMX")
+        return cls
+    if backend_name == "avx2":
+        cls = getattr(kt_kernel_ext.moe, "AVX2MXFP8_MOE", None)
+        if cls is None:
+            raise RuntimeError("AVX2MXFP8_MOE not in .so — rebuild with the AVX2 MXFP8 wiring (PR #2041 + this fix)")
+        return cls
+    raise ValueError(f"unknown backend {backend_name}")
+
+
+def run_backend(backend_name, weights, expert_num, top_k, hidden_size, intermediate_size,
+                layer_idx, qlen, cpu_threads, x, expert_ids, routing, physical_to_logical):
+    cpu_infer = kt_kernel_ext.CPUInfer(cpu_threads)
+    cfg = kt_kernel_ext.moe.MOEConfig(expert_num, top_k, hidden_size, intermediate_size, 0)
+    cfg.layer_idx = layer_idx
+    cfg.max_len = max(qlen, 1)
+    cfg.pool = cpu_infer.backend_
+    cfg.quant_config.bits = 8
+    cfg.quant_config.group_size = 32
+    cfg.quant_config.zero_point = False
+    cfg.swiglu_alpha = 1.702
+    cfg.swiglu_limit = 7.0
+
+    cfg.gate_projs = [[t.data_ptr() for t in weights["gate"]]]
+    cfg.up_projs = [[t.data_ptr() for t in weights["up"]]]
+    cfg.down_projs = [[t.data_ptr() for t in weights["down"]]]
+    cfg.gate_scales = [[t.data_ptr() for t in weights["gate_scale"]]]
+    cfg.up_scales = [[t.data_ptr() for t in weights["up_scale"]]]
+    cfg.down_scales = [[t.data_ptr() for t in weights["down_scale"]]]
+
+    moe = _backend_cls(backend_name)(cfg)
+    cpu_infer.submit(moe.load_weights_task(physical_to_logical.data_ptr()))
+    cpu_infer.sync()
+
+    bsz = torch.tensor([qlen], dtype=torch.int32)
+    y = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+    cpu_infer.submit(
+        moe.forward_task(bsz.data_ptr(), top_k, expert_ids.data_ptr(), routing.data_ptr(),
+                         x.data_ptr(), y.data_ptr(), False)
+    )
+    cpu_infer.sync()
+    return y
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description="MXFP8 MoE AVX2 backend test for MiniMax M3")
+    p.add_argument("--weight-path", required=True)
+    p.add_argument("--layer", type=int, default=3, help="Layer index (default 3 = first MoE layer)")
+    p.add_argument("--qlen", type=int, default=1)
+    p.add_argument("--top-k", type=int, default=4)
+    p.add_argument("--cpu-threads", type=int, default=32)
+    p.add_argument("--max-experts", type=int, default=0, help="Cap experts loaded (0=all)")
+    p.add_argument("--compare-amx", action="store_true",
+                   help="Also run AMX backend and assert numerical equivalence with AVX2.")
+    p.add_argument("--ref-threshold", type=float, default=0.10,
+                   help="Max relative error vs torch reference (default 10%).")
+    p.add_argument("--equiv-threshold", type=float, default=0.01,
+                   help="Max relative error AMX vs AVX2 (default 1%).")
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    torch.manual_seed(42)
+
+    print(f"[AVX2-MXFP8] Loading layer {args.layer} from {args.weight_path}")
+    loader = MXFP8SafeTensorLoader(args.weight_path)
+    weights = loader.load_experts(f"language_model.model.layers.{args.layer}")
+
+    expert_num = len(weights["gate"])
+    if args.max_experts and args.max_experts < expert_num:
+        for k in ("gate", "up", "down", "gate_scale", "up_scale", "down_scale"):
+            weights[k] = weights[k][:args.max_experts]
+        expert_num = args.max_experts
+
+    gate0 = weights["gate"][0]
+    intermediate_size = gate0.shape[0]
+    hidden_size = gate0.shape[1]
+    group_size = hidden_size // weights["gate_scale"][0].shape[1]
+    print(f"[AVX2-MXFP8] expert_num={expert_num} hidden={hidden_size} inter={intermediate_size} gs={group_size}")
+    assert group_size == 32, f"Expected group_size=32, got {group_size}"
+
+    physical_to_logical = torch.arange(expert_num, dtype=torch.int64).contiguous()
+
+    qlen = args.qlen
+    top_k = args.top_k
+    expert_ids = torch.stack([torch.randperm(expert_num)[:top_k] for _ in range(qlen)]).to(torch.int64).contiguous()
+    routing = torch.randn((qlen, top_k), dtype=torch.float32).abs().contiguous()
+    routing = routing / routing.sum(dim=-1, keepdim=True)
+    x = (torch.randn((qlen, hidden_size), dtype=torch.bfloat16) * 0.01).contiguous()
+
+    # ---- AVX2 forward ----
+    print("[AVX2-MXFP8] Running AVX2 backend...")
+    y_avx2 = run_backend("avx2", weights, expert_num, top_k, hidden_size, intermediate_size,
+                         args.layer, qlen, args.cpu_threads, x, expert_ids, routing,
+                         physical_to_logical)
+
+    # ---- Torch reference ----
+    print("[AVX2-MXFP8] Building torch reference (dequant+matmul)...")
+    gate_bf16 = torch.stack([dequantize_mxfp8(weights["gate"][i], weights["gate_scale"][i], group_size)
+                             for i in range(expert_num)])
+    up_bf16 = torch.stack([dequantize_mxfp8(weights["up"][i], weights["up_scale"][i], group_size)
+                           for i in range(expert_num)])
+    down_bf16 = torch.stack([dequantize_mxfp8(weights["down"][i], weights["down_scale"][i], group_size)
+                             for i in range(expert_num)])
+    y_ref = reference_moe_m3(x, expert_ids, routing, gate_bf16, up_bf16, down_bf16,
+                             alpha=1.702, limit=7.0)
+
+    diff_ref = (y_avx2.float() - y_ref.float()).abs()
+    ref_mag = y_ref.float().abs().mean() + 1e-12
+    rel_ref = diff_ref.mean() / ref_mag
+    print(f"[AVX2-MXFP8] vs ref:   mean abs={diff_ref.mean().item():.4e}  max abs={diff_ref.max().item():.4e}  rel={rel_ref.item()*100:.3f}%")
+    pass_ref = rel_ref.item() < args.ref_threshold
+
+    # ---- (optional) AMX vs AVX2 equivalence ----
+    pass_eq = True
+    if args.compare_amx:
+        print("[AVX2-MXFP8] Running AMX backend for equivalence check...")
+        y_amx = run_backend("amx", weights, expert_num, top_k, hidden_size, intermediate_size,
+                            args.layer, qlen, args.cpu_threads, x, expert_ids, routing,
+                            physical_to_logical)
+        diff_eq = (y_amx.float() - y_avx2.float()).abs()
+        amx_mag = y_amx.float().abs().mean() + 1e-12
+        rel_eq = diff_eq.mean() / amx_mag
+        print(f"[AVX2-MXFP8] vs AMX:  mean abs={diff_eq.mean().item():.4e}  max abs={diff_eq.max().item():.4e}  rel={rel_eq.item()*100:.3f}%")
+        pass_eq = rel_eq.item() < args.equiv_threshold
+
+    print(f"[AVX2-MXFP8] vs ref:  {'PASS' if pass_ref else 'FAIL'} (rel < {args.ref_threshold*100:.1f}%)")
+    if args.compare_amx:
+        print(f"[AVX2-MXFP8] vs AMX:  {'PASS' if pass_eq else 'FAIL'} (rel < {args.equiv_threshold*100:.1f}%)")
+    return 0 if (pass_ref and pass_eq) else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

kt-kernel/examples/test_mxfp8_moe_m3.py

@@ -0,0 +1,176 @@
+"""End-to-end MXFP8 MoE validation against the native MiniMax-M3-Preview ckpt.
+
+Loads layer-`LAYER_ID` experts via :class:`MXFP8SafeTensorLoader`, runs the AMX
+MXFP8 backend with swigluoai activation, and compares against a torch reference
+that dequantizes FP8 E4M3fn weights with ue8m0 group scales.
+
+Usage:
+    python test_mxfp8_moe_m3.py --weight-path /mnt/data/models/Minimax-M3-preview [--layer 3]
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/build")
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + "/python")
+
+from kt_kernel import kt_kernel_ext  # noqa: E402
+from kt_kernel.utils.loader import MXFP8SafeTensorLoader  # noqa: E402
+
+
+# ---- Reference implementation ----
+
+def dequantize_mxfp8(weight_u8: torch.Tensor, scale_u8: torch.Tensor, group_size: int = 32) -> torch.Tensor:
+    """Dequantize [N, K] FP8 E4M3fn (as uint8) with [N, K/gs] ue8m0 scales → [N, K] BF16."""
+    n, k = weight_u8.shape
+    assert k % group_size == 0
+    assert scale_u8.shape == (n, k // group_size)
+
+    w_fp32 = weight_u8.view(torch.float8_e4m3fn).float()
+    s_fp32 = (2.0 ** (scale_u8.to(torch.int32) - 127)).float()
+    s_full = s_fp32.repeat_interleave(group_size, dim=-1)
+    return (w_fp32 * s_full).to(torch.bfloat16).contiguous()
+
+
+def reference_mlp_m3(x: torch.Tensor, gate: torch.Tensor, up: torch.Tensor, down: torch.Tensor,
+                     alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor:
+    """Single-expert MLP with swigluoai activation."""
+    g = torch.mm(x.float(), gate.float().t()).clamp(-limit, limit)
+    u = torch.mm(x.float(), up.float().t()).clamp(-limit, limit)
+    act = g * torch.sigmoid(g * alpha) * (u + 1.0)
+    return torch.mm(act, down.float().t())
+
+
+def reference_moe_m3(hidden: torch.Tensor, expert_ids: torch.Tensor, weights: torch.Tensor,
+                     gate_w: torch.Tensor, up_w: torch.Tensor, down_w: torch.Tensor,
+                     alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor:
+    """Full MoE forward: route tokens to experts, compute, weighted sum."""
+    out = torch.zeros(hidden.shape[0], down_w.shape[1], dtype=torch.float32)
+    for tok in range(hidden.shape[0]):
+        for slot in range(expert_ids.shape[1]):
+            eid = int(expert_ids[tok, slot])
+            if eid < 0:
+                continue
+            w = float(weights[tok, slot])
+            y = reference_mlp_m3(hidden[tok:tok+1], gate_w[eid], up_w[eid], down_w[eid], alpha, limit)
+            out[tok] += w * y[0]
+    return out.to(hidden.dtype)
+
+
+# ---- Main ----
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="MXFP8 MoE E2E test for MiniMax M3")
+    p.add_argument("--weight-path", required=True, help="Path to Minimax-M3-preview safetensors directory.")
+    p.add_argument("--layer", type=int, default=3, help="Layer index to validate (default: 3, first MoE layer).")
+    p.add_argument("--qlen", type=int, default=1, help="Number of tokens to test.")
+    p.add_argument("--top-k", type=int, default=4, help="num_experts_per_tok (M3 default 4).")
+    p.add_argument("--cpu-threads", type=int, default=32)
+    p.add_argument("--max-experts", type=int, default=0, help="Cap number of experts loaded (0 = all).")
+    return p.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    torch.manual_seed(42)
+
+    print(f"[M3-MXFP8] Loading layer {args.layer} from {args.weight_path}")
+    loader = MXFP8SafeTensorLoader(args.weight_path)
+    weights = loader.load_experts(f"language_model.model.layers.{args.layer}")
+
+    expert_num = len(weights["gate"])
+    if args.max_experts and args.max_experts < expert_num:
+        for k in ("gate", "up", "down", "gate_scale", "up_scale", "down_scale"):
+            weights[k] = weights[k][:args.max_experts]
+        expert_num = args.max_experts
+    print(f"[M3-MXFP8] expert_num={expert_num}")
+
+    gate0 = weights["gate"][0]
+    down0 = weights["down"][0]
+    intermediate_size = gate0.shape[0]  # w1: [I, K]
+    hidden_size = gate0.shape[1]        # FP8: 1 byte/element, no /2
+    assert down0.shape == (hidden_size, intermediate_size), f"unexpected down shape {down0.shape}"
+
+    group_size = hidden_size // weights["gate_scale"][0].shape[1]
+    print(f"[M3-MXFP8] hidden={hidden_size} inter={intermediate_size} gs={group_size}")
+    assert group_size == 32, f"Expected group_size=32, got {group_size}"
+
+    physical_to_logical = torch.arange(expert_num, dtype=torch.int64).contiguous()
+
+    # ---- C++ MXFP8 forward ----
+    cpu_infer = kt_kernel_ext.CPUInfer(args.cpu_threads)
+    cfg = kt_kernel_ext.moe.MOEConfig(expert_num, args.top_k, hidden_size, intermediate_size, 0)
+    cfg.layer_idx = args.layer
+    cfg.max_len = max(args.qlen, 1)
+    cfg.pool = cpu_infer.backend_
+    cfg.quant_config.bits = 8
+    cfg.quant_config.group_size = group_size
+    cfg.quant_config.zero_point = False
+    cfg.swiglu_alpha = 1.702
+    cfg.swiglu_limit = 7.0
+
+    cfg.gate_projs = [[t.data_ptr() for t in weights["gate"]]]
+    cfg.up_projs = [[t.data_ptr() for t in weights["up"]]]
+    cfg.down_projs = [[t.data_ptr() for t in weights["down"]]]
+    cfg.gate_scales = [[t.data_ptr() for t in weights["gate_scale"]]]
+    cfg.up_scales = [[t.data_ptr() for t in weights["up_scale"]]]
+    cfg.down_scales = [[t.data_ptr() for t in weights["down_scale"]]]
+
+    moe = kt_kernel_ext.moe.AMXMXFP8_KGroup_MOE(cfg)
+    cpu_infer.submit(moe.load_weights_task(physical_to_logical.data_ptr()))
+    cpu_infer.sync()
+    print("[M3-MXFP8] C++ weights loaded")
+
+    qlen = args.qlen
+    top_k = args.top_k
+    bsz = torch.tensor([qlen], dtype=torch.int32)
+    expert_ids = torch.stack([torch.randperm(expert_num)[:top_k] for _ in range(qlen)]).to(torch.int64).contiguous()
+    routing = torch.randn((qlen, top_k), dtype=torch.float32).abs().contiguous()
+    routing = routing / routing.sum(dim=-1, keepdim=True)  # normalize weights
+    x = (torch.randn((qlen, hidden_size), dtype=torch.bfloat16) * 0.01).contiguous()
+    y_cpp = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+
+    cpu_infer.submit(
+        moe.forward_task(
+            bsz.data_ptr(), top_k, expert_ids.data_ptr(), routing.data_ptr(),
+            x.data_ptr(), y_cpp.data_ptr(), False,
+        )
+    )
+    cpu_infer.sync()
+    print("[M3-MXFP8] C++ forward done")
+
+    # ---- Torch reference ----
+    print("[M3-MXFP8] Building torch reference (dequantizing all loaded experts)...")
+    gate_bf16 = torch.stack([dequantize_mxfp8(weights["gate"][i], weights["gate_scale"][i], group_size)
+                             for i in range(expert_num)])
+    up_bf16 = torch.stack([dequantize_mxfp8(weights["up"][i], weights["up_scale"][i], group_size)
+                           for i in range(expert_num)])
+    down_bf16 = torch.stack([dequantize_mxfp8(weights["down"][i], weights["down_scale"][i], group_size)
+                             for i in range(expert_num)])
+
+    y_ref = reference_moe_m3(x, expert_ids, routing, gate_bf16, up_bf16, down_bf16,
+                             alpha=1.702, limit=7.0)
+
+    # ---- Compare ----
+    diff = (y_cpp.float() - y_ref.float()).abs()
+    ref_mag = y_ref.float().abs().mean() + 1e-12
+    rel = diff.mean() / ref_mag
+    print(f"[M3-MXFP8] mean abs diff = {diff.mean().item():.4e}")
+    print(f"[M3-MXFP8] max  abs diff = {diff.max().item():.4e}")
+    print(f"[M3-MXFP8] ref  mean mag = {ref_mag.item():.4e}")
+    print(f"[M3-MXFP8] rel mean diff = {rel.item()*100:.3f}%")
+    print(f"[M3-MXFP8] cpp[:8]  = {y_cpp.flatten()[:8].tolist()}")
+    print(f"[M3-MXFP8] ref[:8]  = {y_ref.flatten()[:8].tolist()}")
+
+    passed = rel.item() < 0.10
+    print(f"[M3-MXFP8] {'PASS' if passed else 'FAIL'} (threshold: 10%)")
+    return 0 if passed else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

kt-kernel/examples/test_rope.cpp

@@ -0,0 +1,123 @@
+#include <cassert>
+#include <iostream>
+#include <random>
+#include <vector>
+
+#include "../operators/rope.hpp"
+
+std::vector<float> create_random_vector(size_t total_size, std::vector<size_t> shape, unsigned int seed = 0) {
+  std::vector<float> vec(total_size);
+  std::mt19937 gen(seed == 0 ? std::random_device{}() : seed);
+  std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
+  // for (size_t i = 0; i < total_size; ++i) {
+  //   vec[i] = 1; // dist(gen);
+  // }
+  for (size_t i = 0; i < shape[0]; ++i) {
+    size_t offset_i = i * shape[1] * shape[2] * shape[3];
+    for (size_t j = 0; j < shape[1]; ++j) {
+      size_t offset_j = j * shape[2] * shape[3];
+      for (size_t k = 0; k < shape[2]; ++k) {
+        size_t offset_k = k * shape[3];
+        for (size_t a = 0; a < shape[3]; ++a) {
+          vec[offset_i + offset_j + offset_k + a] = a;
+        }
+      }
+    }
+  }
+  return vec;
+}
+
+void print_vector_to_file(const std::vector<float>& vec, const char* filename) {
+  FILE* fp = fopen(filename, "w");
+  for (auto x : vec) {
+    fprintf(fp, "%.2f ", x);
+  }
+  fclose(fp);
+}
+
+std::pair<std::vector<float>, std::vector<float>> cpp_torch_rope_with_apply_single(
+    const std::vector<float>& q_in_const, const std::vector<float>& k_in_const,
+    DeepseekV3YarnRotaryEmbedding<float>& rotary_emb, size_t B, size_t H, size_t S, size_t D_rope) {
+  rotary_emb.init(S);
+
+  const float* full_cos_cache_ptr = rotary_emb.cos();
+  const float* full_sin_cache_ptr = rotary_emb.sin();
+
+  std::vector<float> q_out = q_in_const;
+  std::vector<float> k_out = k_in_const;
+
+  size_t stride_head = S * D_rope;
+  size_t stride_batch = H * stride_head;
+
+  for (size_t b = 0; b < B; ++b) {
+    for (size_t h = 0; h < H; ++h) {
+      float* current_k_head_ptr = k_out.data() + b * stride_batch + h * stride_head;
+      Rope<DeepseekV3YarnRotaryEmbedding<float>, float>::apply_multiple(rotary_emb, current_k_head_ptr,
+                                                                        static_cast<int>(D_rope), 0, S);
+      for (size_t s = 0; s < S; ++s) {
+        float* current_q_head_ptr = q_out.data() + b * stride_batch + h * stride_head + s * D_rope;
+
+        Rope<DeepseekV3YarnRotaryEmbedding<float>, float>::apply_single(rotary_emb, current_q_head_ptr,
+                                                                        static_cast<int>(D_rope), s);
+      }
+    }
+  }
+
+  return {q_out, k_out};
+}
+
+int main() {
+  size_t batch_size = 2;
+  size_t num_heads = 16;
+  size_t seq_len = 32;
+  size_t rope_size = 16;
+  float theta = 10000.0f;
+
+  float beta_fast_cfg = 32.0f;
+  float beta_slow_cfg = 1.0f;
+  float factor_cfg = 40.0f;
+  float mscale_cfg = 1.0f;
+  float mscale_all_dim_cfg = 1.0f;
+  size_t original_max_pos_embeddings_cfg = 4096;
+
+  std::cout << "--- Test Parameters ---" << std::endl;
+  std::cout << "Batch Size: " << batch_size << std::endl;
+  std::cout << "Num Heads: " << num_heads << std::endl;
+  std::cout << "Seq Len: " << seq_len << std::endl;
+  std::cout << "Rope Size (dim): " << rope_size << std::endl;
+  std::cout << "Theta (base): " << theta << std::endl;
+  std::cout << "Scaling Factor: " << factor_cfg << std::endl;
+  std::cout << "Original Max Pos Embeddings: " << original_max_pos_embeddings_cfg << std::endl;
+  std::cout << "-----------------------" << std::endl << std::endl;
+
+  DeepseekV3YarnRotaryEmbedding<float> rotary_emb(rope_size, original_max_pos_embeddings_cfg, theta, factor_cfg,
+                                                  original_max_pos_embeddings_cfg, beta_fast_cfg, beta_slow_cfg,
+                                                  mscale_cfg, mscale_all_dim_cfg);
+  std::cout << "DeepseekV3YarnRotaryEmbedding instantiated." << std::endl;
+
+  size_t total_elements_per_tensor = batch_size * num_heads * seq_len * rope_size;
+
+  unsigned int q_seed = 123;
+  unsigned int k_seed = 456;
+  std::vector<float> q_pe_vec =
+      create_random_vector(total_elements_per_tensor, {batch_size, num_heads, seq_len, rope_size}, q_seed);
+  std::vector<float> k_pe_vec =
+      create_random_vector(total_elements_per_tensor, {batch_size, num_heads, seq_len, rope_size}, k_seed);
+
+  std::cout << "Input Q_PE and K_PE vectors created. Total elements per tensor: " << total_elements_per_tensor
+            << std::endl;
+
+  std::cout << std::endl;
+
+  std::cout << "Applying RoPE using cpp_torch_rope_with_apply_single..." << std::endl;
+  auto [q2_vec, k2_vec] =
+      cpp_torch_rope_with_apply_single(q_pe_vec, k_pe_vec, rotary_emb, batch_size, num_heads, seq_len, rope_size);
+  std::cout << "RoPE application finished." << std::endl << std::endl;
+
+  std::cout << std::endl << "test_rope.cpp finished successfully." << std::endl;
+
+  print_vector_to_file(q2_vec, "q_cpp.out");
+  print_vector_to_file(k2_vec, "k_cpp.out");
+
+  return 0;
+}

kt-kernel/examples/test_rope.py

@@ -0,0 +1,89 @@
+import torch
+from torch_attention import apply_rotary_pos_emb, DeepseekV3YarnRotaryEmbedding, DeepseekV3RotaryEmbedding
+
+batch_size  = 1
+num_heads   = 1
+seq_len     = 1024
+rope_size   = 64
+theta       = 10000
+
+max_position_embeddings =  163840
+
+scaling_cfg = {
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "factor": 40,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0,
+    "original_max_position_embeddings": 4096,
+    "type": "yarn"
+}
+
+rotary_emb = DeepseekV3YarnRotaryEmbedding(
+    rope_size,
+    max_position_embeddings=max_position_embeddings,
+    scaling_factor=scaling_cfg["factor"],
+    base=theta,
+    beta_fast=scaling_cfg["beta_fast"],
+    beta_slow=scaling_cfg["beta_slow"],
+    mscale=scaling_cfg["mscale"],
+    mscale_all_dim=scaling_cfg["mscale_all_dim"],
+    original_max_position_embeddings=scaling_cfg["original_max_position_embeddings"],
+)
+
+
+def load_fp16_tensor(file_path, shape):
+    with open(file_path, 'rb') as f:
+        raw_data = f.read()
+    tensor = torch.frombuffer(raw_data, dtype=torch.float16)
+    tensor = tensor.view(shape)  # 根据你的 shape reshape
+    return tensor
+
+def load_fp32_tensor(file_path, shape):
+    with open(file_path, 'rb') as f:
+        raw_data = f.read()
+    tensor = torch.frombuffer(raw_data, dtype=torch.float32)
+    tensor = tensor.view(shape)  # 根据你的 shape reshape
+    return tensor
+
+#q_pe = torch.randn(batch_size, num_heads, seq_len, rope_size, dtype=torch.float32)
+#k_pe = torch.randn_like(q_pe)
+
+q_pe = load_fp16_tensor("csrc/ktransformers_ext/build/before_rope",(batch_size, num_heads, seq_len, rope_size)) 
+# k_pe = torch.ones_like(q_pe) 
+k_pe = load_fp16_tensor("csrc/ktransformers_ext/build/before_rope",(batch_size, num_heads, seq_len, rope_size)) 
+print(q_pe)
+
+check = load_fp16_tensor("csrc/ktransformers_ext/build/after_rope",(batch_size, num_heads, seq_len, rope_size))
+
+
+
+
+def torch_rope(q, k):
+    cos, sin = rotary_emb(q, seq_len=seq_len)
+
+    cos_to_check = load_fp32_tensor("csrc/ktransformers_ext/build/cos",(seq_len, rope_size//2))
+    sin_to_check = load_fp32_tensor("csrc/ktransformers_ext/build/sin",(seq_len, rope_size//2))
+
+
+    
+
+    sin = sin.unsqueeze(0)
+    cos = cos.unsqueeze(0)
+    q2, k2 = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1)
+    return q2, k2
+
+q2, k2 = torch_rope(q_pe, k_pe)
+print(q2,k2)
+print(check)
+
+diff = torch.abs(q2 - check).max()
+
+
+print(diff)
+
+# print(q2,k2)
+
+# print_tensor(q2, 'q_py.out')
+# print_tensor(k2, 'k_py.out')
+

kt-kernel/examples/test_softmax.py

@@ -0,0 +1,24 @@
+
+import torch
+from torch import nn
+
+
+def load_fp16_tensor(file_path, shape):
+    with open(file_path, 'rb') as f:
+        raw_data = f.read()
+    tensor = torch.frombuffer(raw_data, dtype=torch.float16)
+    tensor = tensor.view(shape)  # 根据你的 shape reshape
+    return tensor
+
+a = load_fp16_tensor("csrc/ktransformers_ext/build/before_softmax", (64,1024))
+check = load_fp16_tensor("csrc/ktransformers_ext/build/after_softmax", (64,1024))
+
+
+a = nn.functional.softmax(a, dim=-1, dtype=torch.float16)
+diff = torch.abs(a - check).max()
+
+print(a)
+print(check)
+print(diff)
+
+

kt-kernel/examples/test_write_buffer.py

@@ -0,0 +1,765 @@
+"""
+Test write_weight_scale_to_buffer for AMX MOE operators.
+
+Supports:
+- FP8: FP8 weights (1 byte) + float32 scales (block-wise)
+- FP8_PERCHANNEL: FP8 weights (1 byte) + float32 per-channel scales
+- BF16: Native BF16 weights (2 bytes), no scales
+
+Usage:
+    python test_write_buffer.py          # Run all modes
+    python test_write_buffer.py fp8      # Run FP8 only
+    python test_write_buffer.py fp8_perchannel  # Run FP8 per-channel only
+    python test_write_buffer.py bf16     # Run BF16 only
+"""
+
+import os
+import sys
+import time
+
+import torch
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "build"))
+
+from kt_kernel import kt_kernel_ext
+from kt_kernel_ext import CPUInfer
+
+
+def make_cpu_infer(thread_num=80):
+    return CPUInfer(thread_num)
+
+
+def div_up(a, b):
+    return (a + b - 1) // b
+
+
+def build_config_fp8(cpuinfer, expert_num, num_experts_per_tok, hidden_size, intermediate_size, group_size):
+    cfg = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size)
+    cfg.max_len = 1
+    cfg.quant_config.bits = 8  # FP8
+    cfg.quant_config.group_size = group_size
+    cfg.quant_config.zero_point = False
+    cfg.pool = cpuinfer.backend_
+    return cfg
+
+
+def build_config_fp8_perchannel(cpuinfer, expert_num, num_experts_per_tok, hidden_size, intermediate_size):
+    cfg = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size)
+    cfg.max_len = 1
+    cfg.quant_config.bits = 8  # FP8
+    cfg.quant_config.group_size = 0  # Not used for per-channel
+    cfg.quant_config.zero_point = False
+    cfg.quant_config.per_channel = True
+    cfg.pool = cpuinfer.backend_
+    return cfg
+
+
+def build_config_bf16(cpuinfer, expert_num, num_experts_per_tok, hidden_size, intermediate_size):
+    cfg = kt_kernel_ext.moe.MOEConfig(expert_num, num_experts_per_tok, hidden_size, intermediate_size)
+    cfg.max_len = 1
+    cfg.pool = cpuinfer.backend_
+    return cfg
+
+
+def allocate_weights_fp8(expert_num, hidden_size, intermediate_size, group_size):
+    """Allocate FP8 weights and scales for testing"""
+    # FP8 weights: 1 byte per element
+    per_mat_weight_bytes = hidden_size * intermediate_size
+    # FP8 scales: block-wise (group_size x group_size blocks), stored as float32
+    n_blocks_n_gate_up = div_up(intermediate_size, group_size)
+    n_blocks_k = div_up(hidden_size, group_size)
+    per_mat_scale_elems_gate_up = n_blocks_n_gate_up * n_blocks_k
+
+    # For down: n=hidden_size, k=intermediate_size
+    n_blocks_n_down = n_blocks_k
+    n_blocks_k_down = n_blocks_n_gate_up
+    per_mat_scale_elems_down = n_blocks_n_down * n_blocks_k_down
+
+    gate_q = torch.randint(0, 256, (expert_num * per_mat_weight_bytes,), dtype=torch.uint8)
+    up_q = torch.randint(0, 256, (expert_num * per_mat_weight_bytes,), dtype=torch.uint8)
+    down_q = torch.randint(0, 256, (expert_num * per_mat_weight_bytes,), dtype=torch.uint8)
+
+    gate_scale = torch.randn(expert_num * per_mat_scale_elems_gate_up, dtype=torch.float32)
+    up_scale = torch.randn(expert_num * per_mat_scale_elems_gate_up, dtype=torch.float32)
+    down_scale = torch.randn(expert_num * per_mat_scale_elems_down, dtype=torch.float32)
+
+    return {
+        "gate_q": gate_q,
+        "up_q": up_q,
+        "down_q": down_q,
+        "gate_scale": gate_scale,
+        "up_scale": up_scale,
+        "down_scale": down_scale,
+        "per_mat_weight_bytes": per_mat_weight_bytes,
+        "per_mat_scale_elems_gate_up": per_mat_scale_elems_gate_up,
+        "per_mat_scale_elems_down": per_mat_scale_elems_down,
+    }
+
+
+def allocate_weights_fp8_perchannel(expert_num, hidden_size, intermediate_size):
+    """Allocate FP8 per-channel weights and scales for testing"""
+    per_mat_weight_bytes = hidden_size * intermediate_size
+    per_mat_scale_elems_gate_up = intermediate_size  # one scale per output channel
+    per_mat_scale_elems_down = hidden_size
+
+    gate_q = torch.randint(0, 256, (expert_num * per_mat_weight_bytes,), dtype=torch.uint8)
+    up_q = torch.randint(0, 256, (expert_num * per_mat_weight_bytes,), dtype=torch.uint8)
+    down_q = torch.randint(0, 256, (expert_num * per_mat_weight_bytes,), dtype=torch.uint8)
+
+    gate_scale = torch.randn(expert_num * per_mat_scale_elems_gate_up, dtype=torch.float32)
+    up_scale = torch.randn(expert_num * per_mat_scale_elems_gate_up, dtype=torch.float32)
+    down_scale = torch.randn(expert_num * per_mat_scale_elems_down, dtype=torch.float32)
+
+    return {
+        "gate_q": gate_q,
+        "up_q": up_q,
+        "down_q": down_q,
+        "gate_scale": gate_scale,
+        "up_scale": up_scale,
+        "down_scale": down_scale,
+        "per_mat_weight_bytes": per_mat_weight_bytes,
+        "per_mat_scale_elems_gate_up": per_mat_scale_elems_gate_up,
+        "per_mat_scale_elems_down": per_mat_scale_elems_down,
+    }
+
+
+def allocate_weights_bf16(expert_num, hidden_size, intermediate_size):
+    """Allocate BF16 weights for testing (no scales)"""
+    # BF16 weights: 2 bytes per element
+    per_mat_weight_elems = hidden_size * intermediate_size
+    per_mat_weight_bytes = per_mat_weight_elems * 2  # BF16 = 2 bytes
+
+    gate_proj = torch.randn(expert_num * per_mat_weight_elems, dtype=torch.bfloat16)
+    up_proj = torch.randn(expert_num * per_mat_weight_elems, dtype=torch.bfloat16)
+    down_proj = torch.randn(expert_num * per_mat_weight_elems, dtype=torch.bfloat16)
+
+    return {
+        "gate_proj": gate_proj,
+        "up_proj": up_proj,
+        "down_proj": down_proj,
+        "per_mat_weight_bytes": per_mat_weight_bytes,
+        "per_mat_weight_elems": per_mat_weight_elems,
+    }
+
+
+def test_fp8_write_buffer(gpu_tp_count):
+    """Test write_weight_scale_to_buffer with FP8 weights"""
+    torch.manual_seed(123)
+
+    expert_num = 256
+    gpu_experts = expert_num
+    num_experts_per_tok = 8
+    hidden_size = 3072
+    intermediate_size = 1536
+    group_size = 128
+
+    cpuinfer = make_cpu_infer()
+    cfg = build_config_fp8(cpuinfer, expert_num, num_experts_per_tok, hidden_size, intermediate_size, group_size)
+    weights = allocate_weights_fp8(expert_num, hidden_size, intermediate_size, group_size)
+
+    cfg.gate_proj = weights["gate_q"].data_ptr()
+    cfg.up_proj = weights["up_q"].data_ptr()
+    cfg.down_proj = weights["down_q"].data_ptr()
+    cfg.gate_scale = weights["gate_scale"].data_ptr()
+    cfg.up_scale = weights["up_scale"].data_ptr()
+    cfg.down_scale = weights["down_scale"].data_ptr()
+
+    moe = kt_kernel_ext.moe.AMXFP8_MOE(cfg)
+
+    physical_to_logical_map = torch.arange(expert_num, dtype=torch.int64, device="cpu").contiguous()
+    cpuinfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+    cpuinfer.sync()
+
+    per_mat_weight_bytes = weights["per_mat_weight_bytes"]
+    per_mat_scale_elems_gate_up = weights["per_mat_scale_elems_gate_up"]
+    per_mat_scale_elems_down = weights["per_mat_scale_elems_down"]
+
+    # Calculate sizes per TP part
+    weight_bytes_per_expert_per_tp = per_mat_weight_bytes // gpu_tp_count
+    gpu_n_w13 = intermediate_size // gpu_tp_count
+    gpu_k_w13 = hidden_size
+    scale_elems_per_expert_per_tp_gate_up = div_up(gpu_n_w13, group_size) * div_up(gpu_k_w13, group_size)
+    gpu_n_w2 = hidden_size
+    gpu_k_w2 = intermediate_size // gpu_tp_count
+    scale_elems_per_expert_per_tp_down = div_up(gpu_n_w2, group_size) * div_up(gpu_k_w2, group_size)
+
+    total_weight_bytes_per_tp = gpu_experts * weight_bytes_per_expert_per_tp
+    total_scale_elems_per_tp_gate_up = gpu_experts * scale_elems_per_expert_per_tp_gate_up
+    total_scale_elems_per_tp_down = gpu_experts * scale_elems_per_expert_per_tp_down
+
+    # Create buffer lists
+    w13_weight_bufs = [torch.empty(2 * total_weight_bytes_per_tp, dtype=torch.uint8) for _ in range(gpu_tp_count)]
+    w13_scale_bufs = [
+        torch.empty(2 * total_scale_elems_per_tp_gate_up, dtype=torch.float32) for _ in range(gpu_tp_count)
+    ]
+    w2_weight_bufs = [torch.empty(total_weight_bytes_per_tp, dtype=torch.uint8) for _ in range(gpu_tp_count)]
+    w2_scale_bufs = [torch.empty(total_scale_elems_per_tp_down, dtype=torch.float32) for _ in range(gpu_tp_count)]
+
+    print(f"[FP8] GPU TP count: {gpu_tp_count}, Experts: {expert_num}")
+    print(f"[FP8] Weight bytes per expert per TP: {weight_bytes_per_expert_per_tp}")
+    print(f"[FP8] Scale elements per expert per TP (gate/up): {scale_elems_per_expert_per_tp_gate_up}")
+
+    def get_expert_ptrs(expert_id):
+        w13_weight_ptrs = []
+        w13_scale_ptrs = []
+        w2_weight_ptrs = []
+        w2_scale_ptrs = []
+        for tp_idx in range(gpu_tp_count):
+            w13_weight_expert_offset = expert_id * 2 * weight_bytes_per_expert_per_tp
+            w13_scale_expert_offset = expert_id * 2 * scale_elems_per_expert_per_tp_gate_up
+            w2_weight_expert_offset = expert_id * weight_bytes_per_expert_per_tp
+            w2_scale_expert_offset = expert_id * scale_elems_per_expert_per_tp_down
+
+            w13_weight_ptrs.append(w13_weight_bufs[tp_idx].data_ptr() + w13_weight_expert_offset)
+            w13_scale_ptrs.append(w13_scale_bufs[tp_idx].data_ptr() + w13_scale_expert_offset * 4)
+            w2_weight_ptrs.append(w2_weight_bufs[tp_idx].data_ptr() + w2_weight_expert_offset)
+            w2_scale_ptrs.append(w2_scale_bufs[tp_idx].data_ptr() + w2_scale_expert_offset * 4)
+        return w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs
+
+    # Warm up
+    for _ in range(2):
+        for expert_id in range(gpu_experts):
+            w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs = get_expert_ptrs(expert_id)
+            cpuinfer.submit(
+                moe.write_weight_scale_to_buffer_task(
+                    gpu_tp_count=gpu_tp_count,
+                    expert_id=expert_id,
+                    w13_weight_ptrs=w13_weight_ptrs,
+                    w13_scale_ptrs=w13_scale_ptrs,
+                    w2_weight_ptrs=w2_weight_ptrs,
+                    w2_scale_ptrs=w2_scale_ptrs,
+                )
+            )
+            cpuinfer.sync()
+
+    # Timing
+    begin_time = time.perf_counter_ns()
+    for expert_id in range(gpu_experts):
+        w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs = get_expert_ptrs(expert_id)
+        cpuinfer.submit(
+            moe.write_weight_scale_to_buffer_task(
+                gpu_tp_count=gpu_tp_count,
+                expert_id=expert_id,
+                w13_weight_ptrs=w13_weight_ptrs,
+                w13_scale_ptrs=w13_scale_ptrs,
+                w2_weight_ptrs=w2_weight_ptrs,
+                w2_scale_ptrs=w2_scale_ptrs,
+            )
+        )
+        cpuinfer.sync()
+    end_time = time.perf_counter_ns()
+    elapsed_ms = (end_time - begin_time) / 1e6
+
+    total_bytes = (
+        hidden_size * intermediate_size * gpu_experts * 3
+        + (per_mat_scale_elems_gate_up * 2 + per_mat_scale_elems_down) * gpu_experts * 4
+    )
+    print(f"[FP8] write_weight_scale_to_buffer time: {elapsed_ms:.2f} ms")
+    print(f"[FP8] Throughput: {total_bytes / (elapsed_ms * 1e6):.2f} GB/s")
+
+    # Verify correctness
+    def split_expert_tensor(tensor, chunk):
+        return [tensor[i * chunk : (i + 1) * chunk] for i in range(expert_num)]
+
+    gate_q = weights["gate_q"]
+    up_q = weights["up_q"]
+    down_q = weights["down_q"]
+    gate_scale = weights["gate_scale"]
+    up_scale = weights["up_scale"]
+    down_scale = weights["down_scale"]
+
+    gate_q_experts = split_expert_tensor(gate_q, per_mat_weight_bytes)
+    up_q_experts = split_expert_tensor(up_q, per_mat_weight_bytes)
+    down_q_experts = split_expert_tensor(down_q, per_mat_weight_bytes)
+    gate_scale_experts = split_expert_tensor(gate_scale, per_mat_scale_elems_gate_up)
+    up_scale_experts = split_expert_tensor(up_scale, per_mat_scale_elems_gate_up)
+    down_scale_experts = split_expert_tensor(down_scale, per_mat_scale_elems_down)
+
+    n_blocks_n = div_up(hidden_size, group_size)
+    n_blocks_k = div_up(intermediate_size, group_size)
+    n_blocks_k_per_tp = n_blocks_k // gpu_tp_count
+
+    for tp_idx in range(gpu_tp_count):
+        expected_w13_weights = []
+        expected_w13_scales = []
+        expected_w2_weights = []
+        expected_w2_scales = []
+
+        weight13_per_tp = per_mat_weight_bytes // gpu_tp_count
+        scale13_per_tp = per_mat_scale_elems_gate_up // gpu_tp_count
+
+        for expert_id in range(gpu_experts):
+            start_weight = tp_idx * weight13_per_tp
+            end_weight = (tp_idx + 1) * weight13_per_tp
+            start_scale = tp_idx * scale13_per_tp
+            end_scale = (tp_idx + 1) * scale13_per_tp
+
+            gate_weight_tp = gate_q_experts[expert_id][start_weight:end_weight]
+            gate_scale_tp = gate_scale_experts[expert_id][start_scale:end_scale]
+            up_weight_tp = up_q_experts[expert_id][start_weight:end_weight]
+            up_scale_tp = up_scale_experts[expert_id][start_scale:end_scale]
+
+            down_weight_tp_parts = []
+            down_scale_tp_parts = []
+            tp_slice_weight_size = intermediate_size // gpu_tp_count
+
+            for row_idx in range(hidden_size):
+                row_weight_start = row_idx * intermediate_size
+                tp_weight_offset = row_weight_start + tp_idx * tp_slice_weight_size
+                down_weight_tp_parts.append(
+                    down_q_experts[expert_id][tp_weight_offset : tp_weight_offset + tp_slice_weight_size]
+                )
+
+            for bn in range(n_blocks_n):
+                row_scale_start = bn * n_blocks_k
+                tp_scale_offset = row_scale_start + tp_idx * n_blocks_k_per_tp
+                down_scale_tp_parts.append(
+                    down_scale_experts[expert_id][tp_scale_offset : tp_scale_offset + n_blocks_k_per_tp]
+                )
+
+            down_weight_tp = torch.cat(down_weight_tp_parts)
+            down_scale_tp = torch.cat(down_scale_tp_parts)
+
+            expected_w13_weights.append(gate_weight_tp)
+            expected_w13_weights.append(up_weight_tp)
+            expected_w13_scales.append(gate_scale_tp)
+            expected_w13_scales.append(up_scale_tp)
+            expected_w2_weights.append(down_weight_tp)
+            expected_w2_scales.append(down_scale_tp)
+
+        expected_w13_weight = torch.cat(expected_w13_weights)
+        expected_w13_scale = torch.cat(expected_w13_scales)
+        expected_w2_weight = torch.cat(expected_w2_weights)
+        expected_w2_scale = torch.cat(expected_w2_scales)
+
+        if not torch.equal(w13_weight_bufs[tp_idx], expected_w13_weight):
+            diff_mask = w13_weight_bufs[tp_idx] != expected_w13_weight
+            first_diff_idx = diff_mask.nonzero()[0].item() if diff_mask.any() else -1
+            raise AssertionError(f"[FP8] w13 weight mismatch for TP {tp_idx} at index {first_diff_idx}")
+
+        if not torch.allclose(w13_scale_bufs[tp_idx], expected_w13_scale):
+            raise AssertionError(f"[FP8] w13 scale mismatch for TP {tp_idx}")
+
+        if not torch.equal(w2_weight_bufs[tp_idx], expected_w2_weight):
+            diff_mask = w2_weight_bufs[tp_idx] != expected_w2_weight
+            first_diff_idx = diff_mask.nonzero()[0].item() if diff_mask.any() else -1
+            raise AssertionError(f"[FP8] w2 weight mismatch for TP {tp_idx} at index {first_diff_idx}")
+
+        if not torch.allclose(w2_scale_bufs[tp_idx], expected_w2_scale):
+            raise AssertionError(f"[FP8] w2 scale mismatch for TP {tp_idx}")
+
+    print(f"[FP8] TP={gpu_tp_count} PASSED (verified {gpu_experts} experts across {gpu_tp_count} TP parts)")
+    return True
+
+
+def test_fp8_perchannel_write_buffer(gpu_tp_count):
+    """Test write_weight_scale_to_buffer with FP8 per-channel weights"""
+    torch.manual_seed(123)
+
+    expert_num = 256
+    gpu_experts = expert_num
+    num_experts_per_tok = 8
+    hidden_size = 3072
+    intermediate_size = 1536
+
+    cpuinfer = make_cpu_infer()
+    cfg = build_config_fp8_perchannel(cpuinfer, expert_num, num_experts_per_tok, hidden_size, intermediate_size)
+    weights = allocate_weights_fp8_perchannel(expert_num, hidden_size, intermediate_size)
+
+    cfg.gate_proj = weights["gate_q"].data_ptr()
+    cfg.up_proj = weights["up_q"].data_ptr()
+    cfg.down_proj = weights["down_q"].data_ptr()
+    cfg.gate_scale = weights["gate_scale"].data_ptr()
+    cfg.up_scale = weights["up_scale"].data_ptr()
+    cfg.down_scale = weights["down_scale"].data_ptr()
+
+    moe = kt_kernel_ext.moe.AMXFP8PerChannel_MOE(cfg)
+
+    physical_to_logical_map = torch.arange(expert_num, dtype=torch.int64, device="cpu").contiguous()
+    cpuinfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+    cpuinfer.sync()
+
+    per_mat_weight_bytes = weights["per_mat_weight_bytes"]
+    per_mat_scale_elems_gate_up = weights["per_mat_scale_elems_gate_up"]
+    per_mat_scale_elems_down = weights["per_mat_scale_elems_down"]
+
+    weight_bytes_per_expert_per_tp = per_mat_weight_bytes // gpu_tp_count
+    gpu_n_w13 = intermediate_size // gpu_tp_count
+    scale_elems_per_expert_per_tp_gate_up = gpu_n_w13
+    scale_elems_per_expert_per_tp_down = per_mat_scale_elems_down
+
+    total_weight_bytes_per_tp = gpu_experts * weight_bytes_per_expert_per_tp
+    total_scale_elems_per_tp_gate_up = gpu_experts * scale_elems_per_expert_per_tp_gate_up
+    total_scale_elems_per_tp_down = gpu_experts * scale_elems_per_expert_per_tp_down
+
+    w13_weight_bufs = [torch.empty(2 * total_weight_bytes_per_tp, dtype=torch.uint8) for _ in range(gpu_tp_count)]
+    w13_scale_bufs = [
+        torch.empty(2 * total_scale_elems_per_tp_gate_up, dtype=torch.float32) for _ in range(gpu_tp_count)
+    ]
+    w2_weight_bufs = [torch.empty(total_weight_bytes_per_tp, dtype=torch.uint8) for _ in range(gpu_tp_count)]
+    w2_scale_bufs = [torch.empty(total_scale_elems_per_tp_down, dtype=torch.float32) for _ in range(gpu_tp_count)]
+
+    print(f"[FP8_PERCHANNEL] GPU TP count: {gpu_tp_count}, Experts: {expert_num}")
+    print(f"[FP8_PERCHANNEL] Weight bytes per expert per TP: {weight_bytes_per_expert_per_tp}")
+    print(f"[FP8_PERCHANNEL] Scale elements per expert per TP (gate/up): {scale_elems_per_expert_per_tp_gate_up}")
+
+    def get_expert_ptrs(expert_id):
+        w13_weight_ptrs = []
+        w13_scale_ptrs = []
+        w2_weight_ptrs = []
+        w2_scale_ptrs = []
+        for tp_idx in range(gpu_tp_count):
+            w13_weight_expert_offset = expert_id * 2 * weight_bytes_per_expert_per_tp
+            w13_scale_expert_offset = expert_id * 2 * scale_elems_per_expert_per_tp_gate_up
+            w2_weight_expert_offset = expert_id * weight_bytes_per_expert_per_tp
+            w2_scale_expert_offset = expert_id * scale_elems_per_expert_per_tp_down
+
+            w13_weight_ptrs.append(w13_weight_bufs[tp_idx].data_ptr() + w13_weight_expert_offset)
+            w13_scale_ptrs.append(w13_scale_bufs[tp_idx].data_ptr() + w13_scale_expert_offset * 4)
+            w2_weight_ptrs.append(w2_weight_bufs[tp_idx].data_ptr() + w2_weight_expert_offset)
+            w2_scale_ptrs.append(w2_scale_bufs[tp_idx].data_ptr() + w2_scale_expert_offset * 4)
+        return w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs
+
+    for _ in range(2):
+        for expert_id in range(gpu_experts):
+            w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs = get_expert_ptrs(expert_id)
+            cpuinfer.submit(
+                moe.write_weight_scale_to_buffer_task(
+                    gpu_tp_count=gpu_tp_count,
+                    expert_id=expert_id,
+                    w13_weight_ptrs=w13_weight_ptrs,
+                    w13_scale_ptrs=w13_scale_ptrs,
+                    w2_weight_ptrs=w2_weight_ptrs,
+                    w2_scale_ptrs=w2_scale_ptrs,
+                )
+            )
+            cpuinfer.sync()
+
+    begin_time = time.perf_counter_ns()
+    for expert_id in range(gpu_experts):
+        w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs = get_expert_ptrs(expert_id)
+        cpuinfer.submit(
+            moe.write_weight_scale_to_buffer_task(
+                gpu_tp_count=gpu_tp_count,
+                expert_id=expert_id,
+                w13_weight_ptrs=w13_weight_ptrs,
+                w13_scale_ptrs=w13_scale_ptrs,
+                w2_weight_ptrs=w2_weight_ptrs,
+                w2_scale_ptrs=w2_scale_ptrs,
+            )
+        )
+        cpuinfer.sync()
+    end_time = time.perf_counter_ns()
+    elapsed_ms = (end_time - begin_time) / 1e6
+
+    total_bytes = (
+        hidden_size * intermediate_size * gpu_experts * 3
+        + (per_mat_scale_elems_gate_up * 2 + per_mat_scale_elems_down) * gpu_experts * 4
+    )
+    print(f"[FP8_PERCHANNEL] write_weight_scale_to_buffer time: {elapsed_ms:.2f} ms")
+    print(f"[FP8_PERCHANNEL] Throughput: {total_bytes / (elapsed_ms * 1e6):.2f} GB/s")
+
+    def split_expert_tensor(tensor, chunk):
+        return [tensor[i * chunk : (i + 1) * chunk] for i in range(expert_num)]
+
+    gate_q = weights["gate_q"]
+    up_q = weights["up_q"]
+    down_q = weights["down_q"]
+    gate_scale = weights["gate_scale"]
+    up_scale = weights["up_scale"]
+    down_scale = weights["down_scale"]
+
+    gate_q_experts = split_expert_tensor(gate_q, per_mat_weight_bytes)
+    up_q_experts = split_expert_tensor(up_q, per_mat_weight_bytes)
+    down_q_experts = split_expert_tensor(down_q, per_mat_weight_bytes)
+    gate_scale_experts = split_expert_tensor(gate_scale, per_mat_scale_elems_gate_up)
+    up_scale_experts = split_expert_tensor(up_scale, per_mat_scale_elems_gate_up)
+    down_scale_experts = split_expert_tensor(down_scale, per_mat_scale_elems_down)
+
+    for tp_idx in range(gpu_tp_count):
+        expected_w13_weights = []
+        expected_w13_scales = []
+        expected_w2_weights = []
+        expected_w2_scales = []
+
+        weight13_per_tp = per_mat_weight_bytes // gpu_tp_count
+        scale13_per_tp = per_mat_scale_elems_gate_up // gpu_tp_count
+
+        for expert_id in range(gpu_experts):
+            start_weight = tp_idx * weight13_per_tp
+            end_weight = (tp_idx + 1) * weight13_per_tp
+            start_scale = tp_idx * scale13_per_tp
+            end_scale = (tp_idx + 1) * scale13_per_tp
+
+            gate_weight_tp = gate_q_experts[expert_id][start_weight:end_weight]
+            gate_scale_tp = gate_scale_experts[expert_id][start_scale:end_scale]
+            up_weight_tp = up_q_experts[expert_id][start_weight:end_weight]
+            up_scale_tp = up_scale_experts[expert_id][start_scale:end_scale]
+
+            down_weight_tp_parts = []
+            tp_slice_weight_size = intermediate_size // gpu_tp_count
+
+            for row_idx in range(hidden_size):
+                row_weight_start = row_idx * intermediate_size
+                tp_weight_offset = row_weight_start + tp_idx * tp_slice_weight_size
+                down_weight_tp_parts.append(
+                    down_q_experts[expert_id][tp_weight_offset : tp_weight_offset + tp_slice_weight_size]
+                )
+
+            down_weight_tp = torch.cat(down_weight_tp_parts)
+            down_scale_tp = down_scale_experts[expert_id]
+
+            expected_w13_weights.append(gate_weight_tp)
+            expected_w13_weights.append(up_weight_tp)
+            expected_w13_scales.append(gate_scale_tp)
+            expected_w13_scales.append(up_scale_tp)
+            expected_w2_weights.append(down_weight_tp)
+            expected_w2_scales.append(down_scale_tp)
+
+        expected_w13_weight = torch.cat(expected_w13_weights)
+        expected_w13_scale = torch.cat(expected_w13_scales)
+        expected_w2_weight = torch.cat(expected_w2_weights)
+        expected_w2_scale = torch.cat(expected_w2_scales)
+
+        if not torch.equal(w13_weight_bufs[tp_idx], expected_w13_weight):
+            diff_mask = w13_weight_bufs[tp_idx] != expected_w13_weight
+            first_diff_idx = diff_mask.nonzero()[0].item() if diff_mask.any() else -1
+            raise AssertionError(f"[FP8_PERCHANNEL] w13 weight mismatch for TP {tp_idx} at index {first_diff_idx}")
+
+        if not torch.allclose(w13_scale_bufs[tp_idx], expected_w13_scale):
+            raise AssertionError(f"[FP8_PERCHANNEL] w13 scale mismatch for TP {tp_idx}")
+
+        if not torch.equal(w2_weight_bufs[tp_idx], expected_w2_weight):
+            diff_mask = w2_weight_bufs[tp_idx] != expected_w2_weight
+            first_diff_idx = diff_mask.nonzero()[0].item() if diff_mask.any() else -1
+            raise AssertionError(f"[FP8_PERCHANNEL] w2 weight mismatch for TP {tp_idx} at index {first_diff_idx}")
+
+        if not torch.allclose(w2_scale_bufs[tp_idx], expected_w2_scale):
+            raise AssertionError(f"[FP8_PERCHANNEL] w2 scale mismatch for TP {tp_idx}")
+
+    print(f"[FP8_PERCHANNEL] TP={gpu_tp_count} PASSED (verified {gpu_experts} experts across {gpu_tp_count} TP parts)")
+    return True
+
+
+def test_bf16_write_buffer(gpu_tp_count):
+    """Test write_weight_scale_to_buffer with BF16 weights (no scales)"""
+    torch.manual_seed(123)
+
+    expert_num = 16
+    gpu_experts = expert_num
+    num_experts_per_tok = 8
+    hidden_size = 3072
+    intermediate_size = 1536
+
+    cpuinfer = make_cpu_infer()
+    cfg = build_config_bf16(cpuinfer, expert_num, num_experts_per_tok, hidden_size, intermediate_size)
+    weights = allocate_weights_bf16(expert_num, hidden_size, intermediate_size)
+
+    cfg.gate_proj = weights["gate_proj"].data_ptr()
+    cfg.up_proj = weights["up_proj"].data_ptr()
+    cfg.down_proj = weights["down_proj"].data_ptr()
+    cfg.gate_scale = 0
+    cfg.up_scale = 0
+    cfg.down_scale = 0
+
+    moe = kt_kernel_ext.moe.AMXBF16_MOE(cfg)
+
+    physical_to_logical_map = torch.arange(expert_num, dtype=torch.int64, device="cpu").contiguous()
+    cpuinfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+    cpuinfer.sync()
+
+    per_mat_weight_elems = weights["per_mat_weight_elems"]
+
+    # Calculate sizes per TP part (BF16 = 2 bytes per element)
+    weight_elems_per_expert_per_tp = per_mat_weight_elems // gpu_tp_count
+    weight_bytes_per_expert_per_tp = weight_elems_per_expert_per_tp * 2
+
+    total_weight_bytes_per_tp = gpu_experts * weight_bytes_per_expert_per_tp
+
+    # Create buffer lists (BF16: weights only, no scales)
+    w13_weight_bufs = [torch.empty(2 * total_weight_bytes_per_tp, dtype=torch.uint8) for _ in range(gpu_tp_count)]
+    w2_weight_bufs = [torch.empty(total_weight_bytes_per_tp, dtype=torch.uint8) for _ in range(gpu_tp_count)]
+    # Empty scale buffers (not used for BF16 but needed for interface)
+    w13_scale_bufs = [torch.empty(1, dtype=torch.float32) for _ in range(gpu_tp_count)]
+    w2_scale_bufs = [torch.empty(1, dtype=torch.float32) for _ in range(gpu_tp_count)]
+
+    print(f"[BF16] GPU TP count: {gpu_tp_count}, Experts: {expert_num}")
+    print(f"[BF16] Weight bytes per expert per TP: {weight_bytes_per_expert_per_tp}")
+
+    def get_expert_ptrs(expert_id):
+        w13_weight_ptrs = []
+        w13_scale_ptrs = []
+        w2_weight_ptrs = []
+        w2_scale_ptrs = []
+        for tp_idx in range(gpu_tp_count):
+            w13_weight_expert_offset = expert_id * 2 * weight_bytes_per_expert_per_tp
+            w2_weight_expert_offset = expert_id * weight_bytes_per_expert_per_tp
+
+            w13_weight_ptrs.append(w13_weight_bufs[tp_idx].data_ptr() + w13_weight_expert_offset)
+            w13_scale_ptrs.append(w13_scale_bufs[tp_idx].data_ptr())  # Not used
+            w2_weight_ptrs.append(w2_weight_bufs[tp_idx].data_ptr() + w2_weight_expert_offset)
+            w2_scale_ptrs.append(w2_scale_bufs[tp_idx].data_ptr())  # Not used
+        return w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs
+
+    # Warm up
+    for _ in range(2):
+        for expert_id in range(gpu_experts):
+            w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs = get_expert_ptrs(expert_id)
+            cpuinfer.submit(
+                moe.write_weight_scale_to_buffer_task(
+                    gpu_tp_count=gpu_tp_count,
+                    expert_id=expert_id,
+                    w13_weight_ptrs=w13_weight_ptrs,
+                    w13_scale_ptrs=w13_scale_ptrs,
+                    w2_weight_ptrs=w2_weight_ptrs,
+                    w2_scale_ptrs=w2_scale_ptrs,
+                )
+            )
+            cpuinfer.sync()
+
+    # Timing
+    begin_time = time.perf_counter_ns()
+    for expert_id in range(gpu_experts):
+        w13_weight_ptrs, w13_scale_ptrs, w2_weight_ptrs, w2_scale_ptrs = get_expert_ptrs(expert_id)
+        cpuinfer.submit(
+            moe.write_weight_scale_to_buffer_task(
+                gpu_tp_count=gpu_tp_count,
+                expert_id=expert_id,
+                w13_weight_ptrs=w13_weight_ptrs,
+                w13_scale_ptrs=w13_scale_ptrs,
+                w2_weight_ptrs=w2_weight_ptrs,
+                w2_scale_ptrs=w2_scale_ptrs,
+            )
+        )
+        cpuinfer.sync()
+    end_time = time.perf_counter_ns()
+    elapsed_ms = (end_time - begin_time) / 1e6
+
+    total_bytes = hidden_size * intermediate_size * gpu_experts * 3 * 2  # BF16 = 2 bytes
+    print(f"[BF16] write_weight_scale_to_buffer time: {elapsed_ms:.2f} ms")
+    print(f"[BF16] Throughput: {total_bytes / (elapsed_ms * 1e6):.2f} GB/s")
+
+    # Verify correctness (BF16: weights only, no scales)
+    def split_expert_tensor(tensor, chunk):
+        return [tensor[i * chunk : (i + 1) * chunk] for i in range(expert_num)]
+
+    gate_proj = weights["gate_proj"]
+    up_proj = weights["up_proj"]
+    down_proj = weights["down_proj"]
+
+    # View BF16 as uint8 for byte-level comparison
+    gate_bytes = gate_proj.view(torch.uint8)
+    up_bytes = up_proj.view(torch.uint8)
+    down_bytes = down_proj.view(torch.uint8)
+
+    per_mat_bytes = per_mat_weight_elems * 2  # BF16 = 2 bytes
+    gate_experts = split_expert_tensor(gate_bytes, per_mat_bytes)
+    up_experts = split_expert_tensor(up_bytes, per_mat_bytes)
+    down_experts = split_expert_tensor(down_bytes, per_mat_bytes)
+
+    for tp_idx in range(gpu_tp_count):
+        expected_w13_weights = []
+        expected_w2_weights = []
+
+        weight_bytes_per_tp = per_mat_bytes // gpu_tp_count
+
+        for expert_id in range(gpu_experts):
+            start_weight = tp_idx * weight_bytes_per_tp
+            end_weight = (tp_idx + 1) * weight_bytes_per_tp
+
+            gate_weight_tp = gate_experts[expert_id][start_weight:end_weight]
+            up_weight_tp = up_experts[expert_id][start_weight:end_weight]
+
+            # Down matrix: sliced column-wise (BF16 = 2 bytes per element)
+            down_weight_tp_parts = []
+            tp_slice_elems = intermediate_size // gpu_tp_count
+            tp_slice_bytes = tp_slice_elems * 2
+
+            for row_idx in range(hidden_size):
+                row_byte_start = row_idx * intermediate_size * 2
+                tp_byte_offset = row_byte_start + tp_idx * tp_slice_bytes
+                down_weight_tp_parts.append(down_experts[expert_id][tp_byte_offset : tp_byte_offset + tp_slice_bytes])
+
+            down_weight_tp = torch.cat(down_weight_tp_parts)
+
+            expected_w13_weights.append(gate_weight_tp)
+            expected_w13_weights.append(up_weight_tp)
+            expected_w2_weights.append(down_weight_tp)
+
+        expected_w13_weight = torch.cat(expected_w13_weights)
+        expected_w2_weight = torch.cat(expected_w2_weights)
+
+        if not torch.equal(w13_weight_bufs[tp_idx], expected_w13_weight):
+            diff_mask = w13_weight_bufs[tp_idx] != expected_w13_weight
+            first_diff_idx = diff_mask.nonzero()[0].item() if diff_mask.any() else -1
+            raise AssertionError(f"[BF16] w13 weight mismatch for TP {tp_idx} at index {first_diff_idx}")
+
+        if not torch.equal(w2_weight_bufs[tp_idx], expected_w2_weight):
+            diff_mask = w2_weight_bufs[tp_idx] != expected_w2_weight
+            first_diff_idx = diff_mask.nonzero()[0].item() if diff_mask.any() else -1
+            raise AssertionError(f"[BF16] w2 weight mismatch for TP {tp_idx} at index {first_diff_idx}")
+
+    print(f"[BF16] TP={gpu_tp_count} PASSED (verified {gpu_experts} experts across {gpu_tp_count} TP parts)")
+    return True
+
+
+def test_with_tp(quant_mode: str, gpu_tp_count: int):
+    """Test write_weight_scale_to_buffer with specified mode and TP count"""
+    if quant_mode == "fp8":
+        return test_fp8_write_buffer(gpu_tp_count)
+    elif quant_mode == "fp8_perchannel":
+        return test_fp8_perchannel_write_buffer(gpu_tp_count)
+    elif quant_mode == "bf16":
+        return test_bf16_write_buffer(gpu_tp_count)
+    else:
+        raise ValueError(f"Unsupported quant_mode: {quant_mode}")
+
+
+def main(quant_modes=None):
+    """Run tests for specified quant modes"""
+    if quant_modes is None:
+        quant_modes = ["fp8", "fp8_perchannel", "bf16"]
+
+    tp_values = [1, 2, 4]
+    all_passed = True
+    results = {}
+
+    for quant_mode in quant_modes:
+        print("\n" + "=" * 60)
+        print(f"Testing {quant_mode.upper()} write_weight_scale_to_buffer")
+        print("=" * 60)
+
+        for tp in tp_values:
+            print(f"\n--- Testing {quant_mode.upper()} with gpu_tp_count = {tp} ---")
+            try:
+                test_with_tp(quant_mode, tp)
+                results[(quant_mode, tp)] = "PASSED"
+            except Exception as e:
+                results[(quant_mode, tp)] = f"FAILED: {e}"
+                all_passed = False
+                print(f"[{quant_mode.upper()}] TP={tp} FAILED: {e}")
+
+    print("\n" + "=" * 60)
+    print("SUMMARY")
+    print("=" * 60)
+    for (mode, tp), result in results.items():
+        status = "PASS" if "PASSED" in result else "FAIL"
+        print(f"  [{status}] {mode.upper()} TP={tp}: {result}")
+
+    if all_passed:
+        print("\nALL TESTS PASSED")
+    else:
+        print("\nSOME TESTS FAILED")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1:
+        mode = sys.argv[1].lower()
+        if mode in ["fp8", "fp8_perchannel", "bf16"]:
+            main([mode])
+        else:
+            print(f"Unknown mode: {mode}. Use 'fp8', 'fp8_perchannel' or 'bf16'")
+            sys.exit(1)
+    else:
+        main()

kt-kernel/examples/torch_attention.py

@@ -0,0 +1,318 @@
+
+import math
+import os, sys
+import time
+import subprocess
+import platform
+import json
+from typing import Any, Dict, Optional, Tuple
+import torch
+import torch.nn.init as init
+from torch import nn
+
+class KDeepSeekV3Cache(nn.Module):
+    def __init__(
+        self,
+        # config: PretrainedConfig,
+        page_size: int = 256,
+        kv_lora_rank: int = 128,
+        k_caches: Optional[torch.Tensor] = None,
+        dtype=torch.bfloat16,
+        device=torch.device("cuda:0"),
+        
+    ):
+        super().__init__()
+        # self.config = config
+        self.dtype = dtype
+        self.device = device
+        self.kv_lora_rank = kv_lora_rank
+        self.page_size = page_size
+        self.v_caches = []
+        self.k_caches = k_caches
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+
+        page_idx: torch.Tensor,
+        page_offset: torch.Tensor,
+
+        cache_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.
+        It is VERY important to index using a tensor, otherwise you introduce a copy to the device.
+
+        Parameters:
+            key_states (`torch.Tensor`):
+                The new key states to cache.
+            value_states (`torch.Tensor`):
+                The new value states to cache.
+            layer_idx (`int`):
+                The index of the layer to cache the states for.
+            cache_kwargs (`Dict[str, Any]`, `optional`):
+                Additional arguments for the cache subclass. The `StaticCache` needs the `cache_position` input
+                to know how where to write in the cache.
+
+        Return:
+            A tuple containing the updated key and value states.
+        """
+        k_out = self.k_caches[layer_idx]
+
+        k_out[page_idx, page_offset, :, :self.kv_lora_rank] = key_states.reshape(-1, *key_states.shape[2:])
+        k_out[page_idx, page_offset, :, self.kv_lora_rank:] = value_states.reshape(-1, *value_states.shape[2:])
+        return k_out
+
+        
+    def get_page_table(self, cache_position: torch.Tensor, q_indptr: torch.Tensor, kv_indptr: torch.Tensor, kv_indices: torch.Tensor, bsz_tensors: torch.tensor):
+        page_offset = cache_position % self.page_size  
+        page_idx_local = cache_position // self.page_size  
+        query_ids = torch.zeros_like(cache_position)
+        for i in range(len(q_indptr) - 1):
+            start_idx = q_indptr[i]
+            end_idx = q_indptr[i + 1]
+            query_ids[start_idx:end_idx] = i
+        page_idx = torch.zeros_like(page_idx_local)
+        for i in range(bsz_tensors[0]):
+            query_id = query_ids[i]
+            local_block = page_idx_local[i]
+            start_block = kv_indptr[query_id]
+            if local_block < kv_indptr[query_id + 1] - kv_indptr[query_id]:
+                page_idx[i] = kv_indices[start_block + local_block]
+        
+        return page_idx, page_offset
+
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    b, h, s, d = q.shape
+    q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
+    b, h, s, d = k.shape
+    k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class DeepseekV2RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        DeepseekV2RMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight * hidden_states).to(input_dtype)
+
+
+class DeepseekV2RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+        super().__init__()
+        self.scaling_factor = scaling_factor
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        # For BC we register cos and sin cached
+        self.max_seq_len_cached = max_position_embeddings
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        position_ids_expanded = position_ids[:, None, :].float()
+        # Force float32 since bfloat16 loses precision on long contexts
+        # See https://github.com/huggingface/transformers/pull/29285
+        device_type = x.device.type
+        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)   
+
+class DeepseekV3RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        inv_freq = 1.0 / (
+            self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
+        )
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # Build here to make `torch.jit.trace` work.
+        self._set_cos_sin_cache(
+            seq_len=max_position_embeddings,
+            device=self.inv_freq.device,
+            dtype=torch.get_default_dtype(),
+        )
+        # self.max_seq_len_cached = None
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(
+            self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
+        )
+
+        freqs = torch.outer(t, self.inv_freq.to(t.device))
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        print("emb", emb.shape)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if self.max_seq_len_cached is None: # or seq_len[-1] > self.max_seq_len_cached:
+            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
+
+        return (
+            self.cos_cached[seq_len].to(dtype=x.dtype),
+            self.sin_cached[seq_len].to(dtype=x.dtype),
+        )
+
+# Inverse dim formula to find dim based on number of rotations
+def yarn_find_correction_dim(
+    num_rotations, dim, base=10000, max_position_embeddings=2048
+):
+    return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
+        2 * math.log(base)
+    )
+
+
+# Find dim range bounds based on rotations
+def yarn_find_correction_range(
+    low_rot, high_rot, dim, base=10000, max_position_embeddings=2048
+):
+    low = math.floor(
+        yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
+    )
+    high = math.ceil(
+        yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
+    )
+    return max(low, 0), min(high, dim - 1)  # Clamp values just in case
+
+def yarn_linear_ramp_mask(min, max, dim):
+    if min == max:
+        max += 0.001  # Prevent singularity
+
+    linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
+    ramp_func = torch.clamp(linear_func, 0, 1)
+    return ramp_func
+
+def yarn_get_mscale(scale=1, mscale=1):
+    if scale <= 1:
+        return 1.0
+    return 0.1 * mscale * math.log(scale) + 1.0
+
+class DeepseekV3YarnRotaryEmbedding(DeepseekV3RotaryEmbedding):
+
+    def __init__(
+        self,
+        dim,
+        max_position_embeddings=2048,
+        base=10000,
+        device=None,
+        scaling_factor=1.0,
+        original_max_position_embeddings=4096,
+        beta_fast=32,
+        beta_slow=1,
+        mscale=1,
+        mscale_all_dim=0,
+    ):
+        self.scaling_factor = scaling_factor
+        self.original_max_position_embeddings = original_max_position_embeddings
+        self.beta_fast = beta_fast
+        self.beta_slow = beta_slow
+        self.mscale = mscale
+        self.mscale_all_dim = mscale_all_dim
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        dim = self.dim
+
+        freq_extra = 1.0 / (
+            self.base
+            ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)
+        )
+        freq_inter = 1.0 / (
+            self.scaling_factor
+            * self.base
+            ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)
+        )
+
+        low, high = yarn_find_correction_range(
+            self.beta_fast,
+            self.beta_slow,
+            dim,
+            self.base,
+            self.original_max_position_embeddings,
+        )
+        inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2).to(
+            device=device, dtype=torch.float32
+        )
+        inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        # 判断 seq_len是否是 tensor
+        if isinstance(seq_len,torch.Tensor):
+            t = seq_len
+        else:
+            t = torch.arange(seq_len, device=device, dtype=torch.float32)
+
+        freqs = torch.outer(t, inv_freq)
+
+        _mscale = float(
+            yarn_get_mscale(self.scaling_factor, self.mscale)
+            / yarn_get_mscale(self.scaling_factor, self.mscale_all_dim)
+        )
+
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer(
+            "cos_cached", (emb.cos() * _mscale).to(dtype), persistent=False
+        )
+        self.register_buffer(
+            "sin_cached", (emb.sin() * _mscale).to(dtype), persistent=False
+        )
+