[fix](cli): detect bound ports before launch (#2071)

* [fix](cli): detect bound ports before launch

* [fix](cli): align port reuse check by platform

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

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
 

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)

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

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