From 79b265b2f69a59c7b6d16baaf8b955eebd8ae6c4 Mon Sep 17 00:00:00 2001 From: lutianshu824 <18610028240@126.com> Date: Mon, 6 Jul 2026 18:06:52 +0800 Subject: [PATCH 1/2] fix: normalize compressed RAWINT4 weights (#2075) * fix: normalize compressed RAWINT4 weights * docs: add Hygon DCU ROCm notes --------- Co-authored-by: lutianshu824 --- doc/en/ROCm.md | 38 +++++++++ kt-kernel/python/utils/loader.py | 46 +++++++++++ .../per_commit/test_moe_rawint4_accuracy.py | 78 +++++++++++++++++++ 3 files changed, 162 insertions(+) diff --git a/doc/en/ROCm.md b/doc/en/ROCm.md index 39f48902..5810e8b7 100644 --- a/doc/en/ROCm.md +++ b/doc/en/ROCm.md @@ -42,6 +42,44 @@ pip3 install packaging ninja cpufeature numpy > **Tip:** For other ROCm versions, visit [PyTorch Previous Versions](https://pytorch.org/get-started/previous-versions/) +### Hygon DCU / DTK Notes + +Hygon DCU uses a ROCm-compatible DTK stack. For DCU systems, use the Hygon DCU +PyTorch environment that matches your DTK release instead of installing the +generic PyPI `torch` package or the official AMD ROCm wheel. + +For example, a reported working `gfx936` setup uses a Hygon DCU PyTorch image +from the SourceFind/Hygon developer image portal: + +https://sourcefind.cn/#/image/dcu/pytorch + +The reported environment provides: + +- DTK 26.04, typically under `/opt/dtk` +- PyTorch `2.5.1+das.opt1.dtk2604` +- Python 3.10 + +Before building, verify that the DCU PyTorch package is active: + +```bash +python -c "import torch; print(torch.__version__); print(torch.version.hip); print(torch.__file__)" +``` + +Then build `kt-kernel` without letting pip replace the vendor PyTorch package: + +```bash +export CPUINFER_USE_ROCM=1 +export PYTORCH_ROCM_ARCH=gfx936 +export ROCM_PATH=/opt/dtk # change this if DTK is installed elsewhere + +cd kt-kernel +pip install . --no-build-isolation --no-deps +``` + +> **Tip:** Keep `--no-deps` when building in a vendor PyTorch environment. A +> plain `pip install .` may resolve `kt-kernel`'s normal `torch` dependency and +> shadow or replace the installed DCU PyTorch package with a generic torch wheel. + ### 4. Build ktransformers ```bash diff --git a/kt-kernel/python/utils/loader.py b/kt-kernel/python/utils/loader.py index 25178229..c24f7893 100644 --- a/kt-kernel/python/utils/loader.py +++ b/kt-kernel/python/utils/loader.py @@ -679,6 +679,49 @@ class BF16SafeTensorLoader(SafeTensorLoader): class CompressedSafeTensorLoader(SafeTensorLoader): """Loader for compressed SafeTensor layouts (RAWINT4 weights).""" + @staticmethod + def _normalize_rawint4_weight(weight_tensor, scale_tensor, shape_tensor=None, key: str = "weight_packed"): + """Return byte-packed uint8 RAWINT4 weights expected by kt_kernel_ext.""" + if weight_tensor.dtype == torch.int32: + # compressed-tensors pack-quantized stores 8 int4 values per int32. + # The RAWINT4 kernels consume the same bytes as uint8, two int4 values per byte. + rows, int32_cols = weight_tensor.shape + weight_tensor = weight_tensor.contiguous().view(torch.uint8).view(rows, int32_cols * 4).contiguous() + elif weight_tensor.dtype == torch.uint8: + weight_tensor = weight_tensor.contiguous() + else: + raise TypeError(f"{key} must be torch.uint8 or torch.int32, got {weight_tensor.dtype}") + + if shape_tensor is None: + return weight_tensor + + shape_values = shape_tensor.detach().cpu().tolist() + if len(shape_values) != 2: + raise ValueError(f"{key}.weight_shape must contain [out_features, in_features], got {shape_values}") + + out_features, in_features = (int(shape_values[0]), int(shape_values[1])) + if out_features <= 0 or in_features <= 0: + return weight_tensor + + if in_features % 2 != 0: + return weight_tensor + + expected_weight_shape = (out_features, in_features // 2) + if tuple(weight_tensor.shape) != expected_weight_shape: + return weight_tensor + + if scale_tensor.dim() != 2 or scale_tensor.shape[0] != out_features or scale_tensor.shape[1] <= 0: + raise ValueError( + f"{key} scale shape {tuple(scale_tensor.shape)} is incompatible with weight_shape={shape_values}" + ) + + if in_features % int(scale_tensor.shape[1]) != 0: + raise ValueError( + f"{key} in_features={in_features} is not divisible by scale columns={scale_tensor.shape[1]}" + ) + + return weight_tensor + def load_experts(self, base_key: str, device: str = "cpu"): """Load raw expert weights stored in compressed safetensor format.""" @@ -703,6 +746,7 @@ class CompressedSafeTensorLoader(SafeTensorLoader): for exp_id in range(expert_idx): weight_key = f"{experts_prefix}.{exp_id}.{proj_name}_proj.weight_packed" scale_key = f"{experts_prefix}.{exp_id}.{proj_name}_proj.weight_scale" + shape_key = f"{experts_prefix}.{exp_id}.{proj_name}_proj.weight_shape" if not self.has_tensor(weight_key): raise KeyError(f"Missing tensor: {weight_key}") @@ -711,6 +755,8 @@ class CompressedSafeTensorLoader(SafeTensorLoader): weight_tensor = self.load_tensor(weight_key, device).contiguous() scale_tensor = self.load_tensor(scale_key, device).contiguous() + shape_tensor = self.load_tensor(shape_key, "cpu") if self.has_tensor(shape_key) else None + weight_tensor = self._normalize_rawint4_weight(weight_tensor, scale_tensor, shape_tensor, weight_key) weight_entries.append(weight_tensor) scale_entries.append(scale_tensor) diff --git a/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py b/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py index 36054670..e980be8b 100644 --- a/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py +++ b/kt-kernel/test/per_commit/test_moe_rawint4_accuracy.py @@ -110,6 +110,13 @@ def rawint4_dequantize(qweight, scales, out_features, in_features): 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)) @@ -279,6 +286,77 @@ def test_rawint4_accuracy(): 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() From cb9f47d142a507cac5d74450b30463d2e8d1cf58 Mon Sep 17 00:00:00 2001 From: VectorPeak Date: Mon, 6 Jul 2026 18:25:31 +0800 Subject: [PATCH 2/2] [fix](cli): detect bound ports before launch (#2071) * [fix](cli): detect bound ports before launch * [fix](cli): align port reuse check by platform --- kt-kernel/python/cli/utils/port_checker.py | 23 ++++------ .../test/per_commit/test_port_checker.py | 45 +++++++++++++++++++ 2 files changed, 53 insertions(+), 15 deletions(-) create mode 100644 kt-kernel/test/per_commit/test_port_checker.py diff --git a/kt-kernel/python/cli/utils/port_checker.py b/kt-kernel/python/cli/utils/port_checker.py index ffdf209e..d5cf09c2 100644 --- a/kt-kernel/python/cli/utils/port_checker.py +++ b/kt-kernel/python/cli/utils/port_checker.py @@ -3,6 +3,7 @@ Port availability checking utilities. """ import socket +import sys from typing import Tuple @@ -17,22 +18,14 @@ def is_port_available(host: str, port: int) -> bool: True if port is available, False if occupied """ try: - # Try to bind to the port - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.settimeout(1) + bind_host = "" if host == "0.0.0.0" else host + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + if sys.platform != "win32": + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind((bind_host, port)) + return True - # Use SO_REUSEADDR to allow binding to recently closed ports - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - - # Try to bind - result = sock.connect_ex((host if host != "0.0.0.0" else "127.0.0.1", port)) - sock.close() - - # If connect_ex returns 0, port is occupied - # If it returns error (non-zero), port is available - return result != 0 - - except Exception: + except OSError: # If any error occurs, assume port is not available return False diff --git a/kt-kernel/test/per_commit/test_port_checker.py b/kt-kernel/test/per_commit/test_port_checker.py new file mode 100644 index 00000000..fd551bf7 --- /dev/null +++ b/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()