chore: import upstream snapshot with attribution

kt-kernel/test/per_commit/__init__.py

@@ -0,0 +1,4 @@
+"""Per-commit tests for KT-Kernel.
+
+Tests in this directory are run on every commit in CI.
+"""

kt-kernel/test/per_commit/test_amd_placeholder.py

@@ -0,0 +1,36 @@
+"""AMD/ROCm backend tests for KT-Kernel (Placeholder).
+
+This file is a placeholder for future AMD/ROCm backend tests.
+Currently, KT-Kernel focuses on CPU optimizations (Intel AMX/AVX512).
+
+To implement AMD tests:
+1. Add actual test functions with @pytest.mark.amd
+2. Update the estimated time in register_amd_ci()
+3. Implement AMD/ROCm-specific initialization and validation tests
+"""
+
+import os
+import sys
+
+# Add parent directory to path for CI registration
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_amd_ci
+
+# Register this test for AMD CI (estimated time: 10 seconds, placeholder)
+# Update suite name when implementing: currently using "stage-a-test-1"
+register_amd_ci(est_time=10, suite="stage-a-test-1")
+
+
+def test_amd_placeholder():
+    """Placeholder test for AMD/ROCm backend.
+
+    TODO: Implement actual AMD/ROCm tests when AMD support is added to kt-kernel.
+    """
+    # Currently a no-op placeholder
+    pass
+
+
+if __name__ == "__main__":
+    # Allow running standalone (required by test runner)
+    print("⚠ AMD/ROCm tests are not yet implemented (placeholder)")
+    print("✓ Placeholder test passed")

kt-kernel/test/per_commit/test_basic_cpu.py

@@ -0,0 +1,82 @@
+"""Basic CPU backend tests for KT-Kernel.
+
+These tests verify basic functionality without requiring model files.
+"""
+
+import os
+import sys
+import pytest
+
+# Add parent directory to path for CI registration
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
+# Register this test for CPU CI with estimated runtime of 30 seconds
+register_cpu_ci(est_time=30, suite="default")
+
+# Check if kt_kernel_ext is available
+try:
+    import kt_kernel  # Import kt_kernel first to register kt_kernel_ext
+
+    kt_kernel_ext = kt_kernel.kt_kernel_ext  # Access the extension module
+    HAS_KT_KERNEL = True
+except ImportError:
+    HAS_KT_KERNEL = False
+    kt_kernel_ext = None
+
+
+@pytest.mark.cpu
+def test_kt_kernel_import():
+    """Test that kt_kernel_ext can be imported."""
+    if not HAS_KT_KERNEL:
+        pytest.skip("kt_kernel_ext not built or available")
+
+    assert kt_kernel_ext is not None, "kt_kernel_ext module should be importable"
+
+
+@pytest.mark.cpu
+def test_cpu_infer_initialization():
+    """Test that CPUInfer can be initialized."""
+    if not HAS_KT_KERNEL:
+        pytest.skip("kt_kernel_ext not built or available")
+
+    # Initialize CPUInfer with 4 threads
+    cpuinfer = kt_kernel_ext.CPUInfer(4)
+    assert cpuinfer is not None, "CPUInfer should be initialized successfully"
+
+
+@pytest.mark.cpu
+def test_basic_module_attributes():
+    """Test that kt_kernel_ext has expected attributes."""
+    if not HAS_KT_KERNEL:
+        pytest.skip("kt_kernel_ext not built or available")
+
+    # Check for key attributes/functions
+    assert hasattr(kt_kernel_ext, "CPUInfer"), "kt_kernel_ext should have CPUInfer class"
+
+
+def run_all_tests():
+    """Run all tests in this file (for standalone execution)."""
+    if not HAS_KT_KERNEL:
+        print("⚠ kt_kernel_ext not available, skipping tests")
+        return
+
+    try:
+        test_kt_kernel_import()
+        print("✓ test_kt_kernel_import passed")
+
+        test_cpu_infer_initialization()
+        print("✓ test_cpu_infer_initialization passed")
+
+        test_basic_module_attributes()
+        print("✓ test_basic_module_attributes passed")
+
+        print("\n✓ All tests passed!")
+    except Exception as e:
+        print(f"\n✗ Test failed: {e}")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    # Allow running standalone (required by test runner)
+    run_all_tests()

kt-kernel/test/per_commit/test_convert_kt_to_sglang_adapter.py

@@ -0,0 +1,373 @@
+import importlib.util
+import os
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors.torch import load_file, save_file
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
+
+register_cpu_ci(est_time=5, suite="default")
+
+SCRIPT_PATH = (
+    Path(__file__).resolve().parents[2]
+    / "scripts"
+    / "convert_kt_to_sglang_adapter.py"
+)
+SPEC = importlib.util.spec_from_file_location("convert_kt_to_sglang_adapter", SCRIPT_PATH)
+converter = importlib.util.module_from_spec(SPEC)
+assert SPEC.loader is not None
+SPEC.loader.exec_module(converter)
+
+
+def _write_full_fused_checkpoint(path: Path, *, rank: int = 3) -> dict[str, torch.Tensor]:
+    e, h, i = 2, 5, 7
+    tensors = {
+        "layers.2.experts.gate_lora_a": torch.arange(e * rank * h, dtype=torch.float32).reshape(e, rank, h),
+        "layers.2.experts.gate_lora_b": torch.arange(e * i * rank, dtype=torch.float32).reshape(e, i, rank),
+        "layers.2.experts.up_lora_a": torch.arange(e * rank * h, dtype=torch.float32).reshape(e, rank, h) + 100,
+        "layers.2.experts.up_lora_b": torch.arange(e * i * rank, dtype=torch.float32).reshape(e, i, rank) + 200,
+        "layers.2.experts.down_lora_a": torch.arange(e * rank * i, dtype=torch.float32).reshape(e, rank, i) + 300,
+        "layers.2.experts.down_lora_b": torch.arange(e * h * rank, dtype=torch.float32).reshape(e, h, rank) + 400,
+    }
+    save_file(tensors, str(path / converter.FUSED_EXPERT_LORA_FILE))
+    return tensors
+
+
+def test_convert_fused_expert_lora_shapes_keys_and_config(tmp_path):
+    input_dir = tmp_path / "input"
+    output_dir = tmp_path / "output"
+    input_dir.mkdir()
+    fused = _write_full_fused_checkpoint(input_dir)
+
+    summary = converter.convert_kt_to_sglang_adapter(
+        input_dir,
+        output_dir,
+        base_model_name_or_path="/models/base",
+        lora_alpha=16,
+    )
+
+    out = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
+    assert summary["tensor_count"] == 12
+    assert summary["rank"] == 3
+    assert summary["target_modules"] == ["gate_proj", "up_proj", "down_proj"]
+
+    key = "model.layers.2.mlp.experts.1.gate_proj.lora_A.weight"
+    assert out[key].shape == (3, 5)
+    torch.testing.assert_close(out[key], fused["layers.2.experts.gate_lora_a"][1])
+
+    key = "model.layers.2.mlp.experts.0.down_proj.lora_B.weight"
+    assert out[key].shape == (5, 3)
+    torch.testing.assert_close(out[key], fused["layers.2.experts.down_lora_b"][0])
+
+    config = converter._load_json(output_dir / converter.ADAPTER_CONFIG_FILE)
+    assert config["peft_type"] == "LORA"
+    assert config["r"] == 3
+    assert config["lora_alpha"] == 16
+    assert config["base_model_name_or_path"] == "/models/base"
+    assert config["target_modules"] == ["gate_proj", "up_proj", "down_proj"]
+
+
+def test_merges_existing_adapter_and_prefers_existing_lora_alpha(tmp_path):
+    input_dir = tmp_path / "input"
+    output_dir = tmp_path / "output"
+    input_dir.mkdir()
+    _write_full_fused_checkpoint(input_dir)
+
+    existing_tensor = torch.ones(3, 5)
+    save_file(
+        {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": existing_tensor,
+        },
+        str(input_dir / converter.ADAPTER_MODEL_FILE),
+    )
+    converter._write_json(
+        input_dir / converter.ADAPTER_CONFIG_FILE,
+        {
+            "peft_type": "LORA",
+            "r": 3,
+            "lora_alpha": 9,
+            "target_modules": ["q_proj"],
+            "bias": "none",
+            "task_type": "CAUSAL_LM",
+            "base_model_name_or_path": "old-base",
+        },
+    )
+
+    summary = converter.convert_kt_to_sglang_adapter(
+        input_dir,
+        output_dir,
+        base_model_name_or_path="/models/base",
+        lora_alpha=16,
+    )
+
+    out = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
+    cleaned_key = "model.layers.0.self_attn.q_proj.lora_A.weight"
+    assert cleaned_key in out
+    torch.testing.assert_close(out[cleaned_key], existing_tensor)
+    assert summary["lora_alpha"] == 9
+
+    config = converter._load_json(output_dir / converter.ADAPTER_CONFIG_FILE)
+    assert config["lora_alpha"] == 9
+    assert config["base_model_name_or_path"] == "/models/base"
+    assert config["target_modules"] == ["q_proj", "gate_proj", "up_proj", "down_proj"]
+
+
+def test_writes_split_expert_and_nonexpert_adapters(tmp_path):
+    input_dir = tmp_path / "input"
+    output_dir = tmp_path / "output"
+    expert_dir = tmp_path / "expert"
+    nonexpert_dir = tmp_path / "nonexpert"
+    input_dir.mkdir()
+    _write_full_fused_checkpoint(input_dir)
+
+    q_proj_tensor = torch.ones(3, 5)
+    o_proj_tensor = torch.full((5, 3), 2.0)
+    save_file(
+        {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": q_proj_tensor,
+            "base_model.model.model.layers.0.self_attn.o_proj.lora_B.weight": o_proj_tensor,
+        },
+        str(input_dir / converter.ADAPTER_MODEL_FILE),
+    )
+    converter._write_json(
+        input_dir / converter.ADAPTER_CONFIG_FILE,
+        {
+            "peft_type": "LORA",
+            "r": 3,
+            "lora_alpha": 9,
+            "target_modules": ["q_proj", "o_proj"],
+            "bias": "none",
+            "task_type": "CAUSAL_LM",
+            "base_model_name_or_path": "old-base",
+        },
+    )
+
+    summary = converter.convert_kt_to_sglang_adapter(
+        input_dir,
+        output_dir,
+        base_model_name_or_path="/models/base",
+        expert_output_dir=expert_dir,
+        nonexpert_output_dir=nonexpert_dir,
+    )
+
+    merged = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
+    expert = load_file(str(expert_dir / converter.ADAPTER_MODEL_FILE))
+    nonexpert = load_file(str(nonexpert_dir / converter.ADAPTER_MODEL_FILE))
+
+    assert summary["tensor_count"] == 14
+    assert summary["split_outputs"]["expert"]["tensor_count"] == 12
+    assert summary["split_outputs"]["nonexpert"]["tensor_count"] == 2
+    assert set(merged) == set(expert) | set(nonexpert)
+    assert set(expert).isdisjoint(nonexpert)
+    assert all(".mlp.experts." in key for key in expert)
+    assert not any(".mlp.experts." in key for key in nonexpert)
+
+    cleaned_q_proj_key = "model.layers.0.self_attn.q_proj.lora_A.weight"
+    cleaned_o_proj_key = "model.layers.0.self_attn.o_proj.lora_B.weight"
+    torch.testing.assert_close(nonexpert[cleaned_q_proj_key], q_proj_tensor)
+    torch.testing.assert_close(nonexpert[cleaned_o_proj_key], o_proj_tensor)
+
+    expert_config = converter._load_json(expert_dir / converter.ADAPTER_CONFIG_FILE)
+    nonexpert_config = converter._load_json(nonexpert_dir / converter.ADAPTER_CONFIG_FILE)
+    assert expert_config["target_modules"] == ["gate_proj", "up_proj", "down_proj"]
+    assert nonexpert_config["target_modules"] == ["q_proj", "o_proj"]
+    assert expert_config["base_model_name_or_path"] == "/models/base"
+    assert nonexpert_config["base_model_name_or_path"] == "/models/base"
+
+
+def test_requires_lora_alpha_without_input_config(tmp_path):
+    input_dir = tmp_path / "input"
+    input_dir.mkdir()
+    _write_full_fused_checkpoint(input_dir)
+
+    with pytest.raises(ValueError, match="pass --lora-alpha"):
+        converter.convert_kt_to_sglang_adapter(
+            input_dir,
+            tmp_path / "output",
+            base_model_name_or_path="/models/base",
+        )
+
+
+def test_rejects_inconsistent_rank(tmp_path):
+    input_dir = tmp_path / "input"
+    input_dir.mkdir()
+    save_file(
+        {
+            "layers.0.experts.gate_lora_a": torch.zeros(2, 3, 5),
+            "layers.0.experts.gate_lora_b": torch.zeros(2, 7, 4),
+        },
+        str(input_dir / converter.FUSED_EXPERT_LORA_FILE),
+    )
+
+    with pytest.raises(ValueError, match="Inconsistent LoRA ranks"):
+        converter.convert_kt_to_sglang_adapter(
+            input_dir,
+            tmp_path / "output",
+            base_model_name_or_path="/models/base",
+            lora_alpha=8,
+        )
+
+
+def test_rejects_unexpected_fused_key(tmp_path):
+    input_dir = tmp_path / "input"
+    input_dir.mkdir()
+    save_file(
+        {"layers.0.experts.unknown_lora_a": torch.zeros(2, 3, 5)},
+        str(input_dir / converter.FUSED_EXPERT_LORA_FILE),
+    )
+
+    with pytest.raises(ValueError, match="Unsupported KT fused expert LoRA tensor"):
+        converter.convert_kt_to_sglang_adapter(
+            input_dir,
+            tmp_path / "output",
+            base_model_name_or_path="/models/base",
+            lora_alpha=8,
+        )
+
+
+def test_rejects_nonempty_output_without_overwrite(tmp_path):
+    input_dir = tmp_path / "input"
+    output_dir = tmp_path / "output"
+    input_dir.mkdir()
+    output_dir.mkdir()
+    (output_dir / "existing.txt").write_text("do not remove", encoding="utf-8")
+    _write_full_fused_checkpoint(input_dir)
+
+    with pytest.raises(FileExistsError, match="Output directory is not empty"):
+        converter.convert_kt_to_sglang_adapter(
+            input_dir,
+            output_dir,
+            base_model_name_or_path="/models/base",
+            lora_alpha=8,
+        )
+
+    converter.convert_kt_to_sglang_adapter(
+        input_dir,
+        output_dir,
+        base_model_name_or_path="/models/base",
+        lora_alpha=8,
+        overwrite=True,
+    )
+    assert not (output_dir / "existing.txt").exists()
+    assert (output_dir / converter.ADAPTER_MODEL_FILE).exists()
+
+
+def test_rejects_output_same_as_input(tmp_path):
+    input_dir = tmp_path / "input"
+    input_dir.mkdir()
+    _write_full_fused_checkpoint(input_dir)
+
+    with pytest.raises(ValueError, match="different from input"):
+        converter.convert_kt_to_sglang_adapter(
+            input_dir,
+            input_dir,
+            base_model_name_or_path="/models/base",
+            lora_alpha=8,
+            overwrite=True,
+        )
+
+
+def test_rejects_output_ancestor_of_input(tmp_path):
+    run_dir = tmp_path / "run"
+    input_dir = run_dir / "adapter"
+    output_dir = run_dir
+    input_dir.mkdir(parents=True)
+    _write_full_fused_checkpoint(input_dir)
+
+    with pytest.raises(ValueError, match="ancestor/descendant"):
+        converter.convert_kt_to_sglang_adapter(
+            input_dir,
+            output_dir,
+            base_model_name_or_path="/models/base",
+            lora_alpha=8,
+            overwrite=True,
+        )
+
+
+def test_rejects_output_descendant_of_input(tmp_path):
+    input_dir = tmp_path / "input"
+    output_dir = input_dir / "output"
+    input_dir.mkdir()
+    _write_full_fused_checkpoint(input_dir)
+
+    with pytest.raises(ValueError, match="ancestor/descendant"):
+        converter.convert_kt_to_sglang_adapter(
+            input_dir,
+            output_dir,
+            base_model_name_or_path="/models/base",
+            lora_alpha=8,
+            overwrite=True,
+        )
+
+
+def test_rejects_split_output_ancestor_relationship(tmp_path):
+    input_dir = tmp_path / "input"
+    output_dir = tmp_path / "output"
+    expert_dir = tmp_path / "split" / "expert"
+    nonexpert_dir = tmp_path / "split"
+    input_dir.mkdir()
+    _write_full_fused_checkpoint(input_dir)
+    save_file(
+        {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": torch.ones(3, 5),
+        },
+        str(input_dir / converter.ADAPTER_MODEL_FILE),
+    )
+    converter._write_json(
+        input_dir / converter.ADAPTER_CONFIG_FILE,
+        {
+            "peft_type": "LORA",
+            "r": 3,
+            "lora_alpha": 9,
+            "target_modules": ["q_proj"],
+            "bias": "none",
+            "task_type": "CAUSAL_LM",
+            "base_model_name_or_path": "old-base",
+        },
+    )
+
+    with pytest.raises(ValueError, match="ancestor/descendant"):
+        converter.convert_kt_to_sglang_adapter(
+            input_dir,
+            output_dir,
+            base_model_name_or_path="/models/base",
+            expert_output_dir=expert_dir,
+            nonexpert_output_dir=nonexpert_dir,
+        )
+
+
+def test_rejects_mismatched_nonexpert_rank(tmp_path):
+    input_dir = tmp_path / "input"
+    output_dir = tmp_path / "output"
+    input_dir.mkdir()
+    _write_full_fused_checkpoint(input_dir, rank=3)
+    save_file(
+        {
+            "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": torch.ones(4, 5),
+        },
+        str(input_dir / converter.ADAPTER_MODEL_FILE),
+    )
+    converter._write_json(
+        input_dir / converter.ADAPTER_CONFIG_FILE,
+        {
+            "peft_type": "LORA",
+            "r": 4,
+            "lora_alpha": 9,
+            "target_modules": ["q_proj"],
+            "bias": "none",
+            "task_type": "CAUSAL_LM",
+            "base_model_name_or_path": "old-base",
+        },
+    )
+
+    with pytest.raises(ValueError, match="rank mismatch"):
+        converter.convert_kt_to_sglang_adapter(
+            input_dir,
+            output_dir,
+            base_model_name_or_path="/models/base",
+        )

kt-kernel/test/per_commit/test_convert_kt_to_sglang_adapter_integration.py

@@ -0,0 +1,242 @@
+import importlib.util
+import json
+import os
+import re
+import shutil
+import sys
+import time
+from pathlib import Path
+
+import pytest
+import torch
+from safetensors.torch import load_file
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
+
+register_cpu_ci(est_time=60, suite="default")
+
+KT_ADAPTER_ENV = "KT_LORA_ADAPTER_DIR"
+KT_BASE_MODEL_ENV = "KT_LORA_BASE_MODEL"
+KT_ALPHA_ENV = "KT_LORA_ALPHA"
+KT_LARGE_ADAPTER_ENV = "KT_LORA_LARGE_ADAPTER_DIR"
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+SCRIPT_PATH = (
+    Path(__file__).resolve().parents[2]
+    / "scripts"
+    / "convert_kt_to_sglang_adapter.py"
+)
+SPEC = importlib.util.spec_from_file_location("convert_kt_to_sglang_adapter", SCRIPT_PATH)
+converter = importlib.util.module_from_spec(SPEC)
+assert SPEC.loader is not None
+SPEC.loader.exec_module(converter)
+
+SGLANG_EXPERT_KEY_RE = re.compile(
+    r"^model\.layers\.(\d+)\.mlp\.experts\.(\d+)\."
+    r"(gate_proj|up_proj|down_proj)\.lora_[AB]\.weight$"
+)
+
+
+def _required_adapter_dir(env_name: str) -> Path:
+    value = os.environ.get(env_name)
+    if not value:
+        pytest.skip(f"Set {env_name} to run real adapter integration tests.")
+
+    path = Path(value).expanduser().resolve()
+    if not path.is_dir():
+        pytest.fail(f"{env_name} is not a directory: {path}")
+    if not (path / converter.FUSED_EXPERT_LORA_FILE).is_file():
+        pytest.fail(f"{env_name} must contain {converter.FUSED_EXPERT_LORA_FILE}: {path}")
+    return path
+
+
+def _base_model_name_or_path() -> str:
+    value = os.environ.get(KT_BASE_MODEL_ENV)
+    if not value:
+        pytest.fail(f"Set {KT_BASE_MODEL_ENV} before running real adapter integration tests.")
+    return value
+
+
+def _optional_lora_alpha() -> float | None:
+    value = os.environ.get(KT_ALPHA_ENV)
+    if value in (None, ""):
+        return None
+    try:
+        return float(value)
+    except ValueError:
+        pytest.fail(f"{KT_ALPHA_ENV} must be numeric, got: {value!r}")
+
+
+def _lora_alpha_for_input(input_dir: Path) -> float | None:
+    alpha = _optional_lora_alpha()
+    if (input_dir / converter.ADAPTER_CONFIG_FILE).exists():
+        return alpha
+    if alpha is None:
+        pytest.fail(
+            f"{input_dir} has no {converter.ADAPTER_CONFIG_FILE}; set {KT_ALPHA_ENV}."
+        )
+    return alpha
+
+
+def _convert_real_adapter(input_dir: Path, tmp_path: Path, output_name: str = "output") -> tuple[Path, dict]:
+    output_dir = tmp_path / output_name
+    summary = converter.convert_kt_to_sglang_adapter(
+        input_dir,
+        output_dir,
+        base_model_name_or_path=_base_model_name_or_path(),
+        lora_alpha=_lora_alpha_for_input(input_dir),
+    )
+    return output_dir, summary
+
+
+def _load_json(path: Path) -> dict:
+    with path.open("r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _link_or_copy(source: Path, dest: Path) -> None:
+    try:
+        os.symlink(source, dest)
+    except OSError:
+        shutil.copy2(source, dest)
+
+
+def _assert_config_shape(output_dir: Path, summary: dict) -> dict:
+    config = _load_json(output_dir / converter.ADAPTER_CONFIG_FILE)
+    assert config["peft_type"] == "LORA"
+    assert config["r"] == summary["rank"]
+    assert config["lora_alpha"] == summary["lora_alpha"]
+    assert config["base_model_name_or_path"] == _base_model_name_or_path()
+    assert {"gate_proj", "up_proj", "down_proj"}.issubset(config["target_modules"])
+    return config
+
+
+def _assert_fused_tensors_preserved(input_dir: Path, output_dir: Path) -> int:
+    fused_tensors = load_file(str(input_dir / converter.FUSED_EXPERT_LORA_FILE))
+    output_tensors = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
+
+    checked = 0
+    for input_key, input_tensor in sorted(fused_tensors.items()):
+        match = converter.KT_FUSED_KEY_RE.match(input_key)
+        assert match is not None, input_key
+        layer_idx, kt_name = match.groups()
+        assert kt_name in converter.KT_NAME_MAP, input_key
+        assert input_tensor.dim() == 3, input_key
+
+        proj_name, lora_name, _rank_dim = converter.KT_NAME_MAP[kt_name]
+        for expert_idx in range(input_tensor.shape[0]):
+            output_key = (
+                f"model.layers.{layer_idx}.mlp.experts.{expert_idx}."
+                f"{proj_name}.{lora_name}.weight"
+            )
+            assert SGLANG_EXPERT_KEY_RE.match(output_key), output_key
+            assert output_key in output_tensors
+            assert output_tensors[output_key].shape == input_tensor[expert_idx].shape
+            assert output_tensors[output_key].dtype == input_tensor.dtype
+            assert torch.equal(output_tensors[output_key], input_tensor[expert_idx].cpu())
+            checked += 1
+
+    assert checked > 0
+    return checked
+
+
+@pytest.mark.requires_model
+def test_real_adapter_conversion_preserves_fused_tensors_and_config(tmp_path):
+    input_dir = _required_adapter_dir(KT_ADAPTER_ENV)
+
+    output_dir, summary = _convert_real_adapter(input_dir, tmp_path)
+
+    checked = _assert_fused_tensors_preserved(input_dir, output_dir)
+    config = _assert_config_shape(output_dir, summary)
+    output_tensors = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
+
+    assert summary["tensor_count"] == len(output_tensors)
+    assert checked <= summary["tensor_count"]
+    assert config["target_modules"] == summary["target_modules"]
+
+
+@pytest.mark.requires_model
+def test_real_adapter_directory_merges_existing_adapter_model(tmp_path):
+    input_dir = _required_adapter_dir(KT_ADAPTER_ENV)
+    existing_adapter_path = input_dir / converter.ADAPTER_MODEL_FILE
+    if not existing_adapter_path.exists():
+        pytest.skip(f"{input_dir} has no {converter.ADAPTER_MODEL_FILE} to merge.")
+
+    output_dir, _summary = _convert_real_adapter(input_dir, tmp_path)
+    input_tensors = load_file(str(existing_adapter_path))
+    output_tensors = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
+
+    for input_key, input_tensor in input_tensors.items():
+        output_key = converter._clean_adapter_key(input_key)
+        assert output_key in output_tensors
+        assert output_tensors[output_key].shape == input_tensor.shape
+        assert output_tensors[output_key].dtype == input_tensor.dtype
+        assert torch.equal(output_tensors[output_key], input_tensor.cpu())
+
+
+@pytest.mark.requires_model
+def test_real_fused_conversion_without_input_config_uses_env_alpha(tmp_path):
+    input_dir = _required_adapter_dir(KT_ADAPTER_ENV)
+    alpha = _optional_lora_alpha()
+    if alpha is None:
+        pytest.skip(f"Set {KT_ALPHA_ENV} to validate conversion without input config.")
+
+    no_config_input = tmp_path / "input_without_config"
+    no_config_input.mkdir()
+    _link_or_copy(
+        input_dir / converter.FUSED_EXPERT_LORA_FILE,
+        no_config_input / converter.FUSED_EXPERT_LORA_FILE,
+    )
+    existing_adapter_path = input_dir / converter.ADAPTER_MODEL_FILE
+    if existing_adapter_path.exists():
+        _link_or_copy(existing_adapter_path, no_config_input / converter.ADAPTER_MODEL_FILE)
+
+    output_dir, summary = _convert_real_adapter(no_config_input, tmp_path, "output_without_config")
+
+    assert summary["lora_alpha"] == alpha
+    config = _assert_config_shape(output_dir, summary)
+    assert config["lora_alpha"] == alpha
+    _assert_fused_tensors_preserved(no_config_input, output_dir)
+
+
+@pytest.mark.requires_model
+def test_sglang_lora_config_loader_accepts_converted_adapter(tmp_path):
+    input_dir = _required_adapter_dir(KT_ADAPTER_ENV)
+    output_dir, summary = _convert_real_adapter(input_dir, tmp_path)
+
+    sglang_python = REPO_ROOT / "third_party" / "sglang" / "python"
+    sys.path.insert(0, str(sglang_python))
+    try:
+        from sglang.srt.lora.lora_config import LoRAConfig
+    except Exception as exc:
+        pytest.fail(f"Unable to import SGLang LoRAConfig: {exc}")
+
+    lora_config = LoRAConfig(str(output_dir))
+    output_tensors = load_file(str(output_dir / converter.ADAPTER_MODEL_FILE))
+
+    assert lora_config.r == summary["rank"]
+    assert lora_config.lora_alpha == summary["lora_alpha"]
+    assert lora_config.target_modules == summary["target_modules"]
+    assert len(output_tensors) == summary["tensor_count"]
+
+
+@pytest.mark.requires_model
+def test_large_adapter_conversion_smoke(tmp_path, record_property):
+    input_dir = _required_adapter_dir(KT_LARGE_ADAPTER_ENV)
+
+    start_time = time.perf_counter()
+    output_dir, summary = _convert_real_adapter(input_dir, tmp_path, "large_output")
+    duration_seconds = time.perf_counter() - start_time
+
+    output_path = output_dir / converter.ADAPTER_MODEL_FILE
+    output_tensors = load_file(str(output_path))
+    config = _assert_config_shape(output_dir, summary)
+
+    record_property("conversion_seconds", round(duration_seconds, 3))
+    record_property("output_bytes", output_path.stat().st_size)
+    record_property("tensor_count", summary["tensor_count"])
+
+    assert len(output_tensors) == summary["tensor_count"]
+    assert config["target_modules"] == summary["target_modules"]

kt-kernel/test/per_commit/test_cuda_placeholder.py

@@ -0,0 +1,36 @@
+"""CUDA backend tests for KT-Kernel (Placeholder).
+
+This file is a placeholder for future CUDA backend tests.
+Currently, KT-Kernel focuses on CPU optimizations (Intel AMX/AVX512).
+
+To implement CUDA tests:
+1. Add actual test functions with @pytest.mark.cuda
+2. Update the estimated time in register_cuda_ci()
+3. Implement CUDA-specific initialization and validation tests
+"""
+
+import os
+import sys
+
+# Add parent directory to path for CI registration
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cuda_ci
+
+# Register this test for CUDA CI (estimated time: 10 seconds, placeholder)
+# Update suite name when implementing: currently using "stage-a-test-1"
+register_cuda_ci(est_time=10, suite="stage-a-test-1")
+
+
+def test_cuda_placeholder():
+    """Placeholder test for CUDA backend.
+
+    TODO: Implement actual CUDA tests when CUDA support is added to kt-kernel.
+    """
+    # Currently a no-op placeholder
+    pass
+
+
+if __name__ == "__main__":
+    # Allow running standalone (required by test runner)
+    print("⚠ CUDA tests are not yet implemented (placeholder)")
+    print("✓ Placeholder test passed")

kt-kernel/test/per_commit/test_load_experts_count_guard.py

@@ -0,0 +1,91 @@
+"""Regression tests for SafeTensorLoader.load_experts expert-count guard.
+
+After the discovery loop, ``max_experts_count`` holds the highest expert index
+found, i.e. ``(expert count - 1)``. It is ``-1`` only when no experts exist for
+the key, so the guard must reject ``-1`` (no experts), not ``0`` (exactly one
+expert). The previous ``== 0`` check was an off-by-one that falsely raised
+"No experts found" for a single-expert layer and silently returned empty weight
+lists for the genuine zero-expert case.
+
+The loader module is imported as a top-level module (its directory is added to
+sys.path) so the test does not pull in the compiled kt_kernel_ext extension via
+utils/__init__.py.
+"""
+
+import os
+import sys
+import unittest
+
+# Register this test for CPU CI.
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="default")
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python", "utils"))
+import loader
+
+SafeTensorLoader = loader.SafeTensorLoader
+
+
+class _FakeTensor:
+    """Stand-in for a loaded tensor; load_experts only calls ``.numpy()``."""
+
+    def numpy(self):
+        return 0
+
+
+class _StubSafeTensorLoader(SafeTensorLoader):
+    """SafeTensorLoader backed by an explicit key set instead of real files."""
+
+    def __init__(self, keys):
+        # Bypass the filesystem/safetensors scan in the real __init__.
+        self.tensor_file_map = {k: "stub.safetensors" for k in keys}
+        self.file_handle_map = {}
+        self.tensor_type_map = {}
+        self.tensor_device_map = {}
+
+    def load_tensor(self, key: str, device: str = "cpu"):
+        assert key in self.tensor_file_map, f"unexpected key requested: {key}"
+        return _FakeTensor()
+
+
+def _expert_keys(base_key, n_experts, n_numa=1):
+    """Build the NUMA-sharded expert key set that load_experts expects."""
+    keys = []
+    for proj in ("ffn_up_exps", "ffn_gate_exps", "ffn_down_exps"):
+        for expert in range(n_experts):
+            for numa in range(n_numa):
+                prefix = f"{base_key}.{proj}.{expert}.numa.{numa}"
+                keys.append(f"{prefix}.weight")
+                keys.append(f"{prefix}.scale")
+    return keys
+
+
+class TestLoadExpertsCountGuard(unittest.TestCase):
+    def test_single_expert_is_loaded(self):
+        """A layer with exactly one expert must load, not raise."""
+        loader = _StubSafeTensorLoader(_expert_keys("blk.0", n_experts=1))
+        result = loader.load_experts("blk.0")
+        for proj in ("up", "gate", "down"):
+            self.assertEqual(len(result[proj]), 1, f"{proj}: expected 1 numa group")
+            self.assertEqual(len(result[proj][0]), 1, f"{proj}: expected 1 expert")
+
+    def test_zero_experts_raises(self):
+        """No experts under the key must raise, not silently return empties."""
+        loader = _StubSafeTensorLoader(["blk.0.attn_norm.weight"])
+        with self.assertRaises(ValueError):
+            loader.load_experts("blk.0")
+
+    def test_multiple_experts_and_numa_counts(self):
+        """Counts are correct across several experts and NUMA shards."""
+        loader = _StubSafeTensorLoader(_expert_keys("blk.0", n_experts=3, n_numa=2))
+        result = loader.load_experts("blk.0")
+        for proj in ("up", "gate", "down"):
+            self.assertEqual(len(result[proj]), 2, f"{proj}: expected 2 numa groups")
+            for numa_group in result[proj]:
+                self.assertEqual(len(numa_group), 3, f"{proj}: expected 3 experts")
+
+
+if __name__ == "__main__":
+    unittest.main()

kt-kernel/test/per_commit/test_moe_amx_accuracy_int4.py

@@ -0,0 +1,205 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""AMX MOE INT4 accuracy tests for KT-Kernel.
+
+Tests accuracy of AMX-accelerated INT4 MOE operations against torch reference.
+"""
+
+import os
+import sys
+import pytest
+
+# Add parent directory to path for CI registration
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
+# Register this test for CPU CI with estimated runtime of 120 seconds
+register_cpu_ci(est_time=120, suite="default")
+
+# Check if dependencies are available
+try:
+    import torch
+    import kt_kernel  # Import kt_kernel first to register kt_kernel_ext
+
+    kt_kernel_ext = kt_kernel.kt_kernel_ext  # Access the extension module
+    HAS_DEPS = True
+except ImportError as e:
+    HAS_DEPS = False
+    import_error = str(e)
+
+# Test parameters (from original test_moe_amx.py)
+expert_num = 256
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 25600
+num_experts_per_tok = 8
+qlen = 1
+layer_num = 1
+validation_iter = 2
+physical_to_logical_map = None
+
+
+def act_fn(x):
+    """Activation function for MoE."""
+    return x / (1.0 + torch.exp(-x))
+
+
+def mlp_torch(input, gate_proj, up_proj, down_proj):
+    """PyTorch reference implementation of MLP."""
+    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):
+    """PyTorch reference implementation of MoE."""
+    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
+
+
+@pytest.mark.cpu
+def test_moe_amx_int4_accuracy():
+    """Test AMX INT4 MOE accuracy against PyTorch reference implementation."""
+    if not HAS_DEPS:
+        pytest.skip(f"Dependencies not available: {import_error}")
+
+    global physical_to_logical_map
+    physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous()
+
+    CPUInfer = kt_kernel_ext.CPUInfer(60)
+
+    with torch.inference_mode(mode=True):
+        # Initialize MoE layers
+        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_
+
+        # Initialize INT4 MOE
+        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()
+
+        # Run validation iterations
+        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_data = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            input_data = input_data / 100
+
+            # Run AMX MOE
+            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()
+
+            # Run torch reference
+            t_output = moe_torch(input_data, expert_ids, weights, gate_proj, up_proj, down_proj)
+
+            # Calculate relative difference
+            diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
+            print(f"Iteration {i}, diff = {diff:.6f}")
+
+            # INT4 should have diff < 0.35
+            assert diff < 0.35, f"INT4 accuracy test failed: diff={diff:.6f} >= 0.35"
+
+
+def run_all_tests():
+    """Run all tests in this file (for standalone execution)."""
+    if not HAS_DEPS:
+        print(f"⚠ Dependencies not available: {import_error}")
+        print("Skipping AMX MOE INT4 accuracy tests")
+        return
+
+    try:
+        print("Running AMX MOE INT4 accuracy test...")
+        test_moe_amx_int4_accuracy()
+        print("✓ AMX MOE INT4 accuracy test passed")
+        print("\n✓ All tests passed!")
+    except Exception as e:
+        print(f"\n✗ Test failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    run_all_tests()

kt-kernel/test/per_commit/test_moe_amx_accuracy_int4_1.py

@@ -0,0 +1,205 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""AMX MOE INT4_1 accuracy tests for KT-Kernel.
+
+Tests accuracy of AMX-accelerated INT4_1 MOE operations against torch reference.
+"""
+
+import os
+import sys
+import pytest
+
+# Add parent directory to path for CI registration
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
+# Register this test for CPU CI with estimated runtime of 120 seconds
+register_cpu_ci(est_time=120, suite="default")
+
+# Check if dependencies are available
+try:
+    import torch
+    import kt_kernel  # Import kt_kernel first to register kt_kernel_ext
+
+    kt_kernel_ext = kt_kernel.kt_kernel_ext  # Access the extension module
+    HAS_DEPS = True
+except ImportError as e:
+    HAS_DEPS = False
+    import_error = str(e)
+
+# Test parameters (from original test_moe_amx.py)
+expert_num = 256
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 25600
+num_experts_per_tok = 8
+qlen = 1
+layer_num = 1
+validation_iter = 2
+physical_to_logical_map = None
+
+
+def act_fn(x):
+    """Activation function for MoE."""
+    return x / (1.0 + torch.exp(-x))
+
+
+def mlp_torch(input, gate_proj, up_proj, down_proj):
+    """PyTorch reference implementation of MLP."""
+    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):
+    """PyTorch reference implementation of MoE."""
+    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
+
+
+@pytest.mark.cpu
+def test_moe_amx_int4_1_accuracy():
+    """Test AMX INT4_1 MOE accuracy against PyTorch reference implementation."""
+    if not HAS_DEPS:
+        pytest.skip(f"Dependencies not available: {import_error}")
+
+    global physical_to_logical_map
+    physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous()
+
+    CPUInfer = kt_kernel_ext.CPUInfer(60)
+
+    with torch.inference_mode(mode=True):
+        # Initialize MoE layers
+        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_
+
+        # Initialize INT4_1 MOE
+        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()
+
+        # Run validation iterations
+        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_data = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            input_data = input_data / 100
+
+            # Run AMX MOE
+            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()
+
+            # Run torch reference
+            t_output = moe_torch(input_data, expert_ids, weights, gate_proj, up_proj, down_proj)
+
+            # Calculate relative difference
+            diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
+            print(f"Iteration {i}, diff = {diff:.6f}")
+
+            # INT4_1 should have diff < 0.35
+            assert diff < 0.35, f"INT4_1 accuracy test failed: diff={diff:.6f} >= 0.35"
+
+
+def run_all_tests():
+    """Run all tests in this file (for standalone execution)."""
+    if not HAS_DEPS:
+        print(f"⚠ Dependencies not available: {import_error}")
+        print("Skipping AMX MOE INT4_1 accuracy tests")
+        return
+
+    try:
+        print("Running AMX MOE INT4_1 accuracy test...")
+        test_moe_amx_int4_1_accuracy()
+        print("✓ AMX MOE INT4_1 accuracy test passed")
+        print("\n✓ All tests passed!")
+    except Exception as e:
+        print(f"\n✗ Test failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    run_all_tests()

kt-kernel/test/per_commit/test_moe_amx_accuracy_int4_1k.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""AMX MOE INT4_1K accuracy tests for KT-Kernel.
+
+Tests accuracy of AMX-accelerated INT4_1K group quantization MOE operations against torch reference.
+"""
+
+import os
+import sys
+import pytest
+
+# Add parent directory to path for CI registration
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
+# Register this test for CPU CI with estimated runtime of 120 seconds
+register_cpu_ci(est_time=120, suite="default")
+
+# Check if dependencies are available
+try:
+    import torch
+    import kt_kernel  # Import kt_kernel first to register kt_kernel_ext
+
+    kt_kernel_ext = kt_kernel.kt_kernel_ext  # Access the extension module
+    HAS_DEPS = True
+except ImportError as e:
+    HAS_DEPS = False
+    import_error = str(e)
+
+# Test parameters (from original test_moe_amx.py)
+expert_num = 256
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 25600
+num_experts_per_tok = 8
+qlen = 1
+layer_num = 1
+validation_iter = 2
+k_group_size = 64
+physical_to_logical_map = None
+
+
+def act_fn(x):
+    """Activation function for MoE."""
+    return x / (1.0 + torch.exp(-x))
+
+
+def mlp_torch(input, gate_proj, up_proj, down_proj):
+    """PyTorch reference implementation of MLP."""
+    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):
+    """PyTorch reference implementation of MoE."""
+    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
+
+
+@pytest.mark.cpu
+def test_moe_amx_int4_1k_accuracy():
+    """Test AMX INT4_1K MOE accuracy against PyTorch reference implementation."""
+    if not HAS_DEPS:
+        pytest.skip(f"Dependencies not available: {import_error}")
+
+    global physical_to_logical_map
+    physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous()
+
+    CPUInfer = kt_kernel_ext.CPUInfer(60)
+
+    with torch.inference_mode(mode=True):
+        # Initialize MoE layers
+        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_
+
+        # Configure INT4_1K quantization settings
+        config.quant_config.bits = 4
+        config.quant_config.group_size = k_group_size
+        config.quant_config.zero_point = True
+
+        # Initialize INT4_1K MOE
+        moe = kt_kernel_ext.moe.AMXInt4_1KGroup_MOE(config)
+        CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+        CPUInfer.sync()
+
+        # Run validation iterations
+        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_data = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            input_data = input_data / 100
+
+            # Run AMX MOE
+            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()
+
+            # Run torch reference
+            t_output = moe_torch(input_data, expert_ids, weights, gate_proj, up_proj, down_proj)
+
+            # Calculate relative difference
+            diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
+            print(f"Iteration {i}, diff = {diff:.6f}")
+
+            # INT4_1K should have diff < 0.35
+            assert diff < 0.35, f"INT4_1K accuracy test failed: diff={diff:.6f} >= 0.35"
+
+
+def run_all_tests():
+    """Run all tests in this file (for standalone execution)."""
+    if not HAS_DEPS:
+        print(f"⚠ Dependencies not available: {import_error}")
+        print("Skipping AMX MOE INT4_1K accuracy tests")
+        return
+
+    try:
+        print("Running AMX MOE INT4_1K accuracy test...")
+        test_moe_amx_int4_1k_accuracy()
+        print("✓ AMX MOE INT4_1K accuracy test passed")
+        print("\n✓ All tests passed!")
+    except Exception as e:
+        print(f"\n✗ Test failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    run_all_tests()

kt-kernel/test/per_commit/test_moe_amx_accuracy_int8.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""AMX MOE INT8 accuracy tests for KT-Kernel.
+
+Tests accuracy of AMX-accelerated INT8 MOE operations against torch reference.
+"""
+
+import os
+import sys
+import pytest
+
+# Add parent directory to path for CI registration
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
+# Register this test for CPU CI with estimated runtime of 120 seconds
+register_cpu_ci(est_time=120, suite="default")
+
+# Check if dependencies are available
+try:
+    import torch
+    import kt_kernel  # Import kt_kernel first to register kt_kernel_ext
+
+    kt_kernel_ext = kt_kernel.kt_kernel_ext  # Access the extension module
+    HAS_DEPS = True
+except ImportError as e:
+    HAS_DEPS = False
+    import_error = str(e)
+
+# Test parameters (from original test_moe_amx.py)
+expert_num = 256
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 25600
+num_experts_per_tok = 8
+qlen = 1
+layer_num = 1
+validation_iter = 2
+physical_to_logical_map = None
+
+
+def act_fn(x):
+    """Activation function for MoE."""
+    return x / (1.0 + torch.exp(-x))
+
+
+def mlp_torch(input, gate_proj, up_proj, down_proj):
+    """PyTorch reference implementation of MLP."""
+    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):
+    """PyTorch reference implementation of MoE."""
+    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
+
+
+@pytest.mark.cpu
+def test_moe_amx_int8_accuracy():
+    """Test AMX INT8 MOE accuracy against PyTorch reference implementation."""
+    if not HAS_DEPS:
+        pytest.skip(f"Dependencies not available: {import_error}")
+
+    global physical_to_logical_map
+    physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous()
+
+    CPUInfer = kt_kernel_ext.CPUInfer(60)
+
+    with torch.inference_mode(mode=True):
+        # Initialize MoE layers
+        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_
+
+        # Initialize INT8 MOE
+        moe = kt_kernel_ext.moe.AMXInt8_MOE(config)
+        CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+        CPUInfer.sync()
+
+        # Run validation iterations
+        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_data = torch.randn((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+            input_data = input_data / 100
+
+            # Run AMX MOE
+            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()
+
+            # Run torch reference
+            t_output = moe_torch(input_data, expert_ids, weights, gate_proj, up_proj, down_proj)
+
+            # Calculate relative difference
+            diff = torch.mean(torch.abs(output - t_output)) / torch.mean(torch.abs(t_output))
+            print(f"Iteration {i}, diff = {diff:.6f}")
+
+            # INT8 should have diff < 0.05
+            assert diff < 0.05, f"INT8 accuracy test failed: diff={diff:.6f} >= 0.05"
+
+
+def run_all_tests():
+    """Run all tests in this file (for standalone execution)."""
+    if not HAS_DEPS:
+        print(f"⚠ Dependencies not available: {import_error}")
+        print("Skipping AMX MOE INT8 accuracy tests")
+        return
+
+    try:
+        print("Running AMX MOE INT8 accuracy test...")
+        test_moe_amx_int8_accuracy()
+        print("✓ AMX MOE INT8 accuracy test passed")
+        print("\n✓ All tests passed!")
+    except Exception as e:
+        print(f"\n✗ Test failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    run_all_tests()

kt-kernel/test/per_commit/test_moe_amx_bench_int4.py

@@ -0,0 +1,317 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""AMX MOE INT4 benchmark tests for KT-Kernel.
+
+Benchmarks performance (bandwidth and FLOPS) of AMX-accelerated INT4 MOE operations.
+"""
+
+import os
+import sys
+import time
+import json
+import subprocess
+import platform
+import pytest
+
+# Add parent directory to path for CI registration
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
+# Register this test for CPU CI with estimated runtime of 300 seconds
+register_cpu_ci(est_time=300, suite="default")
+
+# Check if dependencies are available
+try:
+    import torch
+    import kt_kernel  # Import kt_kernel first to register kt_kernel_ext
+
+    kt_kernel_ext = kt_kernel.kt_kernel_ext  # Access the extension module
+    from tqdm import tqdm
+
+    HAS_DEPS = True
+except ImportError as e:
+    HAS_DEPS = False
+    import_error = str(e)
+
+# Test parameters (from original bench_moe_amx.py)
+expert_num = 16
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 25600
+num_experts_per_tok = 8
+layer_num = 2
+qlen = 2048
+warm_up_iter = 1000
+test_iter = 2000
+
+# Worker configuration
+worker_config_dict = {
+    "subpool_count": 2,
+    "subpool_numa_map": [0, 1],
+    "subpool_thread_count": [30, 30],
+}
+CPUINFER_PARAM = 60
+
+
+def get_git_commit():
+    """Get current git commit information."""
+    result = {}
+    try:
+        commit = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("utf-8").strip()
+        commit_msg = subprocess.check_output(["git", "log", "-1", "--pretty=%B"]).decode("utf-8").strip()
+        result["commit"] = commit
+        result["commit_message"] = commit_msg
+
+        dirty_output = subprocess.check_output(["git", "status", "--porcelain"]).decode("utf-8").strip()
+        if dirty_output:
+            result["dirty"] = True
+            result["dirty_files"] = dirty_output.splitlines()
+        else:
+            result["dirty"] = False
+    except Exception as e:
+        result["commit"] = None
+        result["commit_message"] = None
+        result["dirty"] = None
+        result["error"] = str(e)
+    return result
+
+
+def get_system_info():
+    """Get system information including CPU model, memory, cores, and sockets."""
+    info = {}
+    uname = platform.uname()
+    info["system_name"] = uname.system
+    info["node_name"] = uname.node
+
+    # Get CPU model (Linux only)
+    cpu_model = None
+    if os.path.exists("/proc/cpuinfo"):
+        try:
+            with open("/proc/cpuinfo", "r") as f:
+                for line in f:
+                    if "model name" in line:
+                        cpu_model = line.split(":", 1)[1].strip()
+                        break
+        except Exception as e:
+            cpu_model = f"Error: {e}"
+    info["cpu_model"] = cpu_model
+
+    # Get memory size in GB (Linux only)
+    mem_total_gb = None
+    if os.path.exists("/proc/meminfo"):
+        try:
+            with open("/proc/meminfo", "r") as f:
+                for line in f:
+                    if "MemTotal" in line:
+                        mem_kb = float(line.split(":", 1)[1].split()[0])
+                        mem_total_gb = round(mem_kb / (1024 * 1024), 2)
+                        break
+        except Exception as e:
+            mem_total_gb = f"Error: {e}"
+    info["memory_size_GB"] = mem_total_gb
+
+    # Get CPU core count
+    info["cpu_core_count"] = os.cpu_count()
+
+    # Get socket count
+    sockets = set()
+    if os.path.exists("/proc/cpuinfo"):
+        try:
+            with open("/proc/cpuinfo", "r") as f:
+                for line in f:
+                    if "physical id" in line:
+                        sockets.add(line.split(":", 1)[1].strip())
+        except Exception as e:
+            sockets = set()
+    info["cpu_socket_count"] = len(sockets) if len(sockets) > 0 else 1
+
+    return info
+
+
+def record_results(result, filename):
+    """Append results to JSONL file."""
+    with open(filename, "a") as f:
+        f.write(json.dumps(result) + "\n")
+
+
+@pytest.mark.cpu
+def test_moe_amx_int4_benchmark():
+    """Benchmark AMX INT4 MOE performance."""
+    if not HAS_DEPS:
+        pytest.skip(f"Dependencies not available: {import_error}")
+
+    quant_mode = "int4"
+    bytes_per_elem = 0.5
+
+    # Setup output file
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    json_path = os.path.join(script_dir, "bench_moe_amx_int4.jsonl")
+
+    with torch.inference_mode():
+        # Initialize CPUInfer with worker config
+        worker_config = kt_kernel_ext.WorkerPoolConfig()
+        worker_config.subpool_count = worker_config_dict["subpool_count"]
+        worker_config.subpool_numa_map = worker_config_dict["subpool_numa_map"]
+        worker_config.subpool_thread_count = worker_config_dict["subpool_thread_count"]
+        CPUInfer = kt_kernel_ext.CPUInfer(worker_config)
+
+        # Initialize MOE layers
+        moes = []
+        for layer_index in range(layer_num):
+            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()
+            )
+            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.pool = CPUInfer.backend_
+
+            moe = kt_kernel_ext.moe.AMXInt4_MOE(config)
+            CPUInfer.submit(moe.load_weights_task())
+            CPUInfer.sync()
+            moes.append(moe)
+
+        # Generate test data
+        gen_iter = 3000
+        expert_ids = (
+            torch.rand(gen_iter * qlen, expert_num, device="cpu")
+            .argsort(dim=-1)[:, :num_experts_per_tok]
+            .reshape(gen_iter, qlen * num_experts_per_tok)
+            .to("cpu")
+            .contiguous()
+        )
+        weights = (
+            torch.rand((gen_iter, qlen, num_experts_per_tok), dtype=torch.float32, device="cpu").to("cpu").contiguous()
+        )
+        input_tensor = (
+            torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cuda").to("cpu").contiguous()
+        )
+        output_tensor = (
+            torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cuda").to("cpu").contiguous()
+        )
+        bsz_tensor = torch.tensor([qlen], device="cpu")
+
+        # Warm-up iterations
+        print(f"Running warm-up for {warm_up_iter} iterations...")
+        for i in tqdm(range(warm_up_iter), desc="Warm-up"):
+            CPUInfer.submit(
+                moes[i % layer_num].forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids[i % gen_iter].data_ptr(),
+                    weights[i % gen_iter].data_ptr(),
+                    input_tensor[i % layer_num].data_ptr(),
+                    output_tensor[i % layer_num].data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+
+        # Test iterations
+        print(f"Running test for {test_iter} iterations...")
+        start = time.perf_counter()
+        for i in tqdm(range(test_iter), desc="Testing"):
+            CPUInfer.submit(
+                moes[i % layer_num].forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids[i % gen_iter].data_ptr(),
+                    weights[i % gen_iter].data_ptr(),
+                    input_tensor[i % layer_num].data_ptr(),
+                    output_tensor[i % layer_num].data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+        end = time.perf_counter()
+        total_time = end - start
+
+        # Calculate performance metrics
+        time_per_iter_us = total_time / test_iter * 1e6
+        bandwidth = (
+            hidden_size
+            * intermediate_size
+            * 3
+            * num_experts_per_tok
+            * (1 / 8 * 256 * (1 - (31 / 32) ** qlen))
+            * bytes_per_elem
+            * test_iter
+            / total_time
+            / 1e9
+        )  # GB/s
+        flops = (
+            hidden_size * intermediate_size * qlen * 3 * num_experts_per_tok * 2 * test_iter / total_time / 1e12
+        )  # TFLOPS
+
+        print("Quant mode: ", quant_mode)
+        print("Time(s): ", total_time)
+        print("Iteration: ", test_iter)
+        print("Time(us) per iteration: ", time_per_iter_us)
+        print("Bandwidth: ", bandwidth, "GB/s")
+        print("Flops: ", flops, "TFLOPS")
+
+        # Record results
+        result = {
+            "quant_mode": quant_mode,
+            "total_time_seconds": total_time,
+            "iterations": test_iter,
+            "time_per_iteration_us": time_per_iter_us,
+            "bandwidth_GBs": bandwidth,
+            "flops_TFLOPS": flops,
+            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
+            "test_parameters": {
+                "expert_num": expert_num,
+                "hidden_size": hidden_size,
+                "intermediate_size": intermediate_size,
+                "max_len": max_len,
+                "num_experts_per_tok": num_experts_per_tok,
+                "layer_num": layer_num,
+                "qlen": qlen,
+                "warm_up_iter": warm_up_iter,
+                "test_iter": test_iter,
+                "CPUInfer_parameter": CPUINFER_PARAM,
+            },
+        }
+        result.update(get_git_commit())
+        result.update(get_system_info())
+        record_results(result, json_path)
+
+        print(f"Results saved to {json_path}")
+
+
+def run_all_tests():
+    """Run all tests in this file (for standalone execution)."""
+    if not HAS_DEPS:
+        print(f"⚠ Dependencies not available: {import_error}")
+        print("Skipping AMX MOE INT4 benchmark tests")
+        return
+
+    try:
+        print("Running AMX MOE INT4 benchmark test...")
+        test_moe_amx_int4_benchmark()
+        print("✓ AMX MOE INT4 benchmark test passed")
+        print("\n✓ All tests passed!")
+    except Exception as e:
+        print(f"\n✗ Test failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    run_all_tests()

kt-kernel/test/per_commit/test_moe_amx_bench_int4_1.py

@@ -0,0 +1,317 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""AMX MOE INT4 benchmark tests for KT-Kernel.
+
+Benchmarks performance (bandwidth and FLOPS) of AMX-accelerated INT4 MOE operations.
+"""
+
+import os
+import sys
+import time
+import json
+import subprocess
+import platform
+import pytest
+
+# Add parent directory to path for CI registration
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
+# Register this test for CPU CI with estimated runtime of 300 seconds
+register_cpu_ci(est_time=300, suite="default")
+
+# Check if dependencies are available
+try:
+    import torch
+    import kt_kernel  # Import kt_kernel first to register kt_kernel_ext
+
+    kt_kernel_ext = kt_kernel.kt_kernel_ext  # Access the extension module
+    from tqdm import tqdm
+
+    HAS_DEPS = True
+except ImportError as e:
+    HAS_DEPS = False
+    import_error = str(e)
+
+# Test parameters (from original bench_moe_amx.py)
+expert_num = 16
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 25600
+num_experts_per_tok = 8
+layer_num = 2
+qlen = 1024
+warm_up_iter = 1000
+test_iter = 2000
+
+# Worker configuration
+worker_config_dict = {
+    "subpool_count": 2,
+    "subpool_numa_map": [0, 1],
+    "subpool_thread_count": [30, 30],
+}
+CPUINFER_PARAM = 60
+
+
+def get_git_commit():
+    """Get current git commit information."""
+    result = {}
+    try:
+        commit = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("utf-8").strip()
+        commit_msg = subprocess.check_output(["git", "log", "-1", "--pretty=%B"]).decode("utf-8").strip()
+        result["commit"] = commit
+        result["commit_message"] = commit_msg
+
+        dirty_output = subprocess.check_output(["git", "status", "--porcelain"]).decode("utf-8").strip()
+        if dirty_output:
+            result["dirty"] = True
+            result["dirty_files"] = dirty_output.splitlines()
+        else:
+            result["dirty"] = False
+    except Exception as e:
+        result["commit"] = None
+        result["commit_message"] = None
+        result["dirty"] = None
+        result["error"] = str(e)
+    return result
+
+
+def get_system_info():
+    """Get system information including CPU model, memory, cores, and sockets."""
+    info = {}
+    uname = platform.uname()
+    info["system_name"] = uname.system
+    info["node_name"] = uname.node
+
+    # Get CPU model (Linux only)
+    cpu_model = None
+    if os.path.exists("/proc/cpuinfo"):
+        try:
+            with open("/proc/cpuinfo", "r") as f:
+                for line in f:
+                    if "model name" in line:
+                        cpu_model = line.split(":", 1)[1].strip()
+                        break
+        except Exception as e:
+            cpu_model = f"Error: {e}"
+    info["cpu_model"] = cpu_model
+
+    # Get memory size in GB (Linux only)
+    mem_total_gb = None
+    if os.path.exists("/proc/meminfo"):
+        try:
+            with open("/proc/meminfo", "r") as f:
+                for line in f:
+                    if "MemTotal" in line:
+                        mem_kb = float(line.split(":", 1)[1].split()[0])
+                        mem_total_gb = round(mem_kb / (1024 * 1024), 2)
+                        break
+        except Exception as e:
+            mem_total_gb = f"Error: {e}"
+    info["memory_size_GB"] = mem_total_gb
+
+    # Get CPU core count
+    info["cpu_core_count"] = os.cpu_count()
+
+    # Get socket count
+    sockets = set()
+    if os.path.exists("/proc/cpuinfo"):
+        try:
+            with open("/proc/cpuinfo", "r") as f:
+                for line in f:
+                    if "physical id" in line:
+                        sockets.add(line.split(":", 1)[1].strip())
+        except Exception as e:
+            sockets = set()
+    info["cpu_socket_count"] = len(sockets) if len(sockets) > 0 else 1
+
+    return info
+
+
+def record_results(result, filename):
+    """Append results to JSONL file."""
+    with open(filename, "a") as f:
+        f.write(json.dumps(result) + "\n")
+
+
+@pytest.mark.cpu
+def test_moe_amx_int4_1_benchmark():
+    """Benchmark AMX INT4 MOE performance."""
+    if not HAS_DEPS:
+        pytest.skip(f"Dependencies not available: {import_error}")
+
+    quant_mode = "int4"
+    bytes_per_elem = 0.5
+
+    # Setup output file
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    json_path = os.path.join(script_dir, "bench_moe_amx_int4_1.jsonl")
+
+    with torch.inference_mode():
+        # Initialize CPUInfer with worker config
+        worker_config = kt_kernel_ext.WorkerPoolConfig()
+        worker_config.subpool_count = worker_config_dict["subpool_count"]
+        worker_config.subpool_numa_map = worker_config_dict["subpool_numa_map"]
+        worker_config.subpool_thread_count = worker_config_dict["subpool_thread_count"]
+        CPUInfer = kt_kernel_ext.CPUInfer(worker_config)
+
+        # Initialize MOE layers
+        moes = []
+        for layer_index in range(layer_num):
+            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()
+            )
+            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.pool = CPUInfer.backend_
+
+            moe = kt_kernel_ext.moe.AMXInt4_MOE(config)
+            CPUInfer.submit(moe.load_weights_task())
+            CPUInfer.sync()
+            moes.append(moe)
+
+        # Generate test data
+        gen_iter = 3000
+        expert_ids = (
+            torch.rand(gen_iter * qlen, expert_num, device="cpu")
+            .argsort(dim=-1)[:, :num_experts_per_tok]
+            .reshape(gen_iter, qlen * num_experts_per_tok)
+            .to("cpu")
+            .contiguous()
+        )
+        weights = (
+            torch.rand((gen_iter, qlen, num_experts_per_tok), dtype=torch.float32, device="cpu").to("cpu").contiguous()
+        )
+        input_tensor = (
+            torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cuda").to("cpu").contiguous()
+        )
+        output_tensor = (
+            torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cuda").to("cpu").contiguous()
+        )
+        bsz_tensor = torch.tensor([qlen], device="cpu")
+
+        # Warm-up iterations
+        print(f"Running warm-up for {warm_up_iter} iterations...")
+        for i in tqdm(range(warm_up_iter), desc="Warm-up"):
+            CPUInfer.submit(
+                moes[i % layer_num].forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids[i % gen_iter].data_ptr(),
+                    weights[i % gen_iter].data_ptr(),
+                    input_tensor[i % layer_num].data_ptr(),
+                    output_tensor[i % layer_num].data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+
+        # Test iterations
+        print(f"Running test for {test_iter} iterations...")
+        start = time.perf_counter()
+        for i in tqdm(range(test_iter), desc="Testing"):
+            CPUInfer.submit(
+                moes[i % layer_num].forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids[i % gen_iter].data_ptr(),
+                    weights[i % gen_iter].data_ptr(),
+                    input_tensor[i % layer_num].data_ptr(),
+                    output_tensor[i % layer_num].data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+        end = time.perf_counter()
+        total_time = end - start
+
+        # Calculate performance metrics
+        time_per_iter_us = total_time / test_iter * 1e6
+        bandwidth = (
+            hidden_size
+            * intermediate_size
+            * 3
+            * num_experts_per_tok
+            * (1 / 8 * 256 * (1 - (31 / 32) ** qlen))
+            * bytes_per_elem
+            * test_iter
+            / total_time
+            / 1e9
+        )  # GB/s
+        flops = (
+            hidden_size * intermediate_size * qlen * 3 * num_experts_per_tok * 2 * test_iter / total_time / 1e12
+        )  # TFLOPS
+
+        print("Quant mode: ", quant_mode)
+        print("Time(s): ", total_time)
+        print("Iteration: ", test_iter)
+        print("Time(us) per iteration: ", time_per_iter_us)
+        print("Bandwidth: ", bandwidth, "GB/s")
+        print("Flops: ", flops, "TFLOPS")
+
+        # Record results
+        result = {
+            "quant_mode": quant_mode,
+            "total_time_seconds": total_time,
+            "iterations": test_iter,
+            "time_per_iteration_us": time_per_iter_us,
+            "bandwidth_GBs": bandwidth,
+            "flops_TFLOPS": flops,
+            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
+            "test_parameters": {
+                "expert_num": expert_num,
+                "hidden_size": hidden_size,
+                "intermediate_size": intermediate_size,
+                "max_len": max_len,
+                "num_experts_per_tok": num_experts_per_tok,
+                "layer_num": layer_num,
+                "qlen": qlen,
+                "warm_up_iter": warm_up_iter,
+                "test_iter": test_iter,
+                "CPUInfer_parameter": CPUINFER_PARAM,
+            },
+        }
+        result.update(get_git_commit())
+        result.update(get_system_info())
+        record_results(result, json_path)
+
+        print(f"Results saved to {json_path}")
+
+
+def run_all_tests():
+    """Run all tests in this file (for standalone execution)."""
+    if not HAS_DEPS:
+        print(f"Dependencies not available: {import_error}")
+        print("Skipping AMX MOE INT4 benchmark tests")
+        return
+
+    try:
+        print("Running AMX MOE INT4 benchmark test...")
+        test_moe_amx_int4_1_benchmark()
+        print("AMX MOE INT4 benchmark test passed")
+        print("\nAll tests passed!")
+    except Exception as e:
+        print(f"\nTest failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    run_all_tests()

kt-kernel/test/per_commit/test_moe_amx_bench_int4_1k.py

@@ -0,0 +1,329 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""AMX MOE INT4 1K Group benchmark tests for KT-Kernel.
+
+Benchmarks performance (bandwidth and FLOPS) of AMX-accelerated INT4 MOE operations
+with 1K group quantization (AMXInt4_1KGroup_MOE).
+"""
+
+import os
+import sys
+import time
+import json
+import subprocess
+import platform
+import pytest
+
+# Add parent directory to path for CI registration
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
+# Register this test for CPU CI with estimated runtime of 300 seconds
+register_cpu_ci(est_time=300, suite="default")
+
+# Check if dependencies are available
+try:
+    import torch
+    import kt_kernel  # Import kt_kernel first to register kt_kernel_ext
+
+    kt_kernel_ext = kt_kernel.kt_kernel_ext  # Access the extension module
+    from tqdm import tqdm
+
+    HAS_DEPS = True
+except ImportError as e:
+    HAS_DEPS = False
+    import_error = str(e)
+
+# Test parameters (from bench_moe_amx_k.py)
+expert_num = 16
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 25600
+num_experts_per_tok = 8
+layer_num = 2
+qlen = 1024
+warm_up_iter = 1000
+test_iter = 2000
+k_group_size = 128
+
+# Worker configuration
+worker_config_dict = {
+    "subpool_count": 2,
+    "subpool_numa_map": [0, 1],
+    "subpool_thread_count": [30, 30],
+}
+CPUINFER_PARAM = 60
+
+
+def get_git_commit():
+    """Get current git commit information."""
+    result = {}
+    try:
+        commit = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("utf-8").strip()
+        commit_msg = subprocess.check_output(["git", "log", "-1", "--pretty=%B"]).decode("utf-8").strip()
+        result["commit"] = commit
+        result["commit_message"] = commit_msg
+
+        dirty_output = subprocess.check_output(["git", "status", "--porcelain"]).decode("utf-8").strip()
+        if dirty_output:
+            result["dirty"] = True
+            result["dirty_files"] = dirty_output.splitlines()
+        else:
+            result["dirty"] = False
+    except Exception as e:
+        result["commit"] = None
+        result["commit_message"] = None
+        result["dirty"] = None
+        result["error"] = str(e)
+    return result
+
+
+def get_system_info():
+    """Get system information including CPU model, memory, cores, and sockets."""
+    info = {}
+    uname = platform.uname()
+    info["system_name"] = uname.system
+    info["node_name"] = uname.node
+
+    # Get CPU model (Linux only)
+    cpu_model = None
+    if os.path.exists("/proc/cpuinfo"):
+        try:
+            with open("/proc/cpuinfo", "r") as f:
+                for line in f:
+                    if "model name" in line:
+                        cpu_model = line.split(":", 1)[1].strip()
+                        break
+        except Exception as e:
+            cpu_model = f"Error: {e}"
+    info["cpu_model"] = cpu_model
+
+    # Get memory size in GB (Linux only)
+    mem_total_gb = None
+    if os.path.exists("/proc/meminfo"):
+        try:
+            with open("/proc/meminfo", "r") as f:
+                for line in f:
+                    if "MemTotal" in line:
+                        mem_kb = float(line.split(":", 1)[1].split()[0])
+                        mem_total_gb = round(mem_kb / (1024 * 1024), 2)
+                        break
+        except Exception as e:
+            mem_total_gb = f"Error: {e}"
+    info["memory_size_GB"] = mem_total_gb
+
+    # Get CPU core count
+    info["cpu_core_count"] = os.cpu_count()
+
+    # Get socket count
+    sockets = set()
+    if os.path.exists("/proc/cpuinfo"):
+        try:
+            with open("/proc/cpuinfo", "r") as f:
+                for line in f:
+                    if "physical id" in line:
+                        sockets.add(line.split(":", 1)[1].strip())
+        except Exception as e:
+            sockets = set()
+    info["cpu_socket_count"] = len(sockets) if len(sockets) > 0 else 1
+
+    return info
+
+
+def record_results(result, filename):
+    """Append results to JSONL file."""
+    with open(filename, "a") as f:
+        f.write(json.dumps(result) + "\n")
+
+
+@pytest.mark.cpu
+def test_moe_amx_int4_1k_benchmark():
+    """Benchmark AMX INT4 1K Group MOE performance."""
+    if not HAS_DEPS:
+        pytest.skip(f"Dependencies not available: {import_error}")
+
+    quant_mode = "int4_1k"
+    bytes_per_elem = 0.5
+
+    # Setup output file
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    json_path = os.path.join(script_dir, "bench_moe_amx_int4_1k.jsonl")
+
+    with torch.inference_mode():
+        # Initialize CPUInfer with worker config
+        worker_config = kt_kernel_ext.WorkerPoolConfig()
+        worker_config.subpool_count = worker_config_dict["subpool_count"]
+        worker_config.subpool_numa_map = worker_config_dict["subpool_numa_map"]
+        worker_config.subpool_thread_count = worker_config_dict["subpool_thread_count"]
+        CPUInfer = kt_kernel_ext.CPUInfer(worker_config)
+
+        # Physical to logical map for weight loading
+        physical_to_logical_map = torch.tensor(data=range(expert_num), device="cpu", dtype=torch.int64).contiguous()
+
+        # Initialize MOE layers
+        moes = []
+        for layer_index in range(layer_num):
+            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()
+            )
+            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.pool = CPUInfer.backend_
+
+            # Configure quantization for INT4 1K Group
+            config.quant_config.bits = 4
+            config.quant_config.group_size = k_group_size
+            config.quant_config.zero_point = True
+            config.gate_scale = 0
+
+            moe = kt_kernel_ext.moe.AMXInt4_1KGroup_MOE(config)
+            CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+            CPUInfer.sync()
+            moes.append(moe)
+
+        # Generate test data
+        gen_iter = 3000
+        expert_ids = (
+            torch.rand(gen_iter * qlen, expert_num, device="cpu")
+            .argsort(dim=-1)[:, :num_experts_per_tok]
+            .reshape(gen_iter, qlen * num_experts_per_tok)
+            .to("cpu")
+            .contiguous()
+        )
+        weights = (
+            torch.rand((gen_iter, qlen, num_experts_per_tok), dtype=torch.float32, device="cpu").to("cpu").contiguous()
+        )
+        input_tensor = (
+            torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cuda").to("cpu").contiguous()
+        )
+        output_tensor = (
+            torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cuda").to("cpu").contiguous()
+        )
+        bsz_tensor = torch.tensor([qlen], device="cpu")
+
+        # Warm-up iterations
+        print(f"Running warm-up for {warm_up_iter} iterations...")
+        for i in tqdm(range(warm_up_iter), desc="Warm-up"):
+            CPUInfer.submit(
+                moes[i % layer_num].forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids[i % gen_iter].data_ptr(),
+                    weights[i % gen_iter].data_ptr(),
+                    input_tensor[i % layer_num].data_ptr(),
+                    output_tensor[i % layer_num].data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+
+        # Test iterations
+        print(f"Running test for {test_iter} iterations...")
+        start = time.perf_counter()
+        for i in tqdm(range(test_iter), desc="Testing"):
+            CPUInfer.submit(
+                moes[i % layer_num].forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids[i % gen_iter].data_ptr(),
+                    weights[i % gen_iter].data_ptr(),
+                    input_tensor[i % layer_num].data_ptr(),
+                    output_tensor[i % layer_num].data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+        end = time.perf_counter()
+        total_time = end - start
+
+        # Calculate performance metrics
+        time_per_iter_us = total_time / test_iter * 1e6
+        bandwidth = (
+            hidden_size
+            * intermediate_size
+            * 3
+            * num_experts_per_tok
+            * (1 / 8 * 256 * (1 - (31 / 32) ** qlen))
+            * bytes_per_elem
+            * test_iter
+            / total_time
+            / 1e9
+        )  # GB/s
+        flops = (
+            hidden_size * intermediate_size * qlen * 3 * num_experts_per_tok * 2 * test_iter / total_time / 1e12
+        )  # TFLOPS
+
+        print("Quant mode: ", quant_mode)
+        print("Time(s): ", total_time)
+        print("Iteration: ", test_iter)
+        print("Time(us) per iteration: ", time_per_iter_us)
+        print("Bandwidth: ", bandwidth, "GB/s")
+        print("Flops: ", flops, "TFLOPS")
+
+        # Record results
+        result = {
+            "quant_mode": quant_mode,
+            "total_time_seconds": total_time,
+            "iterations": test_iter,
+            "time_per_iteration_us": time_per_iter_us,
+            "bandwidth_GBs": bandwidth,
+            "flops_TFLOPS": flops,
+            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
+            "test_parameters": {
+                "expert_num": expert_num,
+                "hidden_size": hidden_size,
+                "intermediate_size": intermediate_size,
+                "max_len": max_len,
+                "num_experts_per_tok": num_experts_per_tok,
+                "layer_num": layer_num,
+                "qlen": qlen,
+                "warm_up_iter": warm_up_iter,
+                "test_iter": test_iter,
+                "CPUInfer_parameter": CPUINFER_PARAM,
+                "k_group_size": k_group_size,
+            },
+        }
+        result.update(get_git_commit())
+        result.update(get_system_info())
+        record_results(result, json_path)
+
+        print(f"Results saved to {json_path}")
+
+
+def run_all_tests():
+    """Run all tests in this file (for standalone execution)."""
+    if not HAS_DEPS:
+        print(f"Dependencies not available: {import_error}")
+        print("Skipping AMX MOE INT4 1K Group benchmark tests")
+        return
+
+    try:
+        print("Running AMX MOE INT4 1K Group benchmark test...")
+        test_moe_amx_int4_1k_benchmark()
+        print("AMX MOE INT4 1K Group benchmark test passed")
+        print("\nAll tests passed!")
+    except Exception as e:
+        print(f"\nTest failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    run_all_tests()

kt-kernel/test/per_commit/test_moe_amx_bench_int8.py

@@ -0,0 +1,317 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""AMX MOE INT8 benchmark tests for KT-Kernel.
+
+Benchmarks performance (bandwidth and FLOPS) of AMX-accelerated INT8 MOE operations.
+"""
+
+import os
+import sys
+import time
+import json
+import subprocess
+import platform
+import pytest
+
+# Add parent directory to path for CI registration
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from ci.ci_register import register_cpu_ci
+
+# Register this test for CPU CI with estimated runtime of 300 seconds
+register_cpu_ci(est_time=300, suite="default")
+
+# Check if dependencies are available
+try:
+    import torch
+    import kt_kernel  # Import kt_kernel first to register kt_kernel_ext
+
+    kt_kernel_ext = kt_kernel.kt_kernel_ext  # Access the extension module
+    from tqdm import tqdm
+
+    HAS_DEPS = True
+except ImportError as e:
+    HAS_DEPS = False
+    import_error = str(e)
+
+# Test parameters (from original bench_moe_amx.py)
+expert_num = 128
+hidden_size = 7168
+intermediate_size = 2048
+max_len = 25600
+num_experts_per_tok = 0
+layer_num = 2
+qlen = 1
+warm_up_iter = 1000
+test_iter = 2000
+
+# Worker configuration
+worker_config_dict = {
+    "subpool_count": 2,
+    "subpool_numa_map": [0, 1],
+    "subpool_thread_count": [30, 30],
+}
+CPUINFER_PARAM = 60
+
+
+def get_git_commit():
+    """Get current git commit information."""
+    result = {}
+    try:
+        commit = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("utf-8").strip()
+        commit_msg = subprocess.check_output(["git", "log", "-1", "--pretty=%B"]).decode("utf-8").strip()
+        result["commit"] = commit
+        result["commit_message"] = commit_msg
+
+        dirty_output = subprocess.check_output(["git", "status", "--porcelain"]).decode("utf-8").strip()
+        if dirty_output:
+            result["dirty"] = True
+            result["dirty_files"] = dirty_output.splitlines()
+        else:
+            result["dirty"] = False
+    except Exception as e:
+        result["commit"] = None
+        result["commit_message"] = None
+        result["dirty"] = None
+        result["error"] = str(e)
+    return result
+
+
+def get_system_info():
+    """Get system information including CPU model, memory, cores, and sockets."""
+    info = {}
+    uname = platform.uname()
+    info["system_name"] = uname.system
+    info["node_name"] = uname.node
+
+    # Get CPU model (Linux only)
+    cpu_model = None
+    if os.path.exists("/proc/cpuinfo"):
+        try:
+            with open("/proc/cpuinfo", "r") as f:
+                for line in f:
+                    if "model name" in line:
+                        cpu_model = line.split(":", 1)[1].strip()
+                        break
+        except Exception as e:
+            cpu_model = f"Error: {e}"
+    info["cpu_model"] = cpu_model
+
+    # Get memory size in GB (Linux only)
+    mem_total_gb = None
+    if os.path.exists("/proc/meminfo"):
+        try:
+            with open("/proc/meminfo", "r") as f:
+                for line in f:
+                    if "MemTotal" in line:
+                        mem_kb = float(line.split(":", 1)[1].split()[0])
+                        mem_total_gb = round(mem_kb / (1024 * 1024), 2)
+                        break
+        except Exception as e:
+            mem_total_gb = f"Error: {e}"
+    info["memory_size_GB"] = mem_total_gb
+
+    # Get CPU core count
+    info["cpu_core_count"] = os.cpu_count()
+
+    # Get socket count
+    sockets = set()
+    if os.path.exists("/proc/cpuinfo"):
+        try:
+            with open("/proc/cpuinfo", "r") as f:
+                for line in f:
+                    if "physical id" in line:
+                        sockets.add(line.split(":", 1)[1].strip())
+        except Exception as e:
+            sockets = set()
+    info["cpu_socket_count"] = len(sockets) if len(sockets) > 0 else 1
+
+    return info
+
+
+def record_results(result, filename):
+    """Append results to JSONL file."""
+    with open(filename, "a") as f:
+        f.write(json.dumps(result) + "\n")
+
+
+@pytest.mark.cpu
+def test_moe_amx_int8_benchmark():
+    """Benchmark AMX INT8 MOE performance."""
+    if not HAS_DEPS:
+        pytest.skip(f"Dependencies not available: {import_error}")
+
+    quant_mode = "int8"
+    bytes_per_elem = 1.0
+
+    # Setup output file
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    json_path = os.path.join(script_dir, "bench_moe_amx_int8.jsonl")
+
+    with torch.inference_mode():
+        # Initialize CPUInfer with worker config
+        worker_config = kt_kernel_ext.WorkerPoolConfig()
+        worker_config.subpool_count = worker_config_dict["subpool_count"]
+        worker_config.subpool_numa_map = worker_config_dict["subpool_numa_map"]
+        worker_config.subpool_thread_count = worker_config_dict["subpool_thread_count"]
+        CPUInfer = kt_kernel_ext.CPUInfer(worker_config)
+
+        # Initialize MOE layers
+        moes = []
+        for layer_index in range(layer_num):
+            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()
+            )
+            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.pool = CPUInfer.backend_
+
+            moe = kt_kernel_ext.moe.AMXInt8_MOE(config)
+            CPUInfer.submit(moe.load_weights_task())
+            CPUInfer.sync()
+            moes.append(moe)
+
+        # Generate test data
+        gen_iter = 3000
+        expert_ids = (
+            torch.rand(gen_iter * qlen, expert_num, device="cpu")
+            .argsort(dim=-1)[:, :num_experts_per_tok]
+            .reshape(gen_iter, qlen * num_experts_per_tok)
+            .to("cpu")
+            .contiguous()
+        )
+        weights = (
+            torch.rand((gen_iter, qlen, num_experts_per_tok), dtype=torch.float32, device="cpu").to("cpu").contiguous()
+        )
+        input_tensor = (
+            torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cuda").to("cpu").contiguous()
+        )
+        output_tensor = (
+            torch.empty((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cuda").to("cpu").contiguous()
+        )
+        bsz_tensor = torch.tensor([qlen], device="cpu")
+
+        # Warm-up iterations
+        print(f"Running warm-up for {warm_up_iter} iterations...")
+        for i in tqdm(range(warm_up_iter), desc="Warm-up"):
+            CPUInfer.submit(
+                moes[i % layer_num].forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids[i % gen_iter].data_ptr(),
+                    weights[i % gen_iter].data_ptr(),
+                    input_tensor[i % layer_num].data_ptr(),
+                    output_tensor[i % layer_num].data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+
+        # Test iterations
+        print(f"Running test for {test_iter} iterations...")
+        start = time.perf_counter()
+        for i in tqdm(range(test_iter), desc="Testing"):
+            CPUInfer.submit(
+                moes[i % layer_num].forward_task(
+                    bsz_tensor.data_ptr(),
+                    num_experts_per_tok,
+                    expert_ids[i % gen_iter].data_ptr(),
+                    weights[i % gen_iter].data_ptr(),
+                    input_tensor[i % layer_num].data_ptr(),
+                    output_tensor[i % layer_num].data_ptr(),
+                    False,
+                )
+            )
+            CPUInfer.sync()
+        end = time.perf_counter()
+        total_time = end - start
+
+        # Calculate performance metrics
+        time_per_iter_us = total_time / test_iter * 1e6
+        bandwidth = (
+            hidden_size
+            * intermediate_size
+            * 3
+            * num_experts_per_tok
+            * (1 / 8 * 256 * (1 - (31 / 32) ** qlen))
+            * bytes_per_elem
+            * test_iter
+            / total_time
+            / 1e9
+        )  # GB/s
+        flops = (
+            hidden_size * intermediate_size * qlen * 3 * num_experts_per_tok * 2 * test_iter / total_time / 1e12
+        )  # TFLOPS
+
+        print("Quant mode: ", quant_mode)
+        print("Time(s): ", total_time)
+        print("Iteration: ", test_iter)
+        print("Time(us) per iteration: ", time_per_iter_us)
+        print("Bandwidth: ", bandwidth, "GB/s")
+        print("Flops: ", flops, "TFLOPS")
+
+        # Record results
+        result = {
+            "quant_mode": quant_mode,
+            "total_time_seconds": total_time,
+            "iterations": test_iter,
+            "time_per_iteration_us": time_per_iter_us,
+            "bandwidth_GBs": bandwidth,
+            "flops_TFLOPS": flops,
+            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
+            "test_parameters": {
+                "expert_num": expert_num,
+                "hidden_size": hidden_size,
+                "intermediate_size": intermediate_size,
+                "max_len": max_len,
+                "num_experts_per_tok": num_experts_per_tok,
+                "layer_num": layer_num,
+                "qlen": qlen,
+                "warm_up_iter": warm_up_iter,
+                "test_iter": test_iter,
+                "CPUInfer_parameter": CPUINFER_PARAM,
+            },
+        }
+        result.update(get_git_commit())
+        result.update(get_system_info())
+        record_results(result, json_path)
+
+        print(f"Results saved to {json_path}")
+
+
+def run_all_tests():
+    """Run all tests in this file (for standalone execution)."""
+    if not HAS_DEPS:
+        print(f"⚠ Dependencies not available: {import_error}")
+        print("Skipping AMX MOE INT8 benchmark tests")
+        return
+
+    try:
+        print("Running AMX MOE INT8 benchmark test...")
+        test_moe_amx_int8_benchmark()
+        print("✓ AMX MOE INT8 benchmark test passed")
+        print("\n✓ All tests passed!")
+    except Exception as e:
+        print(f"\n✗ Test failed: {e}")
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    run_all_tests()

kt-kernel/test/per_commit/test_moe_avx2_accuracy_bf16.py

@@ -0,0 +1,165 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""AVX2 BF16 MoE accuracy tests for KT-Kernel.
+
+Tests accuracy of AVX2 BF16 MOE operations against torch reference.
+"""
+
+import os
+import sys
+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
+hidden_size = 256
+intermediate_size = 512
+num_experts_per_tok = 2
+max_len = 128
+validation_iter = 3
+CPUINFER_PARAM = 60
+
+
+def act_fn(x):
+    """SiLU activation."""
+    return x / (1.0 + torch.exp(-x))
+
+
+def mlp_torch(input, gate_proj, up_proj, down_proj):
+    """PyTorch reference MLP."""
+    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):
+    """PyTorch reference MoE."""
+    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
+
+
+@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()
+    CPUInfer = kt_kernel_ext.CPUInfer(CPUINFER_PARAM)
+
+    with torch.inference_mode():
+        # Generate BF16 weights
+        gate_proj = (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 10.0).to(torch.bfloat16).contiguous()
+        up_proj = (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 10.0).to(torch.bfloat16).contiguous()
+        down_proj = (torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32) / 10.0).to(torch.bfloat16).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.up_scale = 0
+        config.down_scale = 0
+        config.pool = CPUInfer.backend_
+
+        # Create AVX2 BF16 MOE
+        moe = kt_kernel_ext.moe.AVX2BF16_MOE(config)
+        CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+        CPUInfer.sync()
+
+        print(f"\n--- {label} (qlen={qlen}) ---")
+        for i in range(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_data = (torch.randn((qlen, hidden_size), dtype=torch.float32) / 100.0).to(torch.bfloat16).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+
+            bsz_tensor = torch.tensor([qlen], dtype=torch.int32)
+
+            # Run AVX2 BF16 MOE
+            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()
+
+            # Run torch reference (in float32 for accuracy)
+            t_output = moe_torch(
+                input_data.float(), expert_ids, weights,
+                gate_proj.float(), up_proj.float(), down_proj.float()
+            ).to(torch.bfloat16)
+
+            # Calculate relative difference
+            diff = torch.mean(torch.abs(output.float() - t_output.float())) / (torch.mean(torch.abs(t_output.float())) + 1e-8)
+            print(f"  Iteration {i}: diff = {diff:.6f}")
+
+            # BF16 should be very accurate (< 0.01)
+            assert diff < 0.02, f"AVX2 BF16 accuracy test failed: diff={diff:.6f} >= 0.02"
+
+    print(f"  PASSED")
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("AVX2 BF16 MoE Accuracy Test")
+    print("=" * 60)
+
+    try:
+        # Test decode path (qlen=1)
+        test_avx2_bf16_accuracy(qlen=1, label="Decode")
+
+        # Test prefill path (qlen=16)
+        test_avx2_bf16_accuracy(qlen=16, label="Prefill")
+
+        print("\n" + "=" * 60)
+        print("ALL TESTS PASSED")
+        print("=" * 60)
+    except Exception as e:
+        print(f"\nTEST FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)

kt-kernel/test/per_commit/test_moe_avx2_accuracy_fp8.py

@@ -0,0 +1,258 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""AVX2 FP8 MoE accuracy tests for KT-Kernel."""
+
+import os
+import sys
+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
+intermediate_size = 512
+num_experts_per_tok = 2
+max_len = 128
+group_size = 128
+validation_iter = 3
+CPUINFER_PARAM = 60
+
+
+def fp8_e4m3_quantize(tensor_bf16):
+    """Quantize BF16 tensor to FP8 E4M3 with block-wise scales (128x128)."""
+    n, k = tensor_bf16.shape
+    tensor_fp32 = tensor_bf16.float()
+
+    n_blocks_n = (n + group_size - 1) // group_size
+    n_blocks_k = (k + group_size - 1) // group_size
+
+    fp8_data = torch.zeros(n, k, dtype=torch.uint8)
+    scales = torch.zeros(n_blocks_n, n_blocks_k, dtype=torch.float32)
+
+    # FP8 E4M3 max value: 2^8 * (1 + 7/8) = 448
+    fp8_max = 448.0
+
+    for bn in range(n_blocks_n):
+        for bk in range(n_blocks_k):
+            n_start = bn * group_size
+            n_end = min(n_start + group_size, n)
+            k_start = bk * group_size
+            k_end = min(k_start + group_size, k)
+
+            block = tensor_fp32[n_start:n_end, k_start:k_end]
+            amax = block.abs().max().item()
+            if amax == 0:
+                scale = 1.0
+            else:
+                scale = amax / fp8_max
+            scales[bn, bk] = scale
+
+            # Quantize
+            for i in range(n_end - n_start):
+                for j in range(k_end - k_start):
+                    val = block[i, j].item() / scale
+                    fp8_data[n_start + i, k_start + j] = float_to_fp8_e4m3(val)
+
+    return fp8_data, scales
+
+
+def float_to_fp8_e4m3(val):
+    """Convert float to FP8 E4M3."""
+    if math.isnan(val):
+        return 0x7F
+    sign = 1 if val < 0 else 0
+    val = abs(val)
+    if val == 0:
+        return sign << 7
+    # Clamp to max
+    if val >= 448.0:
+        return (sign << 7) | 0x7E  # max finite
+    # Find exponent
+    exp = int(math.floor(math.log2(val))) + 7
+    if exp <= 0:
+        # Subnormal
+        man = int(round(val * (2**6) * 8))
+        man = min(man, 7)
+        return (sign << 7) | man
+    if exp >= 15:
+        return (sign << 7) | 0x7E  # clamp to max
+    # Normal
+    man = int(round((val / (2**(exp-7)) - 1.0) * 8))
+    man = min(man, 7)
+    return (sign << 7) | (exp << 3) | man
+
+
+def fp8_e4m3_to_float(byte_val):
+    """Convert FP8 E4M3 byte to float."""
+    sign = (byte_val >> 7) & 1
+    exp = (byte_val >> 3) & 0xF
+    man = byte_val & 0x7
+    if exp == 0 and man == 0:
+        return 0.0
+    if exp == 0:
+        val = (2**-6) * (man / 8.0)
+    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
+
+
+def fp8_dequantize(fp8_data, scales):
+    """Dequantize FP8 + scales back to float32."""
+    n, k = fp8_data.shape
+    result = torch.zeros(n, k, dtype=torch.float32)
+    n_blocks_n = scales.shape[0]
+    n_blocks_k = scales.shape[1]
+
+    for i in range(n):
+        for j in range(k):
+            bn = i // group_size
+            bk = j // group_size
+            scale = scales[bn, bk].item()
+            fp8_val = fp8_e4m3_to_float(fp8_data[i, j].item())
+            result[i, j] = fp8_val * scale
+    return result
+
+
+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
+        tokens = sorted_tokens[start_idx:end_idx]
+        out = mlp_torch(tokens, gate_proj[i], up_proj[i], down_proj[i])
+        outputs.append(out)
+        start_idx = end_idx
+    outs = torch.cat(outputs, dim=0) if 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).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)
+
+    with torch.inference_mode():
+        # Generate BF16 weights, quantize to FP8
+        gate_bf16 = (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 10.0).to(torch.bfloat16)
+        up_bf16 = (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 10.0).to(torch.bfloat16)
+        down_bf16 = (torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32) / 10.0).to(torch.bfloat16)
+
+        # Quantize each expert
+        gate_fp8_list, gate_scale_list = [], []
+        up_fp8_list, up_scale_list = [], []
+        down_fp8_list, down_scale_list = [], []
+
+        for e in range(expert_num):
+            gf, gs = fp8_e4m3_quantize(gate_bf16[e])
+            gate_fp8_list.append(gf)
+            gate_scale_list.append(gs)
+            uf, us = fp8_e4m3_quantize(up_bf16[e])
+            up_fp8_list.append(uf)
+            up_scale_list.append(us)
+            df, ds = fp8_e4m3_quantize(down_bf16[e])
+            down_fp8_list.append(df)
+            down_scale_list.append(ds)
+
+        # Stack into contiguous tensors
+        gate_fp8 = torch.stack(gate_fp8_list).contiguous()
+        gate_scales = torch.stack(gate_scale_list).contiguous()
+        up_fp8 = torch.stack(up_fp8_list).contiguous()
+        up_scales = torch.stack(up_scale_list).contiguous()
+        down_fp8 = torch.stack(down_fp8_list).contiguous()
+        down_scales = torch.stack(down_scale_list).contiguous()
+
+        # Dequantize for reference computation
+        gate_deq = torch.stack([fp8_dequantize(gate_fp8_list[e], gate_scale_list[e]) for e in range(expert_num)])
+        up_deq = torch.stack([fp8_dequantize(up_fp8_list[e], up_scale_list[e]) for e in range(expert_num)])
+        down_deq = torch.stack([fp8_dequantize(down_fp8_list[e], down_scale_list[e]) for e in range(expert_num)])
+
+        # 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_fp8.data_ptr()
+        config.up_proj = up_fp8.data_ptr()
+        config.down_proj = down_fp8.data_ptr()
+        config.gate_scale = gate_scales.data_ptr()
+        config.up_scale = up_scales.data_ptr()
+        config.down_scale = down_scales.data_ptr()
+        config.quant_config.bits = 8
+        config.quant_config.group_size = group_size
+        config.quant_config.zero_point = False
+        config.pool = CPUInfer.backend_
+
+        moe = kt_kernel_ext.moe.AVX2FP8_MOE(config)
+        CPUInfer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+        CPUInfer.sync()
+
+        print("\n--- %s (qlen=%d) ---" % (label, qlen))
+        for i in range(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_data = (torch.randn((qlen, hidden_size), dtype=torch.float32) / 100.0).to(torch.bfloat16).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+
+            bsz_tensor = torch.tensor([qlen], dtype=torch.int32)
+            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()
+
+            # Reference: use dequantized FP32 weights
+            t_output = moe_torch(input_data.float(), expert_ids, weights, gate_deq, up_deq, down_deq).to(torch.bfloat16)
+
+            diff = torch.mean(torch.abs(output.float() - t_output.float())) / (torch.mean(torch.abs(t_output.float())) + 1e-8)
+            print("  Iteration %d: diff = %.6f" % (i, diff.item()))
+            assert diff < 0.1, "FP8 accuracy test failed: diff=%.6f >= 0.1" % diff.item()
+
+    print("  PASSED")
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("AVX2 FP8 MoE Accuracy Test")
+    print("=" * 60)
+    try:
+        test_avx2_fp8_accuracy(qlen=1, label="Decode")
+        test_avx2_fp8_accuracy(qlen=16, label="Prefill")
+        print("\n" + "=" * 60)
+        print("ALL TESTS PASSED")
+        print("=" * 60)
+    except Exception as e:
+        print("\nTEST FAILED: %s" % e)
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)

kt-kernel/test/per_commit/test_moe_gptq_int4_accuracy.py

@@ -0,0 +1,317 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""GPTQ INT4 MoE accuracy tests for KT-Kernel x86 backends."""
+
+import importlib.util
+import os
+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
+
+KT_KERNEL_ROOT = Path(__file__).resolve().parents[2]
+
+expert_num = 8
+hidden_size = 256
+intermediate_size = 512
+num_experts_per_tok = 2
+max_len = 128
+group_size = 128
+validation_iter = 3
+CPUINFER_PARAM = 16
+
+
+def load_amx_utils():
+    pkg_root = KT_KERNEL_ROOT / "python"
+    utils_root = pkg_root / "utils"
+
+    if "kt_kernel" not in sys.modules:
+        kt_kernel_pkg = types.ModuleType("kt_kernel")
+        kt_kernel_pkg.__path__ = [str(pkg_root)]
+        kt_kernel_pkg.kt_kernel_ext = kt_kernel_ext
+        sys.modules["kt_kernel"] = kt_kernel_pkg
+
+    if "kt_kernel_ext" not in sys.modules:
+        sys.modules["kt_kernel_ext"] = kt_kernel_ext
+
+    if "kt_kernel.utils" not in sys.modules:
+        utils_pkg = types.ModuleType("kt_kernel.utils")
+        utils_pkg.__path__ = [str(utils_root)]
+        sys.modules["kt_kernel.utils"] = utils_pkg
+
+    module_specs = [
+        ("kt_kernel.experts_base", pkg_root / "experts_base.py"),
+        ("kt_kernel.utils.loader", utils_root / "loader.py"),
+        ("kt_kernel.utils.amx", utils_root / "amx.py"),
+    ]
+    for module_name, module_path in module_specs:
+        if module_name in sys.modules:
+            continue
+        spec = importlib.util.spec_from_file_location(module_name, module_path)
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+        assert spec.loader is not None
+        spec.loader.exec_module(module)
+
+    return sys.modules["kt_kernel.utils.amx"]
+
+
+def gptq_sym_int4_quantize(weight_bf16):
+    """Quantize [N, K] BF16 weight to GPTQ symmetric int4 layout."""
+    n, k = weight_bf16.shape
+    assert k % 8 == 0
+    assert k % group_size == 0
+
+    weight_fp32 = weight_bf16.float()
+    qweight = torch.zeros((k // 8, n), dtype=torch.int32)
+    scales = torch.zeros((k // group_size, n), dtype=torch.float32)
+
+    for ni in range(n):
+        for g in range(k // group_size):
+            k_start = g * group_size
+            k_end = k_start + group_size
+            block = weight_fp32[ni, k_start:k_end]
+            amax = block.abs().max().item()
+            scale = amax / 7.0 if amax > 0 else 1.0
+            scales[g, ni] = scale
+
+            for kk in range(k_start, k_end, 8):
+                packed = 0
+                for nib in range(8):
+                    q = int(round(weight_fp32[ni, kk + nib].item() / scale)) + 8
+                    q = max(0, min(15, q))
+                    packed |= q << (nib * 4)
+                if packed >= 2**31:
+                    packed -= 2**32
+                qweight[kk // 8, ni] = packed
+
+    return qweight, scales
+
+
+def gptq_sym_int4_dequantize(qweight, scales, out_features, in_features):
+    """Dequantize GPTQ qweight/scales back to fp32 [N, K]."""
+    result = torch.zeros((out_features, in_features), dtype=torch.float32)
+    for ni in range(out_features):
+        for g in range(in_features // group_size):
+            scale = scales[g, ni].item()
+            k_start = g * group_size
+            k_end = k_start + group_size
+            for kk in range(k_start, k_end, 8):
+                packed = int(qweight[kk // 8, ni].item())
+                for nib in range(8):
+                    result[ni, kk + nib] = (((packed >> (nib * 4)) & 0xF) - 8) * scale
+    return result
+
+
+def act_fn(x):
+    return x / (1.0 + torch.exp(-x))
+
+
+def mlp_torch(input_data, gate_proj, up_proj, down_proj):
+    gate_buf = torch.mm(input_data, gate_proj.t())
+    up_buf = torch.mm(input_data, up_proj.t())
+    intermediate = act_fn(gate_buf) * up_buf
+    return torch.mm(intermediate, down_proj.t())
+
+
+def moe_torch(input_data, 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_data[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 = sorted_tokens[start_idx:end_idx]
+        out = mlp_torch(tokens, gate_proj[i], up_proj[i], down_proj[i])
+        outputs.append(out)
+        start_idx = end_idx
+    outs = torch.cat(outputs, dim=0) if 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).float().mul_(weights.unsqueeze(-1)).sum(1)).to(new_x.dtype)
+
+
+def available_backends():
+    backends = []
+    if hasattr(kt_kernel_ext.moe, "AVX2GPTQInt4_MOE"):
+        backends.append(("AVX2GPTQInt4_MOE", kt_kernel_ext.moe.AVX2GPTQInt4_MOE, 0.12))
+
+    if hasattr(kt_kernel_ext.moe, "AVXVNNI256GPTQInt4_MOE"):
+        has_avx_vnni = False
+        try:
+            with open("/proc/cpuinfo", "r") as f:
+                has_avx_vnni = any(("avx_vnni" in line or "avxvnni" in line) for line in f if line.startswith("flags"))
+        except OSError:
+            has_avx_vnni = False
+        if has_avx_vnni:
+            backends.append(("AVXVNNI256GPTQInt4_MOE", kt_kernel_ext.moe.AVXVNNI256GPTQInt4_MOE, 0.20))
+    return backends
+
+
+def run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen):
+    physical_to_logical_map = torch.tensor(range(expert_num), dtype=torch.int64).contiguous()
+    cpu_infer = kt_kernel_ext.CPUInfer(CPUINFER_PARAM)
+
+    with torch.inference_mode():
+        gate_bf16 = (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 10.0).to(
+            torch.bfloat16
+        )
+        up_bf16 = (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 10.0).to(
+            torch.bfloat16
+        )
+        down_bf16 = (torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32) / 10.0).to(
+            torch.bfloat16
+        )
+
+        gate_qw_list, gate_scale_list = [], []
+        up_qw_list, up_scale_list = [], []
+        down_qw_list, down_scale_list = [], []
+
+        for e in range(expert_num):
+            qw, sc = gptq_sym_int4_quantize(gate_bf16[e])
+            gate_qw_list.append(qw)
+            gate_scale_list.append(sc)
+
+            qw, sc = gptq_sym_int4_quantize(up_bf16[e])
+            up_qw_list.append(qw)
+            up_scale_list.append(sc)
+
+            qw, sc = gptq_sym_int4_quantize(down_bf16[e])
+            down_qw_list.append(qw)
+            down_scale_list.append(sc)
+
+        gate_qw = torch.stack(gate_qw_list).contiguous()
+        gate_scales = torch.stack(gate_scale_list).contiguous()
+        up_qw = torch.stack(up_qw_list).contiguous()
+        up_scales = torch.stack(up_scale_list).contiguous()
+        down_qw = torch.stack(down_qw_list).contiguous()
+        down_scales = torch.stack(down_scale_list).contiguous()
+
+        gate_deq = torch.stack(
+            [
+                gptq_sym_int4_dequantize(gate_qw_list[e], gate_scale_list[e], intermediate_size, hidden_size)
+                for e in range(expert_num)
+            ]
+        )
+        up_deq = torch.stack(
+            [
+                gptq_sym_int4_dequantize(up_qw_list[e], up_scale_list[e], intermediate_size, hidden_size)
+                for e in range(expert_num)
+            ]
+        )
+        down_deq = torch.stack(
+            [
+                gptq_sym_int4_dequantize(down_qw_list[e], down_scale_list[e], hidden_size, intermediate_size)
+                for e in range(expert_num)
+            ]
+        )
+
+        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_qw.data_ptr()
+        config.up_proj = up_qw.data_ptr()
+        config.down_proj = down_qw.data_ptr()
+        config.gate_scale = gate_scales.data_ptr()
+        config.up_scale = up_scales.data_ptr()
+        config.down_scale = down_scales.data_ptr()
+        config.quant_config.bits = 4
+        config.quant_config.group_size = group_size
+        config.quant_config.zero_point = False
+        config.pool = cpu_infer.backend_
+
+        moe = backend_cls(config)
+        cpu_infer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+        cpu_infer.sync()
+
+        print(f"\n--- {backend_name} (qlen={qlen}) ---")
+        for i in range(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_data = (torch.randn((qlen, hidden_size), dtype=torch.float32) / 100.0).to(torch.bfloat16).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+
+            bsz_tensor = torch.tensor([qlen], dtype=torch.int32)
+            cpu_infer.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,
+                )
+            )
+            cpu_infer.sync()
+
+            ref_output = moe_torch(input_data.float(), expert_ids, weights, gate_deq, up_deq, down_deq).to(
+                torch.bfloat16
+            )
+            diff = torch.mean(torch.abs(output.float() - ref_output.float())) / (
+                torch.mean(torch.abs(ref_output.float())) + 1e-8
+            )
+            print(f"  Iteration {i}: diff = {diff.item():.6f}")
+            assert diff < threshold, f"{backend_name} accuracy test failed: diff={diff.item():.6f} >= {threshold}"
+
+
+def test_gptq_int4_accuracy():
+    backends = available_backends()
+    if not backends:
+        print("Skipping GPTQ INT4 accuracy tests: no x86 GPTQ backend available")
+        return
+
+    for backend_name, backend_cls, threshold in backends:
+        run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen=1)
+        run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen=16)
+
+
+def test_gptq_int4_backend_selection_falls_back_to_avx2_for_large_group_size(monkeypatch):
+    amx_utils = load_amx_utils()
+    fake_avx2_backend = object()
+    fake_avxvnni_backend = object()
+
+    monkeypatch.setattr(amx_utils, "AVX2GPTQInt4_MOE", fake_avx2_backend)
+    monkeypatch.setattr(amx_utils, "AVXVNNI256GPTQInt4_MOE", fake_avxvnni_backend)
+    monkeypatch.setattr(amx_utils, "_HAS_AVX2_GPTQ_INT4_SUPPORT", True)
+    monkeypatch.setattr(amx_utils, "_HAS_AVXVNNI256_GPTQ_INT4_SUPPORT", True)
+    monkeypatch.setattr(amx_utils, "_HOST_HAS_AVX_VNNI", True)
+    monkeypatch.delenv("KT_GPTQ_INT4_BACKEND", raising=False)
+
+    assert amx_utils._select_gptq_int4_backend(512) is fake_avx2_backend
+    assert amx_utils._select_gptq_int4_backend(128) is fake_avxvnni_backend
+
+
+def test_gptq_int4_backend_selection_rejects_forced_avxvnni_with_large_group_size(monkeypatch):
+    amx_utils = load_amx_utils()
+
+    monkeypatch.setattr(amx_utils, "_HAS_AVXVNNI256_GPTQ_INT4_SUPPORT", True)
+    monkeypatch.setattr(amx_utils, "_HOST_HAS_AVX_VNNI", True)
+    monkeypatch.setenv("KT_GPTQ_INT4_BACKEND", "avxvnni")
+
+    with pytest.raises(RuntimeError, match="group_size=512 is unsupported"):
+        amx_utils._select_gptq_int4_backend(512)
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("GPTQ INT4 MoE Accuracy Test")
+    print("=" * 60)
+    test_gptq_int4_accuracy()
+    print("PASSED")

kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py

@@ -0,0 +1,395 @@
+#!/usr/bin/env python
+# coding=utf-8
+"""RAWINT4 MoE accuracy tests for KT-Kernel x86 backends."""
+
+import importlib.util
+import os
+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
+
+KT_KERNEL_ROOT = Path(__file__).resolve().parents[2]
+
+expert_num = 8
+hidden_size = 256
+intermediate_size = 512
+num_experts_per_tok = 2
+max_len = 128
+group_size = 128
+validation_iter = 3
+CPUINFER_PARAM = 16
+
+
+def load_amx_utils():
+    pkg_root = KT_KERNEL_ROOT / "python"
+    utils_root = pkg_root / "utils"
+
+    if "kt_kernel" not in sys.modules:
+        kt_kernel_pkg = types.ModuleType("kt_kernel")
+        kt_kernel_pkg.__path__ = [str(pkg_root)]
+        kt_kernel_pkg.kt_kernel_ext = kt_kernel_ext
+        sys.modules["kt_kernel"] = kt_kernel_pkg
+
+    if "kt_kernel_ext" not in sys.modules:
+        sys.modules["kt_kernel_ext"] = kt_kernel_ext
+
+    if "kt_kernel.utils" not in sys.modules:
+        utils_pkg = types.ModuleType("kt_kernel.utils")
+        utils_pkg.__path__ = [str(utils_root)]
+        sys.modules["kt_kernel.utils"] = utils_pkg
+
+    module_specs = [
+        ("kt_kernel.experts_base", pkg_root / "experts_base.py"),
+        ("kt_kernel.utils.loader", utils_root / "loader.py"),
+        ("kt_kernel.utils.amx", utils_root / "amx.py"),
+    ]
+    for module_name, module_path in module_specs:
+        if module_name in sys.modules:
+            continue
+        spec = importlib.util.spec_from_file_location(module_name, module_path)
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+        assert spec.loader is not None
+        spec.loader.exec_module(module)
+
+    return sys.modules["kt_kernel.utils.amx"]
+
+
+def rawint4_quantize(weight_bf16):
+    """Quantize [N, K] BF16 weight to RAWINT4 layout."""
+    n, k = weight_bf16.shape
+    assert k % 2 == 0
+    assert k % group_size == 0
+
+    weight_fp32 = weight_bf16.float()
+    qweight = torch.zeros((n, k // 2), dtype=torch.uint8)
+    scales = torch.zeros((n, k // group_size), dtype=torch.bfloat16)
+
+    for ni in range(n):
+        for g in range(k // group_size):
+            k_start = g * group_size
+            k_end = k_start + group_size
+            block = weight_fp32[ni, k_start:k_end]
+            amax = block.abs().max().item()
+            scale = amax / 7.0 if amax > 0 else 1.0
+            scales[ni, g] = scale
+
+            for kk in range(k_start, k_end, 2):
+                q0 = int(round(weight_fp32[ni, kk].item() / scale)) + 8
+                q1 = int(round(weight_fp32[ni, kk + 1].item() / scale)) + 8
+                q0 = max(0, min(15, q0))
+                q1 = max(0, min(15, q1))
+                qweight[ni, kk // 2] = (q1 << 4) | q0
+
+    return qweight, scales
+
+
+def rawint4_dequantize(qweight, scales, out_features, in_features):
+    """Dequantize RAWINT4 qweight/scales back to fp32 [N, K]."""
+    result = torch.zeros((out_features, in_features), dtype=torch.float32)
+    for ni in range(out_features):
+        for g in range(in_features // group_size):
+            scale = scales[ni, g].float().item()
+            k_start = g * group_size
+            k_end = k_start + group_size
+            for kk in range(k_start, k_end, 2):
+                packed = int(qweight[ni, kk // 2].item())
+                result[ni, kk] = ((packed & 0x0F) - 8) * scale
+                result[ni, kk + 1] = (((packed >> 4) & 0x0F) - 8) * scale
+    return result
+
+
+def pack_rawint4_uint8_as_int32(qweight):
+    """Pack byte RAWINT4 layout into compressed-tensors int32 storage."""
+    assert qweight.dtype == torch.uint8
+    assert qweight.shape[1] % 4 == 0
+    return qweight.contiguous().view(torch.int32).contiguous()
+
+
+def act_fn(x):
+    return x / (1.0 + torch.exp(-x))
+
+
+def mlp_torch(input_data, gate_proj, up_proj, down_proj):
+    gate_buf = torch.mm(input_data, gate_proj.t())
+    up_buf = torch.mm(input_data, up_proj.t())
+    intermediate = act_fn(gate_buf) * up_buf
+    return torch.mm(intermediate, down_proj.t())
+
+
+def moe_torch(input_data, 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_data[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 = sorted_tokens[start_idx:end_idx]
+        out = mlp_torch(tokens, gate_proj[i], up_proj[i], down_proj[i])
+        outputs.append(out)
+        start_idx = end_idx
+    outs = torch.cat(outputs, dim=0) if 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).float().mul_(weights.unsqueeze(-1)).sum(1)).to(new_x.dtype)
+
+
+def available_backends():
+    backends = []
+    if hasattr(kt_kernel_ext.moe, "AVX2RawInt4_MOE"):
+        backends.append(("AVX2RawInt4_MOE", kt_kernel_ext.moe.AVX2RawInt4_MOE, 0.12))
+
+    if hasattr(kt_kernel_ext.moe, "AVXVNNI256RawInt4_MOE"):
+        has_avx_vnni = False
+        try:
+            with open("/proc/cpuinfo", "r") as f:
+                has_avx_vnni = any(("avx_vnni" in line or "avxvnni" in line) for line in f if line.startswith("flags"))
+        except OSError:
+            has_avx_vnni = False
+        if has_avx_vnni:
+            backends.append(("AVXVNNI256RawInt4_MOE", kt_kernel_ext.moe.AVXVNNI256RawInt4_MOE, 0.20))
+    return backends
+
+
+def run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen):
+    physical_to_logical_map = torch.tensor(range(expert_num), dtype=torch.int64).contiguous()
+    cpu_infer = kt_kernel_ext.CPUInfer(CPUINFER_PARAM)
+
+    with torch.inference_mode():
+        gate_bf16 = (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 10.0).to(
+            torch.bfloat16
+        )
+        up_bf16 = (torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32) / 10.0).to(
+            torch.bfloat16
+        )
+        down_bf16 = (torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32) / 10.0).to(
+            torch.bfloat16
+        )
+
+        gate_qw_list, gate_scale_list = [], []
+        up_qw_list, up_scale_list = [], []
+        down_qw_list, down_scale_list = [], []
+
+        for e in range(expert_num):
+            qw, sc = rawint4_quantize(gate_bf16[e])
+            gate_qw_list.append(qw)
+            gate_scale_list.append(sc)
+
+            qw, sc = rawint4_quantize(up_bf16[e])
+            up_qw_list.append(qw)
+            up_scale_list.append(sc)
+
+            qw, sc = rawint4_quantize(down_bf16[e])
+            down_qw_list.append(qw)
+            down_scale_list.append(sc)
+
+        gate_qw = torch.stack(gate_qw_list).contiguous()
+        gate_scales = torch.stack(gate_scale_list).contiguous()
+        up_qw = torch.stack(up_qw_list).contiguous()
+        up_scales = torch.stack(up_scale_list).contiguous()
+        down_qw = torch.stack(down_qw_list).contiguous()
+        down_scales = torch.stack(down_scale_list).contiguous()
+
+        gate_deq = torch.stack(
+            [
+                rawint4_dequantize(gate_qw_list[e], gate_scale_list[e], intermediate_size, hidden_size)
+                for e in range(expert_num)
+            ]
+        )
+        up_deq = torch.stack(
+            [
+                rawint4_dequantize(up_qw_list[e], up_scale_list[e], intermediate_size, hidden_size)
+                for e in range(expert_num)
+            ]
+        )
+        down_deq = torch.stack(
+            [
+                rawint4_dequantize(down_qw_list[e], down_scale_list[e], hidden_size, intermediate_size)
+                for e in range(expert_num)
+            ]
+        )
+
+        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_qw.data_ptr()
+        config.up_proj = up_qw.data_ptr()
+        config.down_proj = down_qw.data_ptr()
+        config.gate_scale = gate_scales.data_ptr()
+        config.up_scale = up_scales.data_ptr()
+        config.down_scale = down_scales.data_ptr()
+        config.quant_config.bits = 4
+        config.quant_config.group_size = group_size
+        config.quant_config.zero_point = False
+        config.pool = cpu_infer.backend_
+
+        moe = backend_cls(config)
+        cpu_infer.submit(moe.load_weights_task(physical_to_logical_map.data_ptr()))
+        cpu_infer.sync()
+
+        print(f"\n--- {backend_name} (qlen={qlen}) ---")
+        for i in range(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_data = (torch.randn((qlen, hidden_size), dtype=torch.float32) / 100.0).to(torch.bfloat16).contiguous()
+            output = torch.empty((qlen, hidden_size), dtype=torch.bfloat16).contiguous()
+
+            bsz_tensor = torch.tensor([qlen], dtype=torch.int32)
+            cpu_infer.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,
+                )
+            )
+            cpu_infer.sync()
+
+            ref_output = moe_torch(input_data.float(), expert_ids, weights, gate_deq, up_deq, down_deq).to(
+                torch.bfloat16
+            )
+            diff = torch.mean(torch.abs(output.float() - ref_output.float())) / (
+                torch.mean(torch.abs(ref_output.float())) + 1e-8
+            )
+            print(f"  Iteration {i}: diff = {diff.item():.6f}")
+            assert diff < threshold, f"{backend_name} accuracy test failed: diff={diff.item():.6f} >= {threshold}"
+
+
+def test_rawint4_accuracy():
+    backends = available_backends()
+    if not backends:
+        print("Skipping RAWINT4 accuracy tests: no x86 RAWINT4 backend available")
+        return
+
+    for backend_name, backend_cls, threshold in backends:
+        run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen=1)
+        run_backend_accuracy_test(backend_name, backend_cls, threshold, qlen=16)
+
+
+def test_compressed_loader_normalizes_int32_pack_quantized_weights():
+    load_amx_utils()
+    loader_mod = sys.modules["kt_kernel.utils.loader"]
+
+    weight_bf16 = (torch.randn((intermediate_size, hidden_size), dtype=torch.float32) / 10.0).to(torch.bfloat16)
+    qweight, scales = rawint4_quantize(weight_bf16)
+    packed_int32 = pack_rawint4_uint8_as_int32(qweight)
+    weight_shape = torch.tensor([intermediate_size, hidden_size], dtype=torch.int32)
+
+    normalized = loader_mod.CompressedSafeTensorLoader._normalize_rawint4_weight(
+        packed_int32, scales, weight_shape, "test.weight_packed"
+    )
+
+    assert normalized.dtype == torch.uint8
+    assert normalized.shape == qweight.shape
+    assert torch.equal(normalized, qweight)
+
+
+def test_compressed_loader_accepts_uint8_rawint4_weights():
+    load_amx_utils()
+    loader_mod = sys.modules["kt_kernel.utils.loader"]
+
+    weight_bf16 = (torch.randn((intermediate_size, hidden_size), dtype=torch.float32) / 10.0).to(torch.bfloat16)
+    qweight, scales = rawint4_quantize(weight_bf16)
+    weight_shape = torch.tensor([intermediate_size, hidden_size], dtype=torch.int32)
+
+    normalized = loader_mod.CompressedSafeTensorLoader._normalize_rawint4_weight(
+        qweight, scales, weight_shape, "test.weight_packed"
+    )
+
+    assert normalized.dtype == torch.uint8
+    assert normalized.shape == qweight.shape
+    assert torch.equal(normalized, qweight)
+
+
+def test_compressed_loader_ignores_invalid_weight_shape_metadata():
+    load_amx_utils()
+    loader_mod = sys.modules["kt_kernel.utils.loader"]
+
+    weight_bf16 = (torch.randn((intermediate_size, hidden_size), dtype=torch.float32) / 10.0).to(torch.bfloat16)
+    qweight, scales = rawint4_quantize(weight_bf16)
+    packed_int32 = pack_rawint4_uint8_as_int32(qweight)
+    invalid_shape = torch.tensor([-1752796263, -1707567530], dtype=torch.int32)
+
+    normalized = loader_mod.CompressedSafeTensorLoader._normalize_rawint4_weight(
+        packed_int32, scales, invalid_shape, "test.weight_packed"
+    )
+
+    assert normalized.dtype == torch.uint8
+    assert normalized.shape == qweight.shape
+    assert torch.equal(normalized, qweight)
+
+
+def test_compressed_loader_ignores_odd_weight_shape_metadata():
+    load_amx_utils()
+    loader_mod = sys.modules["kt_kernel.utils.loader"]
+
+    weight_bf16 = (torch.randn((intermediate_size, hidden_size), dtype=torch.float32) / 10.0).to(torch.bfloat16)
+    qweight, scales = rawint4_quantize(weight_bf16)
+    packed_int32 = pack_rawint4_uint8_as_int32(qweight)
+    invalid_shape = torch.tensor([241597647, 1216029047], dtype=torch.int32)
+
+    normalized = loader_mod.CompressedSafeTensorLoader._normalize_rawint4_weight(
+        packed_int32, scales, invalid_shape, "test.weight_packed"
+    )
+
+    assert normalized.dtype == torch.uint8
+    assert normalized.shape == qweight.shape
+    assert torch.equal(normalized, qweight)
+
+
+def test_rawint4_backend_selection_falls_back_to_avx2_for_large_group_size(monkeypatch):
+    amx_utils = load_amx_utils()
+    fake_amx_backend = object()
+    fake_avx2_backend = object()
+    fake_avxvnni_backend = object()
+
+    monkeypatch.setattr(amx_utils, "AMXInt4_KGroup_MOE", fake_amx_backend)
+    monkeypatch.setattr(amx_utils, "AVX2RawInt4_MOE", fake_avx2_backend)
+    monkeypatch.setattr(amx_utils, "AVXVNNI256RawInt4_MOE", fake_avxvnni_backend)
+    monkeypatch.setattr(amx_utils, "_HAS_RAWINT4_SUPPORT", False)
+    monkeypatch.setattr(amx_utils, "_HAS_AVX2_RAWINT4_SUPPORT", True)
+    monkeypatch.setattr(amx_utils, "_HAS_AVXVNNI256_RAW_INT4_SUPPORT", True)
+    monkeypatch.setattr(amx_utils, "_HOST_HAS_AVX_VNNI", True)
+    monkeypatch.delenv("KT_RAWINT4_BACKEND", raising=False)
+
+    assert amx_utils._select_rawint4_backend(512) is fake_avx2_backend
+    assert amx_utils._select_rawint4_backend(128) is fake_avxvnni_backend
+
+
+def test_rawint4_backend_selection_rejects_forced_avxvnni_with_large_group_size(monkeypatch):
+    amx_utils = load_amx_utils()
+
+    monkeypatch.setattr(amx_utils, "_HAS_AVXVNNI256_RAW_INT4_SUPPORT", True)
+    monkeypatch.setattr(amx_utils, "_HOST_HAS_AVX_VNNI", True)
+    monkeypatch.setenv("KT_RAWINT4_BACKEND", "avxvnni")
+
+    with pytest.raises(RuntimeError, match="group_size=512 is unsupported"):
+        amx_utils._select_rawint4_backend(512)
+
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("RAWINT4 MoE Accuracy Test")
+    print("=" * 60)
+    test_rawint4_accuracy()
+    print("PASSED")

kt-kernel/test/per_commit/test_port_checker.py

@@ -0,0 +1,45 @@
+import importlib.util
+import socket
+from pathlib import Path
+import unittest
+from unittest.mock import MagicMock, patch
+
+from ci.ci_register import register_cpu_ci
+
+
+register_cpu_ci(est_time=0.1, suite="default")
+
+
+PORT_CHECKER_PATH = Path(__file__).resolve().parents[2] / "python" / "cli" / "utils" / "port_checker.py"
+SPEC = importlib.util.spec_from_file_location("port_checker", PORT_CHECKER_PATH)
+assert SPEC is not None and SPEC.loader is not None
+port_checker = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(port_checker)
+
+
+class TestPortChecker(unittest.TestCase):
+    def test_bound_port_is_not_available_before_listen(self):
+        holder = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        try:
+            holder.bind(("127.0.0.1", 0))
+            port = holder.getsockname()[1]
+
+            self.assertFalse(port_checker.is_port_available("127.0.0.1", port))
+            self.assertEqual(port_checker.find_available_port("127.0.0.1", port, max_attempts=1), (False, port))
+        finally:
+            holder.close()
+
+    def test_non_windows_bind_check_uses_reuseaddr(self):
+        sock = MagicMock()
+        sock.__enter__.return_value = sock
+
+        with patch.object(port_checker.sys, "platform", "linux"):
+            with patch.object(port_checker.socket, "socket", return_value=sock):
+                self.assertTrue(port_checker.is_port_available("127.0.0.1", 12345))
+
+        sock.setsockopt.assert_called_once_with(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        sock.bind.assert_called_once_with(("127.0.0.1", 12345))
+
+
+if __name__ == "__main__":
+    unittest.main()