diff --git a/python/sglang/srt/models/xllm.py b/python/sglang/srt/models/xllm.py index 10cc99332..69ec06355 100644 --- a/python/sglang/srt/models/xllm.py +++ b/python/sglang/srt/models/xllm.py @@ -661,10 +661,10 @@ def _validate_mova_config( "the released checkpoints persist float32 dtype metadata but " "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( - "Native xLLM/K2 Horizon serving does not support quantized " - "model weights" + "Native xLLM/K2 Horizon serving supports only " + "compressed-tensors quantized model weights" ) runtime = get_exec() @@ -1686,6 +1686,14 @@ class XllmModel(nn.Module): class XllmForCausalLM(nn.Module): 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__( self, config: PretrainedConfig, diff --git a/test/registered/unit/models/test_xllm.py b/test/registered/unit/models/test_xllm.py index b7607dae3..f5107ed69 100644 --- a/test/registered/unit/models/test_xllm.py +++ b/test/registered/unit/models/test_xllm.py @@ -266,15 +266,35 @@ def test_native_xllm_requires_bfloat16(monkeypatch): _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) + 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) with ( 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):