diff --git a/tests/model_executor/test_k2_horizon_rope_fold.py b/tests/model_executor/test_k2_horizon_rope_fold.py new file mode 100644 index 0000000000..07e490a041 --- /dev/null +++ b/tests/model_executor/test_k2_horizon_rope_fold.py @@ -0,0 +1,636 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +import vllm.model_executor.models.k2_horizon as k2_horizon +from tests.quantization.utils import is_quant_method_supported +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.distributed import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, +) +from vllm.model_executor.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.models.k2_horizon import ( + K2HorizonDecoderLayer, + K2HorizonModel, + K2HorizonRMSNorm, + _rope_weight_perm, + apply_partial_rope, + is_rope_weights_folding_supported, +) +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port +from vllm.utils.torch_utils import set_default_torch_dtype, set_random_seed + +DEVICE = current_platform.device_type +DTYPE = torch.bfloat16 +SHAPES = [ + (128, 64), + (128, 128), + (192, 128), + (96, 32), +] + + +def get_gptj_rope( + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + num_heads: int, + num_kv_heads: int, + head_dim: int, + rope_head_dim: int, + max_position: int, + base: float, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + rotary_emb = get_rope( + rope_head_dim, + max_position=max_position, + is_neox_style=True, + rope_parameters={"rope_theta": base}, + dtype=dtype, + ) + return apply_partial_rope( + rotary_emb, + positions, + q, + k, + num_heads, + num_kv_heads, + head_dim, + rope_head_dim, + fold_rope_weights=False, + ) + + +def _apply_head_perm( + x: torch.Tensor, num: int, head_dim: int, idx: torch.Tensor +) -> torch.Tensor: + return ( + x.reshape(*x.shape[:-1], num, head_dim)[..., idx].reshape(*x.shape).contiguous() + ) + + +def get_neox_rope( + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + num_heads: int, + num_kv_heads: int, + head_dim: int, + rope_head_dim: int, + max_position: int, + base: float, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + rotary_emb = get_rope( + head_dim, + max_position=max_position, + is_neox_style=True, + rope_parameters={"rope_theta": base, "rope_dim": rope_head_dim}, + dtype=dtype, + ) + return apply_partial_rope( + rotary_emb, + positions, + q, + k, + num_heads, + num_kv_heads, + head_dim, + rope_head_dim, + fold_rope_weights=True, + ) + + +def fold_qk_proj_weight( + weight: torch.Tensor, head_dim: int, idx: torch.Tensor +) -> torch.Tensor: + hidden = weight.shape[-1] + return weight.view(-1, head_dim, hidden)[:, idx, :].reshape(-1, hidden).contiguous() + + +def fold_qk_channels( + vec: torch.Tensor, head_dim: int, idx: torch.Tensor +) -> torch.Tensor: + return vec.view(-1, head_dim)[:, idx].reshape(-1).contiguous() + + +@pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES) +def test_rope_fold_matches_permute_then_rope( + head_dim, rope_head_dim, default_vllm_config +): + dtype = DTYPE + set_random_seed(0) + num_tokens = 17 + num_heads = 8 + num_kv_heads = 2 + hidden_size = 512 + max_position = 4096 + base = 10000.0 + + torch.set_default_dtype(dtype) + + positions = torch.randint(0, max_position, (num_tokens,), device=DEVICE) + + scale = hidden_size**-0.5 + x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=DEVICE) + w_q = ( + torch.randn(num_heads * head_dim, hidden_size, dtype=dtype, device=DEVICE) + * scale + ) + w_k = ( + torch.randn(num_kv_heads * head_dim, hidden_size, dtype=dtype, device=DEVICE) + * scale + ) + + idx = _rope_weight_perm(head_dim, rope_head_dim).to(DEVICE) + + q = torch.nn.functional.linear(x, w_q) + k = torch.nn.functional.linear(x, w_k) + q_gptj, k_gptj = get_gptj_rope( + positions, + q, + k, + num_heads, + num_kv_heads, + head_dim, + rope_head_dim, + max_position, + base, + dtype, + ) + + q_folded = torch.nn.functional.linear(x, fold_qk_proj_weight(w_q, head_dim, idx)) + k_folded = torch.nn.functional.linear(x, fold_qk_proj_weight(w_k, head_dim, idx)) + q_neox, k_neox = get_neox_rope( + positions, + q_folded, + k_folded, + num_heads, + num_kv_heads, + head_dim, + rope_head_dim, + max_position, + base, + dtype, + ) + + q_gptj_perm = _apply_head_perm(q_gptj, num_heads, head_dim, idx) + k_gptj_perm = _apply_head_perm(k_gptj, num_kv_heads, head_dim, idx) + + atol, rtol = 1e-2, 1e-2 + + # Two complementary checks. First, an element-wise check that the folded + # NeoX path equals the GPT-J path up to the channel permutation P + # (neox_rope(Pq, Pk) == P @ gptj_rope(q, k)), hence permuting q_gptj/k_gptj. + # Second, the scores check below: P cancels in q @ k.T. + torch.testing.assert_close(q_neox, q_gptj_perm, atol=atol, rtol=rtol) + torch.testing.assert_close(k_neox, k_gptj_perm, atol=atol, rtol=rtol) + + def scores(qt, kt): + qh = qt.reshape(num_tokens, num_heads, head_dim) + kh = kt.reshape(num_tokens, num_kv_heads, head_dim) + g = num_heads // num_kv_heads + qh = qh.reshape(num_tokens, num_kv_heads, g, head_dim).float() + kh = kh.float() + return torch.einsum("tkgd,skd->tksg", qh, kh) + + torch.testing.assert_close( + scores(q_neox, k_neox), scores(q_gptj, k_gptj), atol=atol, rtol=rtol + ) + + +@pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES) +@pytest.mark.parametrize("num_heads", [8, 2]) +def test_qk_bias_fold_matches_permute( + head_dim, rope_head_dim, num_heads, default_vllm_config +): + dtype = DTYPE + set_random_seed(0) + num_tokens = 17 + hidden_size = 512 + torch.set_default_dtype(dtype) + + scale = hidden_size**-0.5 + x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=DEVICE) + w = ( + torch.randn(num_heads * head_dim, hidden_size, dtype=dtype, device=DEVICE) + * scale + ) + b = torch.randn(num_heads * head_dim, dtype=dtype, device=DEVICE) + + idx = _rope_weight_perm(head_dim, rope_head_dim).to(DEVICE) + + ref = _apply_head_perm( + torch.nn.functional.linear(x, w, b), num_heads, head_dim, idx + ) + folded = torch.nn.functional.linear( + x, + fold_qk_proj_weight(w, head_dim, idx), + fold_qk_channels(b, head_dim, idx), + ) + + atol, rtol = 1e-2, 1e-2 + torch.testing.assert_close(folded, ref, atol=atol, rtol=rtol) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="requires a GPU (CUDA/ROCm)" +) +@pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES) +@pytest.mark.parametrize("num_heads", [8, 2]) +def test_qk_norm_scale_fold_matches_permute( + head_dim, rope_head_dim, num_heads, default_vllm_config +): + dtype = DTYPE + set_random_seed(0) + num_tokens = 13 + torch.set_default_dtype(dtype) + + hidden = num_heads * head_dim + idx = _rope_weight_perm(head_dim, rope_head_dim).to(DEVICE) + + x = torch.randn(num_tokens, hidden, dtype=dtype, device=DEVICE) + + norm = K2HorizonRMSNorm(hidden_size=hidden, n_groups=num_heads).to(DEVICE) + with torch.no_grad(): + norm.weight.copy_(torch.randn(hidden, dtype=dtype, device=DEVICE)) + + y_ref = _apply_head_perm(norm(x.clone()), num_heads, head_dim, idx) + + folded_norm = K2HorizonRMSNorm(hidden_size=hidden, n_groups=num_heads).to(DEVICE) + with torch.no_grad(): + folded_norm.weight.copy_(fold_qk_channels(norm.weight, head_dim, idx)) + # Permute x to emulate the folded q_proj/k_proj: at runtime the norm sees the + # already-permuted projection output (P @ x), not the raw activations. + y_folded = folded_norm(_apply_head_perm(x.clone(), num_heads, head_dim, idx)) + + atol, rtol = 1e-2, 1e-2 + torch.testing.assert_close(y_folded, y_ref, atol=atol, rtol=rtol) + + +class _Fp8QKVOnlyConfig(Fp8Config): + def get_quant_method(self, layer, prefix): + if isinstance(layer, LinearBase) and not ( + prefix.endswith("qkv_proj") or prefix.endswith("qk_proj") + ): + return UnquantizedLinearMethod() + return super().get_quant_method(layer, prefix) + + +class _DummyAttention(nn.Module): + def __init__(self, *args, **kwargs) -> None: + super().__init__() + + def forward(self, q, k, v): + return q + + +@pytest.fixture +def dist_init(): + with set_current_vllm_config(VllmConfig()): + init_distributed_environment( + world_size=1, + rank=0, + distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}", + local_rank=0, + backend="gloo", + ) + initialize_model_parallel(1, 1) + yield + cleanup_dist_env_and_memory() + + +def _make_hf_config( + *, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + rope_head_dim: int, + attention_bias: bool, + query_key_norm: bool, +) -> SimpleNamespace: + return SimpleNamespace( + hidden_size=hidden_size, + max_position_embeddings=4096, + dual_chunk_attention_config=None, + mlp_only_layers=[], + num_experts=0, # dense layer: avoids the MoE path + decoder_sparse_step=1, + mova_num_experts=0, + num_attention_heads=num_heads, + num_key_value_heads=num_kv_heads, + rope_parameters={"rope_theta": 10000.0}, + rope_head_dim=rope_head_dim, + query_key_norm=query_key_norm, + rms_norm_eps=1e-6, + attention_bias=attention_bias, + head_dim=head_dim, + attention_gate_func=None, + intermediate_size=2 * hidden_size, + hidden_act="silu", + layernorm_num_groups=1, + vocab_size=1000, + num_hidden_layers=1, + tie_word_embeddings=False, + ) + + +def _build_single_layer_model( + monkeypatch, hf_config: SimpleNamespace, *, quant_config=None +) -> K2HorizonModel: + monkeypatch.setattr(k2_horizon, "Attention", _DummyAttention) + + vllm_config = VllmConfig() + vllm_config.model_config = SimpleNamespace( + hf_text_config=hf_config, dtype=torch.get_default_dtype() + ) + vllm_config.quant_config = quant_config + + with set_current_vllm_config(vllm_config): + layer = K2HorizonDecoderLayer(vllm_config=vllm_config, prefix="model.layers.0") + + model = K2HorizonModel.__new__(K2HorizonModel) + nn.Module.__init__(model) + model.config = hf_config + model.quant_config = quant_config + model.num_redundant_experts = 0 + model.layers = nn.ModuleList([layer]) + return model + + +def _fold_rows(weight: torch.Tensor, head_dim: int, idx: torch.Tensor) -> torch.Tensor: + hidden = weight.shape[-1] + return weight.view(-1, head_dim, hidden)[:, idx, :].reshape(-1, hidden).contiguous() + + +def _fold_channels(vec: torch.Tensor, head_dim: int, idx: torch.Tensor) -> torch.Tensor: + return vec.view(-1, head_dim)[:, idx].reshape(-1).contiguous() + + +def _non_qk_weights(prefix, *, hidden_size, intermediate_size, q_size, device=None): + return [ + ( + f"{prefix}.self_attn.o_proj.weight", + torch.randn(hidden_size, q_size, device=device), + ), + ( + f"{prefix}.mlp.gate_proj.weight", + torch.randn(intermediate_size, hidden_size, device=device), + ), + ( + f"{prefix}.mlp.up_proj.weight", + torch.randn(intermediate_size, hidden_size, device=device), + ), + ( + f"{prefix}.mlp.down_proj.weight", + torch.randn(hidden_size, intermediate_size, device=device), + ), + ( + f"{prefix}.input_layernorm.weight", + torch.randn(hidden_size, device=device), + ), + ( + f"{prefix}.post_attention_layernorm.weight", + torch.randn(hidden_size, device=device), + ), + ] + + +@pytest.mark.usefixtures("dist_init") +@pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES) +@pytest.mark.parametrize("attention_bias", [False, True]) +def test_load_weights_dense_folds_qk( + head_dim, + rope_head_dim, + attention_bias, + monkeypatch, +): + dtype = DTYPE + set_random_seed(0) + num_heads = 8 + num_kv_heads = 2 + hidden_size = 128 + intermediate_size = 2 * hidden_size + q_size = num_heads * head_dim + kv_size = num_kv_heads * head_dim + prefix = "layers.0" + + with set_default_torch_dtype(dtype): + hf_config = _make_hf_config( + hidden_size=hidden_size, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + attention_bias=attention_bias, + query_key_norm=True, + ) + model = _build_single_layer_model(monkeypatch, hf_config, quant_config=None) + + assert model.layers[0].self_attn.fold_rope_weights is True + + raw = { + "q_proj.weight": torch.randn(q_size, hidden_size), + "k_proj.weight": torch.randn(kv_size, hidden_size), + "v_proj.weight": torch.randn(kv_size, hidden_size), + "q_norm.weight": torch.randn(q_size), + "k_norm.weight": torch.randn(kv_size), + } + if attention_bias: + raw.update( + { + "q_proj.bias": torch.randn(q_size), + "k_proj.bias": torch.randn(kv_size), + "v_proj.bias": torch.randn(kv_size), + "o_proj.bias": torch.randn(hidden_size), + } + ) + + weights: list[tuple[str, torch.Tensor]] = [ + (f"{prefix}.self_attn.{name}", tensor) for name, tensor in raw.items() + ] + weights += _non_qk_weights( + prefix, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + q_size=q_size, + ) + + loaded = model.load_weights(iter(weights)) + + params = dict(model.named_parameters()) + assert loaded == set(params) + + idx = _rope_weight_perm(head_dim, rope_head_dim) + + qkv = params[f"{prefix}.self_attn.qkv_proj.weight"] + q_part = qkv[:q_size] + k_part = qkv[q_size : q_size + kv_size] + v_part = qkv[q_size + kv_size :] + + q_norm = params[f"{prefix}.self_attn.q_norm.weight"] + k_norm = params[f"{prefix}.self_attn.k_norm.weight"] + + assert torch.equal(q_part, _fold_rows(raw["q_proj.weight"], head_dim, idx)) + assert torch.equal(k_part, _fold_rows(raw["k_proj.weight"], head_dim, idx)) + assert torch.equal(v_part, raw["v_proj.weight"]) # v never folded + assert torch.equal(q_norm, _fold_channels(raw["q_norm.weight"], head_dim, idx)) + assert torch.equal(k_norm, _fold_channels(raw["k_norm.weight"], head_dim, idx)) + assert qkv.dtype == dtype + + if attention_bias: + qkv_bias = params[f"{prefix}.self_attn.qkv_proj.bias"] + q_bias = qkv_bias[:q_size] + k_bias = qkv_bias[q_size : q_size + kv_size] + v_bias = qkv_bias[q_size + kv_size :] + assert torch.equal(q_bias, _fold_channels(raw["q_proj.bias"], head_dim, idx)) + assert torch.equal(k_bias, _fold_channels(raw["k_proj.bias"], head_dim, idx)) + assert torch.equal(v_bias, raw["v_proj.bias"]) # v bias never folded + + +@pytest.mark.skipif( + not is_quant_method_supported("fp8"), + reason="fp8 is not supported.", +) +@pytest.mark.usefixtures("dist_init") +@pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES) +@pytest.mark.parametrize("attention_bias", [False, True]) +def test_load_weights_fp8_qk_no_fold( + head_dim, + rope_head_dim, + attention_bias, + monkeypatch, +): + set_random_seed(0) + num_heads = 8 + num_kv_heads = 2 + hidden_size = 128 + intermediate_size = 2 * hidden_size + q_size = num_heads * head_dim + kv_size = num_kv_heads * head_dim + prefix = "layers.0" + + def rand_fp8(*shape): + return (torch.randn(*shape, device=DEVICE) * 32.0).to(torch.float8_e4m3fn) + + def rand_scale(): + return torch.rand(1, device=DEVICE, dtype=torch.float32) + 0.1 + + with set_default_torch_dtype(torch.bfloat16), torch.device(DEVICE): + hf_config = _make_hf_config( + hidden_size=hidden_size, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + attention_bias=attention_bias, + query_key_norm=True, + ) + quant_config = _Fp8QKVOnlyConfig( + is_checkpoint_fp8_serialized=True, activation_scheme="dynamic" + ) + model = _build_single_layer_model( + monkeypatch, hf_config, quant_config=quant_config + ) + + qkv_proj = model.layers[0].self_attn.qkv_proj + assert isinstance(qkv_proj.quant_method, Fp8LinearMethod) + assert model.layers[0].self_attn.fold_rope_weights is False + + raw_w = { + "q_proj.weight": rand_fp8(q_size, hidden_size), + "k_proj.weight": rand_fp8(kv_size, hidden_size), + "v_proj.weight": rand_fp8(kv_size, hidden_size), + } + raw_s = { + "q_proj.weight_scale": rand_scale(), + "k_proj.weight_scale": rand_scale(), + "v_proj.weight_scale": rand_scale(), + } + q_norm = torch.randn(q_size, device=DEVICE) + k_norm = torch.randn(kv_size, device=DEVICE) + + weights: list[tuple[str, torch.Tensor]] = [ + (f"{prefix}.self_attn.{name}", tensor) + for name, tensor in {**raw_w, **raw_s}.items() + ] + weights += [ + (f"{prefix}.self_attn.q_norm.weight", q_norm), + (f"{prefix}.self_attn.k_norm.weight", k_norm), + ] + biases = {} + if attention_bias: + biases = { + "q_proj.bias": torch.randn(q_size, device=DEVICE), + "k_proj.bias": torch.randn(kv_size, device=DEVICE), + "v_proj.bias": torch.randn(kv_size, device=DEVICE), + "o_proj.bias": torch.randn(hidden_size, device=DEVICE), + } + weights += [ + (f"{prefix}.self_attn.{name}", tensor) + for name, tensor in biases.items() + ] + weights += _non_qk_weights( + prefix, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + q_size=q_size, + device=DEVICE, + ) + + loaded = model.load_weights(iter(weights)) + + params = dict(model.named_parameters()) + assert loaded == set(params) + + qkv = params[f"{prefix}.self_attn.qkv_proj.weight"] + weight_scale = params[f"{prefix}.self_attn.qkv_proj.weight_scale"] + + assert qkv.dtype == torch.float8_e4m3fn + assert torch.equal(qkv[:q_size], raw_w["q_proj.weight"]) + assert torch.equal(qkv[q_size : q_size + kv_size], raw_w["k_proj.weight"]) + assert torch.equal(qkv[q_size + kv_size :], raw_w["v_proj.weight"]) + + assert torch.equal( + weight_scale, + torch.cat( + [ + raw_s["q_proj.weight_scale"], + raw_s["k_proj.weight_scale"], + raw_s["v_proj.weight_scale"], + ] + ), + ) + + assert torch.equal(params[f"{prefix}.self_attn.q_norm.weight"], q_norm) + assert torch.equal(params[f"{prefix}.self_attn.k_norm.weight"], k_norm) + + if attention_bias: + qkv_bias = params[f"{prefix}.self_attn.qkv_proj.bias"] + assert torch.equal(qkv_bias[:q_size], biases["q_proj.bias"]) + assert torch.equal(qkv_bias[q_size : q_size + kv_size], biases["k_proj.bias"]) + assert torch.equal(qkv_bias[q_size + kv_size :], biases["v_proj.bias"]) + + +def test_is_rope_weights_folding_supported(default_vllm_config): + dense = SimpleNamespace(quant_method=UnquantizedLinearMethod()) + quantized = SimpleNamespace(quant_method=object()) + + assert is_rope_weights_folding_supported(dense, None) is True + assert is_rope_weights_folding_supported(quantized, None) is False + assert is_rope_weights_folding_supported(dense, {"chunk_size": 8}) is False + assert is_rope_weights_folding_supported(SimpleNamespace(), None) is False diff --git a/vllm/model_executor/models/k2_horizon.py b/vllm/model_executor/models/k2_horizon.py index db435eaf05..8ef517ac32 100644 --- a/vllm/model_executor/models/k2_horizon.py +++ b/vllm/model_executor/models/k2_horizon.py @@ -71,6 +71,7 @@ from vllm.model_executor.layers.linear import ( QKVParallelLinear, ReplicatedLinear, RowParallelLinear, + UnquantizedLinearMethod, ) from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization import QuantizationConfig @@ -119,6 +120,96 @@ def interleaved_to_split(x): return x.reshape(*x.shape[:-1], -1, 2).transpose(-1, -2).reshape(*x.shape[:-1], -1) +def is_rope_weights_folding_supported( + qk_proj: nn.Module, + dual_chunk_attention_config: dict[str, Any] | None, +) -> bool: + """Whether the partial-RoPE weight-fold path is safe for this attention. + + Folding permutes the q/k projection's output channels at load time, which + only produces correct numerics when: + + * the q/k projection is dense/unquantized and + * dual-chunk attention is disabled -- the P.RoPE equivalence has not been + validated with dual-chunk attention. + + When either condition fails we fall back to the runtime channel-permutation + RoPE path instead of folding. + """ + is_dense_qk_proj = isinstance( + getattr(qk_proj, "quant_method", None), UnquantizedLinearMethod + ) + return is_dense_qk_proj and dual_chunk_attention_config is None + + +def apply_partial_rope( + rotary_emb, + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + num_heads: int, + num_kv_heads: int, + head_dim: int, + rope_head_dim: int, + fold_rope_weights: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + if fold_rope_weights or rope_head_dim == head_dim: + return rotary_emb(positions, q, k) + + q = q.reshape(*q.shape[:-1], num_heads, head_dim) + k = k.reshape(*k.shape[:-1], num_kv_heads, head_dim) + + q_rope, q_nope = torch.split( + split_to_interleaved(q), + split_size_or_sections=[rope_head_dim, head_dim - rope_head_dim], + dim=-1, + ) + k_rope, k_nope = torch.split( + split_to_interleaved(k), + split_size_or_sections=[rope_head_dim, head_dim - rope_head_dim], + dim=-1, + ) + + q_rope, k_rope = rotary_emb( + positions, interleaved_to_split(q_rope), interleaved_to_split(k_rope) + ) + + q = interleaved_to_split( + torch.cat([split_to_interleaved(q_rope), q_nope], dim=-1) + ).reshape(*q.shape[:-2], -1) + k = interleaved_to_split( + torch.cat([split_to_interleaved(k_rope), k_nope], dim=-1) + ).reshape(*k.shape[:-2], -1) + + return q, k + + +def _rope_weight_perm(head_dim: int, rope_head_dim: int) -> torch.Tensor: + """Per-head gather index that folds K2Horizon's partial-RoPE channel reordering. + + The K2Horizon checkpoint stores each head's channels in a GPT-J-interleaved + convention with rope/nope channels interleaved. Applying this fixed + (position-independent) permutation to the q/k weight rows lets the runtime + use vLLM's *native* NeoX partial-RoPE path (``head_size = head_dim``, + ``rotary_dim = rope_head_dim``) with no per-forward permutes/all-gather. + """ + D, R = head_dim, rope_head_dim + assert R % 2 == 0 and R <= D, ( + f"partial NeoX RoPE requires an even rope_head_dim <= head_dim, " + f"got R={R}, D={D}" + ) + h = D // 2 + idx = torch.cat( + [ + torch.arange(0, R // 2), # rope first-half channels + torch.arange(h, h + R // 2), # rope second-half channels + torch.arange(R // 2, h), # nope remainder (first half) + torch.arange(h + R // 2, D), # nope remainder (second half) + ] + ) + return idx + + def calc_router_weights( router_logits: torch.Tensor, e_score_correction_bias: torch.Tensor | None, @@ -529,12 +620,30 @@ class K2HorizonAttention(nn.Module): ) self.rope_head_dim = rope_head_dim or self.head_dim - self.rotary_emb = get_rope( - self.rope_head_dim, - max_position=max_position_embeddings, - rope_parameters=rope_parameters, - dual_chunk_attention_config=dual_chunk_attention_config, + # Fold the partial-RoPE channel permutation into the q/k weights only + # when supported (dense projections + no dual-chunk attention); + # otherwise keep the runtime channel-permutation path + # (see ``is_rope_weights_folding_supported``). + self.fold_rope_weights = is_rope_weights_folding_supported( + self.qkv_proj, dual_chunk_attention_config ) + if self.fold_rope_weights: + rope_parameters = dict(rope_parameters) + rope_parameters["rope_dim"] = self.rope_head_dim + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + is_neox_style=True, + rope_parameters=rope_parameters, + dual_chunk_attention_config=dual_chunk_attention_config, + ) + else: + self.rotary_emb = get_rope( + self.rope_head_dim, + max_position=max_position_embeddings, + rope_parameters=rope_parameters, + dual_chunk_attention_config=dual_chunk_attention_config, + ) self.attn = Attention( self.num_heads, self.head_dim, @@ -584,40 +693,17 @@ class K2HorizonAttention(nn.Module): q = self.q_norm(q) k = self.k_norm(k) - if self.rope_head_dim == self.head_dim: - q, k = self.rotary_emb(positions, q, k) - else: - q = q.reshape(*q.shape[:-1], self.num_heads, self.head_dim) - k = k.reshape(*k.shape[:-1], self.num_kv_heads, self.head_dim) - - q_rope, q_nope = torch.split( - split_to_interleaved(q), - split_size_or_sections=[ - self.rope_head_dim, - self.head_dim - self.rope_head_dim, - ], - dim=-1, - ) - k_rope, k_nope = torch.split( - split_to_interleaved(k), - split_size_or_sections=[ - self.rope_head_dim, - self.head_dim - self.rope_head_dim, - ], - dim=-1, - ) - - q_rope, k_rope = self.rotary_emb( - positions, interleaved_to_split(q_rope), interleaved_to_split(k_rope) - ) - - q = interleaved_to_split( - torch.cat([split_to_interleaved(q_rope), q_nope], dim=-1) - ).reshape(*q.shape[:-2], -1) - - k = interleaved_to_split( - torch.cat([split_to_interleaved(k_rope), k_nope], dim=-1) - ).reshape(*k.shape[:-2], -1) + q, k = apply_partial_rope( + self.rotary_emb, + positions, + q, + k, + self.num_heads, + self.num_kv_heads, + self.head_dim, + self.rope_head_dim, + self.fold_rope_weights, + ) attn_output = self.attn(q, k, v) @@ -732,12 +818,30 @@ class K2HorizonMoVAAttention(nn.Module): ) self.rope_head_dim = rope_head_dim or self.head_dim - self.rotary_emb = get_rope( - self.rope_head_dim, - max_position=max_position_embeddings, - rope_parameters=rope_parameters, - dual_chunk_attention_config=dual_chunk_attention_config, + # Fold the partial-RoPE channel permutation into the q/k weights only + # when supported (dense projections + no dual-chunk attention); + # otherwise keep the runtime channel-permutation path + # (see ``is_rope_weights_folding_supported``). + self.fold_rope_weights = is_rope_weights_folding_supported( + self.qk_proj, dual_chunk_attention_config ) + if self.fold_rope_weights: + rope_parameters = dict(rope_parameters) + rope_parameters["rope_dim"] = self.rope_head_dim + self.rotary_emb = get_rope( + self.head_dim, + max_position=max_position_embeddings, + is_neox_style=True, + rope_parameters=rope_parameters, + dual_chunk_attention_config=dual_chunk_attention_config, + ) + else: + self.rotary_emb = get_rope( + self.rope_head_dim, + max_position=max_position_embeddings, + rope_parameters=rope_parameters, + dual_chunk_attention_config=dual_chunk_attention_config, + ) self.attn = Attention( self.num_heads, self.head_dim, @@ -838,40 +942,17 @@ class K2HorizonMoVAAttention(nn.Module): q = self.q_norm(q) k = self.k_norm(k) - if self.rope_head_dim == self.head_dim: - q, k = self.rotary_emb(positions, q, k) - else: - q = q.reshape(*q.shape[:-1], self.num_heads, self.head_dim) - k = k.reshape(*k.shape[:-1], self.num_kv_heads, self.head_dim) - - q_rope, q_nope = torch.split( - split_to_interleaved(q), - split_size_or_sections=[ - self.rope_head_dim, - self.head_dim - self.rope_head_dim, - ], - dim=-1, - ) - k_rope, k_nope = torch.split( - split_to_interleaved(k), - split_size_or_sections=[ - self.rope_head_dim, - self.head_dim - self.rope_head_dim, - ], - dim=-1, - ) - - q_rope, k_rope = self.rotary_emb( - positions, interleaved_to_split(q_rope), interleaved_to_split(k_rope) - ) - - q = interleaved_to_split( - torch.cat([split_to_interleaved(q_rope), q_nope], dim=-1) - ).reshape(*q.shape[:-2], -1) - - k = interleaved_to_split( - torch.cat([split_to_interleaved(k_rope), k_nope], dim=-1) - ).reshape(*k.shape[:-2], -1) + q, k = apply_partial_rope( + self.rotary_emb, + positions, + q, + k, + self.num_heads, + self.num_kv_heads, + self.head_dim, + self.rope_head_dim, + self.fold_rope_weights, + ) attn_output = self.attn(q, k, v) @@ -1109,6 +1190,27 @@ class K2HorizonModel(nn.Module, EagleModelMixin): ("gate_up_proj", "up_proj", 1), ] + # Per-head gather index used to fold K2Horizon's partial-RoPE channel + # reordering into the q_proj/k_proj rows and the q_norm/k_norm scales, + # so the runtime can use vLLM's native partial NeoX RoPE. + head_dim = self.config.head_dim or ( + self.config.hidden_size // self.config.num_attention_heads + ) + rope_head_dim = self.config.rope_head_dim or head_dim + rope_perm_idx = _rope_weight_perm(head_dim, rope_head_dim) + + # Attention modules whose q/k projections are dense: only these fold the + # RoPE channel permutation into their q/k weights, biases and q/k norm + # scales at load time. Attentions with packed/quantized q/k projections + # keep the runtime channel-permutation path, so their weights must be + # loaded verbatim (folding would misalign per-channel scales). + rope_fold_prefixes = tuple( + f"{module_name}." + for module_name, module in self.named_modules() + if isinstance(module, (K2HorizonAttention, K2HorizonMoVAAttention)) + and getattr(module, "fold_rope_weights", False) + ) + # Skip loading extra parameters for GPTQ/modelopt models. # ignore_suffixes = ( # ".bias", @@ -1139,6 +1241,14 @@ class K2HorizonModel(nn.Module, EagleModelMixin): if name not in params_dict: continue + # Fold the per-head RoPE channel permutation into the norm scales. + if self.config.query_key_norm and name.startswith(rope_fold_prefixes): + loaded_weight = ( + loaded_weight.view(-1, head_dim)[:, rope_perm_idx] + .reshape(-1) + .contiguous() + ) + param = params_dict[name] tp_rank = get_tensor_model_parallel_rank() tp_size = get_tensor_model_parallel_world_size() @@ -1203,6 +1313,27 @@ class K2HorizonModel(nn.Module, EagleModelMixin): if ".self_attn." in name and param_name == "gate_up_proj": assert weight_name == "gate_proj" continue + + # Fold the per-head RoPE channel permutation into the dense + # q_proj / k_proj rows (W' = P W), and biases, before TP + # sharding. + if shard_id in ("q", "k") and name.startswith(rope_fold_prefixes): + if name.endswith(".weight"): + hidden = loaded_weight.shape[-1] + loaded_weight = ( + loaded_weight.view(-1, head_dim, hidden)[ + :, rope_perm_idx, : + ] + .reshape(-1, hidden) + .contiguous() + ) + elif self.config.attention_bias and name.endswith(".bias"): + loaded_weight = ( + loaded_weight.view(-1, head_dim)[:, rope_perm_idx] + .reshape(-1) + .contiguous() + ) + if ( ".self_attn." in name and param_name == "qkv_proj"