mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-10-18 22:55:21 +08:00
losewayy
a02efe36b5
* [feat](kt-kernel): add AVX2 RAWINT4 routed-expert LoRA SFT backend * [fix](kt-kernel): fix RAWINT4 SFT wrapper and clamp gradients Allow the RAWINT4 production wrapper on AVX2 runtimes and reject shared backward pools before construction. Preserve FP32 clamp decisions across BF16 caches, remove unused configuration state, and cover rectangular dimensions, multiblock dX, clamp boundaries, and production wrapper tasks with tighter LoRA gradient checks. * [fix](kt-kernel): plumb swiglu_limit through the RAWINT4 SFT python path The production wrapper hard-zeroed swiglu_limit for every format except MXFP4, and both SFT python gates rejected non-zero values for RAWINT4 even though MOESFTConfig carries the field and the AVX2 kernel implements the clamp in forward plus the gradient mask in backward. RAWINT4_SFT now accepts a finite non-negative swiglu_limit end to end, while the AMX K2 part (which computes plain SiLU) fails fast instead of silently dropping the clamp. The production-wrapper contract test is parametrized over swiglu_limit 0 and 2, a model-config plumb test covers plain and nested text_config shapes, the __main__ entry exercises both limits, and bench_rawint4_sft.py gains --swiglu-limit.
366 lines
11 KiB
Python
366 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
SFT_PATH = Path(__file__).resolve().parents[2] / "python" / "sft"
|
|
PACKAGE_NAME = "kt_sft_backend_under_test"
|
|
package = ModuleType(PACKAGE_NAME)
|
|
package.__path__ = [str(SFT_PATH)]
|
|
sys.modules[PACKAGE_NAME] = package
|
|
spec = importlib.util.spec_from_file_location(
|
|
f"{PACKAGE_NAME}.backend",
|
|
SFT_PATH / "backend.py",
|
|
)
|
|
assert spec is not None and spec.loader is not None
|
|
backend = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = backend
|
|
spec.loader.exec_module(backend)
|
|
|
|
|
|
def _install_fake_runtime(
|
|
monkeypatch,
|
|
*,
|
|
variant: str,
|
|
kernel: str,
|
|
layout: str = backend.INT8_WEIGHT_LAYOUT,
|
|
has_int8_sft: bool = True,
|
|
has_fp8_sft: bool = False,
|
|
fp8_kernel: str = backend.FP8_KERNEL,
|
|
fp8_layout: str = backend.FP8_WEIGHT_LAYOUT,
|
|
has_rawint4_sft: bool = False,
|
|
rawint4_kernel: str = backend.RAWINT4_KERNEL,
|
|
rawint4_layout: str = backend.RAWINT4_WEIGHT_LAYOUT,
|
|
):
|
|
moe = SimpleNamespace()
|
|
if has_int8_sft:
|
|
moe.AMXInt8_SFT_MOE = object()
|
|
if has_fp8_sft:
|
|
moe.AMXFP8_SFT_MOE = object()
|
|
if has_rawint4_sft:
|
|
moe.AMXInt4_KGroup_SFT_MOE = object()
|
|
extension = SimpleNamespace(
|
|
__cpu_variant__=variant,
|
|
__int8_kernel__=kernel,
|
|
__int8_weight_layout__=layout,
|
|
__fp8_kernel__=fp8_kernel,
|
|
__fp8_weight_layout__=fp8_layout,
|
|
__rawint4_kernel__=rawint4_kernel,
|
|
__rawint4_weight_layout__=rawint4_layout,
|
|
moe=moe,
|
|
)
|
|
fake_package = SimpleNamespace(
|
|
kt_kernel_ext=extension,
|
|
__cpu_variant__=variant,
|
|
)
|
|
monkeypatch.setitem(sys.modules, "kt_kernel", fake_package)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("variant", "kernel"),
|
|
[
|
|
("amx", "amx-int8"),
|
|
("avx512_bf16", "avx512-vnni"),
|
|
("avx512_bf16", "onednn-vnni"),
|
|
],
|
|
)
|
|
def test_int8_runtime_reports_effective_kernel(monkeypatch, variant, kernel):
|
|
_install_fake_runtime(monkeypatch, variant=variant, kernel=kernel)
|
|
|
|
runtime = backend.get_int8_runtime()
|
|
|
|
assert runtime.cpu_variant == variant
|
|
assert runtime.kernel == kernel
|
|
assert runtime.weight_layout == backend.INT8_WEIGHT_LAYOUT
|
|
|
|
|
|
def test_int8_runtime_rejects_extension_without_sft_symbol(monkeypatch):
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="avx512_bf16",
|
|
kernel="avx512-vnni",
|
|
has_int8_sft=False,
|
|
)
|
|
with pytest.raises(RuntimeError, match="does not provide INT8 SFT"):
|
|
backend.get_int8_runtime()
|
|
|
|
|
|
def test_int8_runtime_rejects_extension_without_runtime_metadata(monkeypatch):
|
|
extension = SimpleNamespace(moe=SimpleNamespace(AMXInt8_SFT_MOE=object()))
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"kt_kernel",
|
|
SimpleNamespace(kt_kernel_ext=extension, __cpu_variant__="avx512_bf16"),
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="predates production INT8 runtime metadata"):
|
|
backend.get_int8_runtime()
|
|
|
|
|
|
def test_int8_runtime_rejects_unsupported_isa(monkeypatch):
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="avx2",
|
|
kernel="unsupported",
|
|
)
|
|
with pytest.raises(RuntimeError, match="requires an AMX-INT8 or AVX512"):
|
|
backend.get_int8_runtime()
|
|
|
|
|
|
def test_int8_runtime_rejects_weight_layout_mismatch(monkeypatch):
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="amx",
|
|
kernel="amx-int8",
|
|
layout="different-layout",
|
|
)
|
|
with pytest.raises(RuntimeError, match="weight-layout mismatch"):
|
|
backend.get_int8_runtime()
|
|
|
|
|
|
def test_fp8_runtime_requires_native_sft_symbol(monkeypatch):
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="amx",
|
|
kernel="amx-int8",
|
|
has_fp8_sft=False,
|
|
)
|
|
with pytest.raises(RuntimeError, match="does not provide native FP8 SFT"):
|
|
backend.get_fp8_runtime()
|
|
|
|
|
|
def test_fp8_runtime_reports_native_layout(monkeypatch):
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="amx",
|
|
kernel="amx-int8",
|
|
has_fp8_sft=True,
|
|
)
|
|
runtime = backend.get_fp8_runtime()
|
|
assert runtime.cpu_variant == "amx"
|
|
assert runtime.kernel == backend.FP8_KERNEL
|
|
assert runtime.weight_layout == backend.FP8_WEIGHT_LAYOUT
|
|
|
|
|
|
def test_fp8_runtime_rejects_extension_without_runtime_metadata(monkeypatch):
|
|
extension = SimpleNamespace(moe=SimpleNamespace(AMXFP8_SFT_MOE=object()))
|
|
monkeypatch.setitem(
|
|
sys.modules,
|
|
"kt_kernel",
|
|
SimpleNamespace(kt_kernel_ext=extension, __cpu_variant__="avx512_bf16"),
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="predates native FP8 runtime metadata"):
|
|
backend.get_fp8_runtime()
|
|
|
|
|
|
def test_fp8_runtime_rejects_unsupported_kernel(monkeypatch):
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="avx2",
|
|
kernel="unsupported",
|
|
has_fp8_sft=True,
|
|
fp8_kernel="unsupported",
|
|
)
|
|
with pytest.raises(RuntimeError, match="requires the AVX512-BF16\\+VBMI"):
|
|
backend.get_fp8_runtime()
|
|
|
|
|
|
def test_fp8_runtime_rejects_weight_layout_mismatch(monkeypatch):
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="avx512_bf16",
|
|
kernel="avx512-vnni",
|
|
has_fp8_sft=True,
|
|
fp8_layout="different-layout",
|
|
)
|
|
with pytest.raises(RuntimeError, match="weight-layout mismatch"):
|
|
backend.get_fp8_runtime()
|
|
|
|
|
|
@pytest.mark.parametrize("configured", ["auto", "FP8", "fp8"])
|
|
def test_fp8_backend_normalization(configured):
|
|
assert (
|
|
backend.normalize_sft_backend(
|
|
configured,
|
|
expert_weight_format="fp8",
|
|
)
|
|
== backend.FP8_BACKEND
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("configured", ["auto", "RAWINT4", "rawint4"])
|
|
def test_rawint4_backend_normalization(configured):
|
|
assert (
|
|
backend.normalize_sft_backend(
|
|
configured,
|
|
expert_weight_format="rawint4",
|
|
)
|
|
== backend.RAWINT4_BACKEND
|
|
)
|
|
|
|
|
|
def test_rawint4_legacy_backend_alias_warns():
|
|
with pytest.warns(FutureWarning, match="deprecated"):
|
|
normalized = backend.normalize_sft_backend(
|
|
"AMXINT4_KGroup",
|
|
expert_weight_format="rawint4",
|
|
)
|
|
assert normalized == backend.RAWINT4_BACKEND
|
|
|
|
|
|
def test_rawint4_runtime_reports_exact_contract(monkeypatch):
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="amx",
|
|
kernel="amx-int8",
|
|
has_rawint4_sft=True,
|
|
)
|
|
|
|
runtime = backend.get_rawint4_runtime()
|
|
|
|
assert runtime.cpu_variant == "amx"
|
|
assert runtime.kernel == backend.RAWINT4_KERNEL
|
|
assert runtime.weight_layout == backend.RAWINT4_WEIGHT_LAYOUT
|
|
|
|
|
|
def test_rawint4_runtime_rejects_missing_symbol(monkeypatch):
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="amx",
|
|
kernel="amx-int8",
|
|
has_rawint4_sft=False,
|
|
)
|
|
with pytest.raises(RuntimeError, match="does not provide RAWINT4 SFT"):
|
|
backend.get_rawint4_runtime()
|
|
|
|
|
|
def test_rawint4_runtime_reports_avx2_kernel(monkeypatch):
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="avx2",
|
|
kernel="unsupported",
|
|
has_rawint4_sft=True,
|
|
rawint4_kernel=backend.RAWINT4_AVX2_KERNEL,
|
|
)
|
|
|
|
runtime = backend.get_rawint4_runtime()
|
|
|
|
assert runtime.cpu_variant == "avx2"
|
|
assert runtime.kernel == backend.RAWINT4_AVX2_KERNEL
|
|
assert runtime.weight_layout == backend.RAWINT4_WEIGHT_LAYOUT
|
|
|
|
|
|
def test_rawint4_runtime_reports_amx_kernel_on_avx512_bf16(monkeypatch):
|
|
# The avx512_bf16 wheel tier is built with USE_AMX_AVX_KERNEL (without
|
|
# AMX intrinsics) and binds the AMX KGroup part on vector fallbacks, so
|
|
# it reports the AMX kernel tier rather than the AVX2 tier.
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="avx512_bf16",
|
|
kernel="unsupported",
|
|
has_rawint4_sft=True,
|
|
rawint4_kernel=backend.RAWINT4_KERNEL,
|
|
)
|
|
|
|
runtime = backend.get_rawint4_runtime()
|
|
|
|
assert runtime.cpu_variant == "avx512_bf16"
|
|
assert runtime.kernel == backend.RAWINT4_KERNEL
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"variant",
|
|
["avx2", "avx512_base", "avx512_vnni", "avx512_vbmi"],
|
|
)
|
|
def test_rawint4_runtime_uses_avx2_kernel_on_non_amx_variants(monkeypatch, variant):
|
|
# Every non-AMX x86_64 wheel without USE_AMX_AVX_KERNEL binds the same
|
|
# AVX2 SFT kernel under the AMXInt4_KGroup_SFT_MOE symbol, so any such
|
|
# variant whose extension reports the AVX2 kernel tier is accepted.
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant=variant,
|
|
kernel="unsupported",
|
|
has_rawint4_sft=True,
|
|
rawint4_kernel=backend.RAWINT4_AVX2_KERNEL,
|
|
)
|
|
|
|
runtime = backend.get_rawint4_runtime()
|
|
|
|
assert runtime.cpu_variant == variant
|
|
assert runtime.kernel == backend.RAWINT4_AVX2_KERNEL
|
|
|
|
|
|
def test_rawint4_runtime_rejects_kernel_tier_mismatch(monkeypatch):
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="avx2",
|
|
kernel="unsupported",
|
|
has_rawint4_sft=True,
|
|
rawint4_kernel=backend.RAWINT4_KERNEL,
|
|
)
|
|
with pytest.raises(RuntimeError, match="AMX KGroup or AVX2 kernel"):
|
|
backend.get_rawint4_runtime()
|
|
|
|
|
|
def test_rawint4_runtime_rejects_unsupported_kernel_tier(monkeypatch):
|
|
_install_fake_runtime(
|
|
monkeypatch,
|
|
variant="avx2",
|
|
kernel="unsupported",
|
|
has_rawint4_sft=True,
|
|
rawint4_kernel="unsupported",
|
|
)
|
|
with pytest.raises(RuntimeError, match="AMX KGroup or AVX2 kernel"):
|
|
backend.get_rawint4_runtime()
|
|
|
|
|
|
def _sap4_rawint4_quantization_config():
|
|
return {
|
|
"quant_method": "compressed-tensors",
|
|
"format": "pack-quantized",
|
|
"quantization_status": "compressed",
|
|
"config_groups": {
|
|
"group_0": {
|
|
"input_activations": None,
|
|
"output_activations": None,
|
|
"targets": ["Linear"],
|
|
"weights": {
|
|
"actorder": None,
|
|
"block_structure": None,
|
|
"dynamic": False,
|
|
"group_size": 32,
|
|
"num_bits": 4,
|
|
"observer": "minmax",
|
|
"observer_kwargs": {},
|
|
"strategy": "group",
|
|
"symmetric": True,
|
|
"type": "int",
|
|
},
|
|
}
|
|
},
|
|
"ignore": ["lm_head"],
|
|
}
|
|
|
|
|
|
def test_rawint4_checkpoint_contract_accepts_real_metadata_shape():
|
|
contract = backend.get_rawint4_checkpoint_contract(
|
|
SimpleNamespace(text_config=SimpleNamespace(quantization_config=_sap4_rawint4_quantization_config()))
|
|
)
|
|
assert contract.bits == 4
|
|
assert contract.group_size == 32
|
|
assert contract.signed is True
|
|
assert contract.zero_point is False
|
|
|
|
|
|
@pytest.mark.parametrize("zero_point", [False, True, None, 0])
|
|
def test_rawint4_checkpoint_contract_rejects_explicit_zero_point(zero_point):
|
|
quantization_config = _sap4_rawint4_quantization_config()
|
|
quantization_config["config_groups"]["group_0"]["weights"]["zero_point"] = zero_point
|
|
with pytest.raises(ValueError, match="must not declare zero-point metadata"):
|
|
backend.get_rawint4_checkpoint_contract({"quantization_config": quantization_config})
|