[feat] : Add MACA backend support for kt-kernel and fix CPU MoE tests (#2044)

* Add MACA backend support for kt-kernel

* Add MACA event API mappings

* Fix AMX build flags and CPU MoE tests


---------

Co-authored-by: <Engle_Chaveztih@sociologist.com>

kt-kernel/CMakeLists.txt

@@ -16,6 +16,7 @@ option(LLAMA_AVX512_FANCY_SIMD "llama: enable AVX512-VL, AVX512-BW, AVX512-DQ, A
 option(KTRANSFORMERS_USE_CUDA "ktransformers: use CUDA" OFF)
 option(KTRANSFORMERS_USE_MUSA "ktransformers: use MUSA" OFF)
 option(KTRANSFORMERS_USE_ROCM "ktransformers: use ROCM" OFF)
+option(KTRANSFORMERS_USE_MACA "ktransformers: use MACA" OFF)
 option(KTRANSFORMERS_CUDA_STATIC_RUNTIME "ktransformers: statically link CUDA runtime" ON)
 option(KTRANSFORMERS_CPU_USE_KML "ktransformers: CPU use KML" OFF)
 option(KTRANSFORMERS_CPU_USE_AMX_AVX512 "ktransformers: CPU use AMX or AVX512" OFF)
@@ -329,6 +330,16 @@ endif()
 
 list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake")
 
+if(NOT DEFINED MACA_PATH OR MACA_PATH STREQUAL "")
+    if(DEFINED ENV{MACA_PATH} AND EXISTS "$ENV{MACA_PATH}")
+        set(MACA_PATH $ENV{MACA_PATH})
+    elseif(EXISTS /opt/maca)
+        set(MACA_PATH /opt/maca)
+    else()
+        set(MACA_PATH /usr/local/maca)
+    endif()
+endif()
+
 if(KTRANSFORMERS_CPU_MOE_AMD)
     set(BLIS_ROOT "" CACHE PATH "Root directory of BLIS installation")
     set(_BLIS_SEARCH_DIRS)
@@ -371,6 +382,11 @@ if(HOST_IS_X86)
             list(APPEND ARCH_FLAGS -mamx-tile -mamx-bf16 -mamx-int8)
             message(STATUS "AMX enabled")
         endif()
+        if(NOT MSVC)
+            # AMX/SFT kernels use 256-bit AVX512 masked load/store intrinsics.
+            # GCC requires AVX512VL in addition to AVX512BW for those intrinsics.
+            list(APPEND ARCH_FLAGS -mavx512vl)
+        endif()
         # add_executable(amx-test ${CMAKE_CURRENT_SOURCE_DIR}/operators/amx/amx-test.cpp)
         # target_link_libraries(amx-test llama)
         if(KTRANSFORMERS_CPU_DEBUG)
@@ -458,6 +474,33 @@ elseif(KTRANSFORMERS_USE_MUSA)
         message(STATUS "MUSA Toolkit found")
         add_compile_definitions(KTRANSFORMERS_USE_MUSA=1)
     endif()
+elseif(KTRANSFORMERS_USE_MACA)
+    find_path(MACA_INCLUDE_DIR
+        NAMES mcr/maca.h
+        HINTS "${MACA_PATH}"
+        PATH_SUFFIXES include
+    )
+    find_library(MACA_RUNTIME_LIBRARY
+        NAMES mcruntime
+        HINTS "${MACA_PATH}"
+        PATH_SUFFIXES lib lib64
+    )
+    find_library(MACA_BLAS_LIBRARY
+        NAMES mcblas
+        HINTS "${MACA_PATH}"
+        PATH_SUFFIXES lib lib64
+    )
+
+    if(NOT MACA_INCLUDE_DIR OR NOT MACA_RUNTIME_LIBRARY OR NOT MACA_BLAS_LIBRARY)
+        message(FATAL_ERROR "KTRANSFORMERS_USE_MACA=ON but MACA SDK was not found. Set MACA_PATH to the MACA SDK root.")
+    endif()
+
+    message(STATUS "MACA SDK found at ${MACA_PATH}")
+    message(STATUS "MACA include: ${MACA_INCLUDE_DIR}")
+    message(STATUS "MACA runtime: ${MACA_RUNTIME_LIBRARY}")
+    message(STATUS "MACA BLAS: ${MACA_BLAS_LIBRARY}")
+    include_directories(${MACA_INCLUDE_DIR})
+    add_compile_definitions(KTRANSFORMERS_USE_MACA=1)
 elseif(KTRANSFORMERS_CPU_USE_KML)
     message(STATUS "KML CPU detected")
 else()
@@ -661,6 +704,7 @@ endif()
 
 
 if(KTRANSFORMERS_USE_CUDA)
+    add_compile_definitions(USE_CUDA=1)
     # Link CUDA runtime (static or dynamic)
     if(KTRANSFORMERS_CUDA_STATIC_RUNTIME)
         # Platform-aware static library path
@@ -694,7 +738,14 @@ if(KTRANSFORMERS_USE_ROCM)
     message(STATUS "Building for HIP")
 endif()
 if(KTRANSFORMERS_USE_MUSA)
+    add_compile_definitions(USE_MUSA=1)
     target_link_libraries(${PROJECT_NAME} PRIVATE MUSA::musart)
+    message(STATUS "Building for MUSA")
+endif()
+if(KTRANSFORMERS_USE_MACA)
+    add_compile_definitions(USE_MACA=1)
+    target_link_libraries(${PROJECT_NAME} PRIVATE ${MACA_RUNTIME_LIBRARY} ${MACA_BLAS_LIBRARY})
+    message(STATUS "Building for MACA")
 endif()
 
 

kt-kernel/cpu_backend/cpuinfer.h

@@ -24,6 +24,8 @@
 #elif KTRANSFORMERS_USE_ROCM
 #define __HIP_PLATFORM_AMD__
 #include "vendors/hip.h"
+#elif KTRANSFORMERS_USE_MACA
+#include "vendors/maca.h"
 #endif
 
 #include "./vendors/vendor.h"
@@ -83,7 +85,8 @@ class CPUInfer {
   }
 #ifndef KTRANSFORMERS_CPU_ONLY
   void submit_with_cuda_stream(intptr_t user_cuda_stream, std::pair<intptr_t, intptr_t> params) {
-#if defined(KTRANSFORMERS_USE_CUDA)
+#if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) || \
+    defined(KTRANSFORMERS_USE_MACA)
     void (*func)(void*) = (void (*)(void*))params.first;
     void* args = (void*)params.second;
     *((CPUInfer**)args) = this;
@@ -100,6 +103,7 @@ class CPUInfer {
   static void sync_(void* sync_args) {
     SyncArgs* args = (SyncArgs*)sync_args;
     args->cpuinfer->task_queue_->sync(args->allow_n_pending);
+    delete args;
   }
 
   void sync(size_t allow_n_pending = 0) {
@@ -108,7 +112,8 @@ class CPUInfer {
   }
 #ifndef KTRANSFORMERS_CPU_ONLY
   void sync_with_cuda_stream(intptr_t user_cuda_stream, size_t allow_n_pending = 0) {
-#if defined(KTRANSFORMERS_USE_CUDA)
+#if defined(KTRANSFORMERS_USE_CUDA) || defined(KTRANSFORMERS_USE_MUSA) || defined(KTRANSFORMERS_USE_ROCM) || \
+    defined(KTRANSFORMERS_USE_MACA)
     SyncArgs* args = new SyncArgs{this, allow_n_pending};
     cudaLaunchHostFunc((cudaStream_t)user_cuda_stream, (cudaHostFn_t)&sync_, (void*)args);
 #endif
@@ -119,4 +124,4 @@ class CPUInfer {
   TaskQueue* task_queue_;
 };
 
-#endif
+#endif

kt-kernel/cpu_backend/vendors/maca.h

@@ -0,0 +1,145 @@
+#pragma once
+
+#include <common/maca_bfloat16.h>
+#include <common/maca_fp16.h>
+#include <common/mc_library_types.h>
+#include <mcblas/mcblas.h>
+#include <mcr/maca.h>
+
+#define CUBLAS_COMPUTE_16F MCBLAS_COMPUTE_16F
+#define CUBLAS_COMPUTE_32F MCBLAS_COMPUTE_32F
+#define CUBLAS_COMPUTE_32F_FAST_16F MCBLAS_COMPUTE_32F_FAST_16F
+#define CUBLAS_COMPUTE_32F_FAST_16BF MCBLAS_COMPUTE_32F_FAST_16BF
+#define CUBLAS_COMPUTE_32F_FAST_TF32 MCBLAS_COMPUTE_32F_FAST_TF32
+#define CUBLAS_GEMM_DEFAULT MCBLAS_GEMM_DEFAULT
+#define CUBLAS_GEMM_DEFAULT_TENSOR_OP MCBLAS_GEMM_DEFAULT_TENSOR_OP
+#define CUBLAS_OP_N MCBLAS_OP_N
+#define CUBLAS_OP_T MCBLAS_OP_T
+#define CUBLAS_STATUS_SUCCESS MCBLAS_STATUS_SUCCESS
+#define CUBLAS_TF32_TENSOR_OP_MATH MCBLAS_TF32_TENSOR_OP_MATH
+#define CUDA_R_16F MACA_R_16F
+#define CUDA_R_16BF MACA_R_16BF
+#define CUDA_R_32F MACA_R_32F
+#define cublasComputeType_t mcblasComputeType_t
+#define cublasCreate mcblasCreate
+#define cublasDestroy mcblasDestroy
+#define cublasGemmEx mcblasGemmEx
+#define cublasGemmBatchedEx mcblasGemmBatchedEx
+#define cublasGemmStridedBatchedEx mcblasGemmStridedBatchedEx
+#define cublasHandle_t mcblasHandle_t
+#define cublasSetMathMode mcblasSetMathMode
+#define cublasSetStream mcblasSetStream
+#define cublasSgemm mcblasSgemm
+#define cublasStatus_t mcblasStatus_t
+#define cublasOperation_t mcblasOperation_t
+#define cublasGetStatusString mcblasGetStatusString
+#define cudaDataType_t macaDataType_t
+#define cudaDeviceCanAccessPeer mcDeviceCanAccessPeer
+#define cudaDeviceDisablePeerAccess mcDeviceDisablePeerAccess
+#define cudaDeviceEnablePeerAccess mcDeviceEnablePeerAccess
+#define cudaDeviceProp mcDeviceProp_t
+#define cudaDeviceSynchronize mcDeviceSynchronize
+#define cudaError_t mcError_t
+#define cudaErrorPeerAccessAlreadyEnabled mcErrorPeerAccessAlreadyEnabled
+#define cudaErrorPeerAccessNotEnabled mcErrorPeerAccessNotEnabled
+#define cudaEventCreate mcEventCreate
+#define cudaEventCreateWithFlags mcEventCreateWithFlags
+#define cudaEventDisableTiming mcEventDisableTiming
+#define cudaEventElapsedTime mcEventElapsedTime
+#define cudaEventRecord mcEventRecord
+#define cudaEventSynchronize mcEventSynchronize
+#define cudaEvent_t mcEvent_t
+#define cudaEventDestroy mcEventDestroy
+#define cudaFree mcFree
+#define cudaFreeHost mcFreeHost
+#define cudaGetDevice mcGetDevice
+#define cudaGetDeviceCount mcGetDeviceCount
+#define cudaGetDeviceProperties mcGetDeviceProperties
+#define cudaGetErrorString mcGetErrorString
+#define cudaGetLastError mcGetLastError
+#define cudaHostRegister mcHostRegister
+#define cudaHostRegisterPortable mcHostRegisterPortable
+#define cudaHostRegisterReadOnly mcHostRegisterReadOnly
+#define cudaHostUnregister mcHostUnregister
+#define cudaLaunchHostFunc mcLaunchHostFunc
+#define cudaMalloc mcMalloc
+#define cudaMallocHost(ptr, size) mcMallocHost(ptr, size)
+#define cudaMallocManaged mcMallocManaged
+#define cudaMemcpy mcMemcpy
+#define cudaMemcpyAsync mcMemcpyAsync
+#define cudaMemcpyPeerAsync mcMemcpyPeerAsync
+#define cudaMemcpy2DAsync mcMemcpy2DAsync
+#define cudaMemcpyDeviceToDevice mcMemcpyDeviceToDevice
+#define cudaMemcpyDeviceToHost mcMemcpyDeviceToHost
+#define cudaMemcpyHostToDevice mcMemcpyHostToDevice
+#define cudaMemcpyKind mcMemcpyKind
+#define cudaMemset mcMemset
+#define cudaMemsetAsync mcMemsetAsync
+#define cudaMemGetInfo mcMemGetInfo
+#define cudaOccupancyMaxPotentialBlockSize mcOccupancyMaxPotentialBlockSize
+#define cudaSetDevice mcSetDevice
+#define cudaStreamCreateWithFlags mcStreamCreateWithFlags
+#define cudaStreamDestroy mcStreamDestroy
+#define cudaStreamNonBlocking mcStreamNonBlocking
+#define cudaStreamPerThread mcStreamPerThread
+#define cudaStreamSynchronize mcStreamSynchronize
+#define cudaStreamWaitEvent mcStreamWaitEvent
+#define cudaStream_t mcStream_t
+#define cudaHostFn_t mcHostFn_t
+#define cudaSuccess mcSuccess
+#define nv_bfloat16 maca_bfloat16
+
+// Additional mappings for MACA virtual memory pool.
+#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED mcDeviceAttributeVirtualMemoryManagementSupported
+#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE mcMemAccessFlagsProtReadWrite
+#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED MC_MEM_ALLOC_GRANULARITY_RECOMMENDED
+#define CU_MEM_ALLOCATION_TYPE_PINNED mcMemAllocationTypePinned
+#define CU_MEM_LOCATION_TYPE_DEVICE mcMemLocationTypeDevice
+#define CUdevice mcDevice_t
+#define CUdeviceptr mcDeviceptr_t
+#define CUmemAccessDesc mcMemAccessDesc
+#define CUmemAllocationProp mcMemAllocationProp
+#define CUmemGenericAllocationHandle mcMemGenericAllocationHandle
+#define cuDeviceGet mcDeviceGet
+#define cuDeviceGetAttribute mcDeviceGetAttribute
+#define cuMemAddressFree mcMemAddressFree
+#define cuMemAddressReserve mcMemAddressReserve
+#define cuMemCreate mcMemCreate
+#define cuMemGetAllocationGranularity mcMemGetAllocationGranularity
+#define cuMemMap mcMemMap
+#define cuMemRelease mcMemRelease
+#define cuMemSetAccess mcMemSetAccess
+#define cuMemUnmap mcMemUnmap
+#define cudaFuncAttributeMaxDynamicSharedMemorySize mcFuncAttributeMaxDynamicSharedMemorySize
+#define cudaFuncSetAttribute mcFuncSetAttribute
+#define cudaMemcpy3DPeerParms mcMemcpy3DPeerParms
+#define make_cudaExtent make_mcExtent
+#define make_cudaPitchedPtr make_mcPitchedPtr
+
+// Additional mappings for MACA graphs.
+#define CUDA_SUCCESS mcSuccess
+#define CUresult mcError_t
+#define cuGetErrorString(error, pstr) (*(pstr) = mcGetErrorString(error), mcSuccess)
+#define cudaErrorGraphExecUpdateFailure mcErrorGraphExecUpdateFailure
+#define cudaErrorInvalidDeviceFunction mcErrorInvalidDeviceFunction
+#define cudaGraphDestroy mcGraphDestroy
+#define cudaGraphExecDestroy mcGraphExecDestroy
+#define cudaGraphExec_t mcGraphExec_t
+#define cudaGraphExecUpdate mcGraphExecUpdate
+#define cudaGraphExecUpdateResultInfo mcGraphExecUpdateResultInfo
+#define cudaGraphGetNodes mcGraphGetNodes
+#define cudaGraphInstantiate mcGraphInstantiate
+#define cudaGraphKernelNodeGetParams mcGraphKernelNodeGetParams
+#define cudaGraphKernelNodeSetParams mcGraphKernelNodeSetParams
+#define cudaGraphLaunch mcGraphLaunch
+#define cudaGraphNodeGetType mcGraphNodeGetType
+#define cudaGraphNode_t mcGraphNode_t
+#define cudaGraphNodeType mcGraphNodeType
+#define cudaGraphNodeTypeKernel mcGraphNodeTypeKernel
+#define cudaGraph_t mcGraph_t
+#define cudaKernelNodeParams mcKernelNodeParams
+#define cudaStreamCaptureModeRelaxed mcStreamCaptureModeRelaxed
+#define cudaStreamBeginCapture mcStreamBeginCapture
+#define cudaStreamEndCapture mcStreamEndCapture
+
+typedef maca_bfloat16 nv_bfloat16;

kt-kernel/cpu_backend/vendors/vendor.h

@@ -8,6 +8,8 @@
 #include "hip.h"
 #elif USE_MUSA
 #include "musa.h"
+#elif USE_MACA
+#include "maca.h"
 #endif
 
-#endif  // CPUINFER_VENDOR_VENDOR_H
+#endif  // CPUINFER_VENDOR_VENDOR_H

kt-kernel/setup.py

@@ -34,10 +34,12 @@ Environment knobs (export before running pip install .):
   CPUINFER_NATIVE=ON               (override LLAMA_NATIVE)
 
 
-GPU backends (if ever added later, keep placeholders):
+GPU backends:
   CPUINFER_USE_CUDA=0/1           -DKTRANSFORMERS_USE_CUDA
   CPUINFER_USE_ROCM=0/1           -DKTRANSFORMERS_USE_ROCM
   CPUINFER_USE_MUSA=0/1           -DKTRANSFORMERS_USE_MUSA
+  CPUINFER_USE_MACA=0/1           -DKTRANSFORMERS_USE_MACA
+  MACA_PATH=/opt/maca             MACA SDK root
 
 Usage:
   pip install .
@@ -102,6 +104,12 @@ def _forward_str_env(cmake_args: list[str], env_name: str, cmake_flag: str) -> b
 
 REPO_ROOT = Path(__file__).parent.resolve()
 
+# setuptools resolves package_dir and inplace extension copy paths relative to
+# the current working directory. Keep direct invocations like
+# `python kt-kernel/setup.py build_ext --inplace` equivalent to running from
+# inside kt-kernel.
+os.chdir(REPO_ROOT)
+
 CPU_FEATURE_MAP = {
     "FANCY": "-DLLAMA_NATIVE=OFF -DLLAMA_FMA=ON -DLLAMA_F16C=ON -DLLAMA_AVX=ON -DLLAMA_AVX2=ON -DLLAMA_AVX512=ON -DLLAMA_AVX512_FANCY_SIMD=ON",
     "AVX512": "-DLLAMA_NATIVE=OFF -DLLAMA_FMA=ON -DLLAMA_F16C=ON -DLLAMA_AVX=ON -DLLAMA_AVX2=ON -DLLAMA_AVX512=ON",
@@ -490,6 +498,19 @@ class CMakeBuild(build_ext):
                     return cand
             return None
 
+        def find_maca_path() -> str | None:
+            for cand in [
+                os.environ.get("MACA_PATH"),
+                "/opt/maca",
+                "/usr/local/maca",
+            ]:
+                if not cand:
+                    continue
+                root = Path(cand)
+                if (root / "include" / "mcr" / "maca.h").exists():
+                    return str(root)
+            return None
+
         # Note: We no longer set CMAKE_CUDA_ARCHITECTURES by default.
         # If users want to specify CUDA archs, they can set env CPUINFER_CUDA_ARCHS
         # (e.g. "89" or "86;89") or pass it via CMAKE_ARGS.
@@ -497,9 +518,30 @@ class CMakeBuild(build_ext):
         # Normalize CPUINFER_USE_CUDA: if unset, auto-detect; otherwise respect truthy/falsey values
         cuda_env = _env_get_bool("CPUINFER_USE_CUDA", None)
         if cuda_env is None:
-            auto_cuda = detect_cuda_toolkit()
-            os.environ["CPUINFER_USE_CUDA"] = "1" if auto_cuda else "0"
-            print(f"-- CPUINFER_USE_CUDA not set; auto-detected CUDA toolkit: {'YES' if auto_cuda else 'NO'}")
+            requested_non_cuda_gpu = any(
+                _env_get_bool(name, False)
+                for name in ("CPUINFER_USE_ROCM", "CPUINFER_USE_MUSA", "CPUINFER_USE_MACA")
+            )
+            if requested_non_cuda_gpu:
+                os.environ["CPUINFER_USE_CUDA"] = "0"
+                print("-- CPUINFER_USE_CUDA not set; another GPU backend was requested, disabling CUDA auto-detect")
+            else:
+                auto_cuda = detect_cuda_toolkit()
+                os.environ["CPUINFER_USE_CUDA"] = "1" if auto_cuda else "0"
+                print(f"-- CPUINFER_USE_CUDA not set; auto-detected CUDA toolkit: {'YES' if auto_cuda else 'NO'}")
+
+        enabled_gpu_backends = [
+            name
+            for name, env_name in (
+                ("CUDA", "CPUINFER_USE_CUDA"),
+                ("ROCM", "CPUINFER_USE_ROCM"),
+                ("MUSA", "CPUINFER_USE_MUSA"),
+                ("MACA", "CPUINFER_USE_MACA"),
+            )
+            if _env_get_bool(env_name, False)
+        ]
+        if len(enabled_gpu_backends) > 1:
+            raise RuntimeError(f"GPU backends are mutually exclusive, but enabled: {', '.join(enabled_gpu_backends)}")
 
         # Base CMake args
         cmake_args = [
@@ -647,6 +689,12 @@ class CMakeBuild(build_ext):
             cmake_args.append("-DKTRANSFORMERS_USE_ROCM=ON")
         if _env_get_bool("CPUINFER_USE_MUSA", False):
             cmake_args.append("-DKTRANSFORMERS_USE_MUSA=ON")
+        if _env_get_bool("CPUINFER_USE_MACA", False):
+            cmake_args.append("-DKTRANSFORMERS_USE_MACA=ON")
+            maca_path = find_maca_path()
+            if maca_path and not os.environ.get("MACA_PATH"):
+                cmake_args.append(f"-DMACA_PATH={maca_path}")
+            print("-- Enabling MACA backend (-DKTRANSFORMERS_USE_MACA=ON)")
 
         # Respect user extra CMAKE_ARGS (space separated)
         extra = os.environ.get("CMAKE_ARGS")
@@ -711,11 +759,21 @@ else:
     print(f"-- Version: {VERSION}")
 
 # Package name is always kt-kernel
-# The CUDA-enabled wheel includes both CPU multi-variant support and CUDA capabilities
+# GPU-enabled wheels include both CPU multi-variant support and the selected GPU runtime.
 PACKAGE_NAME = "kt-kernel"
-cuda_enabled = _env_get_bool("CPUINFER_USE_CUDA", False)
-if cuda_enabled:
-    print(f"-- Building kt-kernel with CUDA support (+ CPU multi-variant)")
+gpu_backend = None
+for _backend_name, _env_name in (
+    ("CUDA", "CPUINFER_USE_CUDA"),
+    ("ROCM", "CPUINFER_USE_ROCM"),
+    ("MUSA", "CPUINFER_USE_MUSA"),
+    ("MACA", "CPUINFER_USE_MACA"),
+):
+    if _env_get_bool(_env_name, False):
+        gpu_backend = _backend_name
+        break
+
+if gpu_backend:
+    print(f"-- Building kt-kernel with {gpu_backend} support (+ CPU multi-variant)")
 else:
     print(f"-- Building kt-kernel (CPU-only multi-variant)")
 

kt-kernel/test/per_commit/test_moe_avx2_accuracy_bf16.py

@@ -11,8 +11,12 @@ import time
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
 
+import pytest
 import torch
 from kt_kernel import kt_kernel_ext
+from ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=120, suite="default")
 
 # Small test parameters for fast validation
 expert_num = 8
@@ -69,6 +73,8 @@ def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj):
     return t_output
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("qlen,label", [(1, "Decode"), (16, "Prefill")])
 def test_avx2_bf16_accuracy(qlen, label):
     """Test AVX2 BF16 MoE accuracy."""
     physical_to_logical_map = torch.tensor(range(expert_num), device="cpu", dtype=torch.int64).contiguous()

kt-kernel/test/per_commit/test_moe_avx2_accuracy_fp8.py

@@ -8,8 +8,12 @@ import math
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
 
+import pytest
 import torch
 from kt_kernel import kt_kernel_ext
+from ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=120, suite="default")
 
 expert_num = 8
 hidden_size = 256
@@ -94,8 +98,10 @@ def fp8_e4m3_to_float(byte_val):
         return 0.0
     if exp == 0:
         val = (2**-6) * (man / 8.0)
-    elif exp == 15:
-        return float("nan")
+    elif exp == 15 and man == 7:
+        # Match the AVX2 LUT: E4M3 has finite exp=15 values up to 0x7e,
+        # and the NaN sentinel is treated as zero to avoid propagation.
+        return 0.0
     else:
         val = (2**(exp-7)) * (1.0 + man / 8.0)
     return -val if sign else val
@@ -151,6 +157,8 @@ def moe_torch(input, expert_ids, weights, gate_proj, up_proj, down_proj):
     return (new_x.view(*expert_ids.shape, -1).float().mul_(weights.unsqueeze(-1)).sum(1)).to(new_x.dtype)
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("qlen,label", [(1, "Decode"), (16, "Prefill")])
 def test_avx2_fp8_accuracy(qlen, label):
     physical_to_logical_map = torch.tensor(range(expert_num), dtype=torch.int64).contiguous()
     CPUInfer = kt_kernel_ext.CPUInfer(CPUINFER_PARAM)

kt-kernel/test/per_commit/test_moe_gptq_int4_accuracy.py

@@ -8,8 +8,13 @@ import sys
 import types
 from pathlib import Path
 
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python"))
 
+register_cpu_ci(est_time=120, suite="default")
+
 import pytest
 import torch
 import kt_kernel_ext