[Model] Add K2 Horizon FP8 checkpoint support (#38033)

Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
Yash Akhauri
2026-09-05 13:18:26 +08:00
committed by GitHub
co-authored by Xiaoyu Zhang
parent e980c1a2f1
commit 756d0e0a85
2 changed files with 34 additions and 6 deletions
+11 -3
View File
@@ -661,10 +661,10 @@ def _validate_mova_config(
"the released checkpoints persist float32 dtype metadata but " "the released checkpoints persist float32 dtype metadata but "
"their weights and validated runtime contract are BF16." "their weights and validated runtime contract are BF16."
) )
if quant_config is not None: if quant_config is not None and quant_config.get_name() != "compressed_tensors":
raise ValueError( raise ValueError(
"Native xLLM/K2 Horizon serving does not support quantized " "Native xLLM/K2 Horizon serving supports only "
"model weights" "compressed-tensors quantized model weights"
) )
runtime = get_exec() runtime = get_exec()
@@ -1686,6 +1686,14 @@ class XllmModel(nn.Module):
class XllmForCausalLM(nn.Module): class XllmForCausalLM(nn.Module):
fall_back_to_pt_during_load = False fall_back_to_pt_during_load = False
# Quantized checkpoints store these projections separately. This mapping
# lets quantization configs resolve fused runtime modules and their ignore
# lists consistently.
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
"gate_up_proj": ["gate_proj", "up_proj"],
}
def __init__( def __init__(
self, self,
config: PretrainedConfig, config: PretrainedConfig,
+23 -3
View File
@@ -266,15 +266,35 @@ def test_native_xllm_requires_bfloat16(monkeypatch):
_validate_mova_config(config, quant_config=None) _validate_mova_config(config, quant_config=None)
def test_native_xllm_rejects_quantized_weights(monkeypatch): def test_native_xllm_accepts_compressed_tensors_quantization(monkeypatch):
config = XllmConfig(num_values=0, num_experts=0) config = XllmConfig(num_values=0, num_experts=0)
quant_config = SimpleNamespace(get_name=lambda: "compressed_tensors")
monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.bfloat16)
with get_context().override_server_args(**_native_runtime_config()):
_validate_mova_config(config, quant_config=quant_config)
def test_native_xllm_rejects_other_quantization(monkeypatch):
config = XllmConfig(num_values=0, num_experts=0)
quant_config = SimpleNamespace(get_name=lambda: "awq")
monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.bfloat16) monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.bfloat16)
with ( with (
get_context().override_server_args(**_native_runtime_config()), get_context().override_server_args(**_native_runtime_config()),
pytest.raises(ValueError, match="does not support quantized"), pytest.raises(ValueError, match="supports only compressed-tensors"),
): ):
_validate_mova_config(config, quant_config=object()) _validate_mova_config(config, quant_config=quant_config)
def test_native_xllm_declares_quantized_fused_module_mapping():
assert XllmForCausalLM.packed_modules_mapping == {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
"gate_up_proj": ["gate_proj", "up_proj"],
}
assert K2HorizonForCausalLM.packed_modules_mapping == (
XllmForCausalLM.packed_modules_mapping
)
def test_native_xllm_accepts_bfloat16_without_expert_remapping(monkeypatch): def test_native_xllm_accepts_bfloat16_without_expert_remapping(monkeypatch):