[K2 Horizon] fold partial-RoPE permutation into q/k (and norm) weights (#55335)
Signed-off-by: Alexandra Sidorova <asidorov@amd.com>
Alexandra Sidorova committed
Sep 30, 2026 at 12:03 UTC
b81984879ab28d0f7c22d21c688689613510ef92
2 files changed
+845
-78
tests/model_executor/test_k2_horizon_rope_fold.py
new
+636
@@ -0,0 +1,636 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from types import SimpleNamespace
4
+
5
+import pytest
6
+import torch
7
+from torch import nn
8
+
9
+import vllm.model_executor.models.k2_horizon as k2_horizon
10
+from tests.quantization.utils import is_quant_method_supported
11
+from vllm.config import VllmConfig, set_current_vllm_config
12
+from vllm.distributed import (
13
+ cleanup_dist_env_and_memory,
14
+ init_distributed_environment,
15
+ initialize_model_parallel,
16
+)
17
+from vllm.model_executor.layers.linear import (
18
+ LinearBase,
19
+ UnquantizedLinearMethod,
20
+)
21
+from vllm.model_executor.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod
22
+from vllm.model_executor.layers.rotary_embedding import get_rope
23
+from vllm.model_executor.models.k2_horizon import (
24
+ K2HorizonDecoderLayer,
25
+ K2HorizonModel,
26
+ K2HorizonRMSNorm,
27
+ _rope_weight_perm,
28
+ apply_partial_rope,
29
+ is_rope_weights_folding_supported,
30
+)
31
+from vllm.platforms import current_platform
32
+from vllm.utils.network_utils import get_open_port
33
+from vllm.utils.torch_utils import set_default_torch_dtype, set_random_seed
34
+
35
+DEVICE = current_platform.device_type
36
+DTYPE = torch.bfloat16
37
+SHAPES = [
38
+ (128, 64),
39
+ (128, 128),
40
+ (192, 128),
41
+ (96, 32),
42
+]
43
+
44
+
45
+def get_gptj_rope(
46
+ positions: torch.Tensor,
47
+ q: torch.Tensor,
48
+ k: torch.Tensor,
49
+ num_heads: int,
50
+ num_kv_heads: int,
51
+ head_dim: int,
52
+ rope_head_dim: int,
53
+ max_position: int,
54
+ base: float,
55
+ dtype: torch.dtype,
56
+) -> tuple[torch.Tensor, torch.Tensor]:
57
+ rotary_emb = get_rope(
58
+ rope_head_dim,
59
+ max_position=max_position,
60
+ is_neox_style=True,
61
+ rope_parameters={"rope_theta": base},
62
+ dtype=dtype,
63
+ )
64
+ return apply_partial_rope(
65
+ rotary_emb,
66
+ positions,
67
+ q,
68
+ k,
69
+ num_heads,
70
+ num_kv_heads,
71
+ head_dim,
72
+ rope_head_dim,
73
+ fold_rope_weights=False,
74
+ )
75
+
76
+
77
+def _apply_head_perm(
78
+ x: torch.Tensor, num: int, head_dim: int, idx: torch.Tensor
79
+) -> torch.Tensor:
80
+ return (
81
+ x.reshape(*x.shape[:-1], num, head_dim)[..., idx].reshape(*x.shape).contiguous()
82
+ )
83
+
84
+
85
+def get_neox_rope(
86
+ positions: torch.Tensor,
87
+ q: torch.Tensor,
88
+ k: torch.Tensor,
89
+ num_heads: int,
90
+ num_kv_heads: int,
91
+ head_dim: int,
92
+ rope_head_dim: int,
93
+ max_position: int,
94
+ base: float,
95
+ dtype: torch.dtype,
96
+) -> tuple[torch.Tensor, torch.Tensor]:
97
+ rotary_emb = get_rope(
98
+ head_dim,
99
+ max_position=max_position,
100
+ is_neox_style=True,
101
+ rope_parameters={"rope_theta": base, "rope_dim": rope_head_dim},
102
+ dtype=dtype,
103
+ )
104
+ return apply_partial_rope(
105
+ rotary_emb,
106
+ positions,
107
+ q,
108
+ k,
109
+ num_heads,
110
+ num_kv_heads,
111
+ head_dim,
112
+ rope_head_dim,
113
+ fold_rope_weights=True,
114
+ )
115
+
116
+
117
+def fold_qk_proj_weight(
118
+ weight: torch.Tensor, head_dim: int, idx: torch.Tensor
119
+) -> torch.Tensor:
120
+ hidden = weight.shape[-1]
121
+ return weight.view(-1, head_dim, hidden)[:, idx, :].reshape(-1, hidden).contiguous()
122
+
123
+
124
+def fold_qk_channels(
125
+ vec: torch.Tensor, head_dim: int, idx: torch.Tensor
126
+) -> torch.Tensor:
127
+ return vec.view(-1, head_dim)[:, idx].reshape(-1).contiguous()
128
+
129
+
130
+@pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES)
131
+def test_rope_fold_matches_permute_then_rope(
132
+ head_dim, rope_head_dim, default_vllm_config
133
+):
134
+ dtype = DTYPE
135
+ set_random_seed(0)
136
+ num_tokens = 17
137
+ num_heads = 8
138
+ num_kv_heads = 2
139
+ hidden_size = 512
140
+ max_position = 4096
141
+ base = 10000.0
142
+
143
+ torch.set_default_dtype(dtype)
144
+
145
+ positions = torch.randint(0, max_position, (num_tokens,), device=DEVICE)
146
+
147
+ scale = hidden_size**-0.5
148
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=DEVICE)
149
+ w_q = (
150
+ torch.randn(num_heads * head_dim, hidden_size, dtype=dtype, device=DEVICE)
151
+ * scale
152
+ )
153
+ w_k = (
154
+ torch.randn(num_kv_heads * head_dim, hidden_size, dtype=dtype, device=DEVICE)
155
+ * scale
156
+ )
157
+
158
+ idx = _rope_weight_perm(head_dim, rope_head_dim).to(DEVICE)
159
+
160
+ q = torch.nn.functional.linear(x, w_q)
161
+ k = torch.nn.functional.linear(x, w_k)
162
+ q_gptj, k_gptj = get_gptj_rope(
163
+ positions,
164
+ q,
165
+ k,
166
+ num_heads,
167
+ num_kv_heads,
168
+ head_dim,
169
+ rope_head_dim,
170
+ max_position,
171
+ base,
172
+ dtype,
173
+ )
174
+
175
+ q_folded = torch.nn.functional.linear(x, fold_qk_proj_weight(w_q, head_dim, idx))
176
+ k_folded = torch.nn.functional.linear(x, fold_qk_proj_weight(w_k, head_dim, idx))
177
+ q_neox, k_neox = get_neox_rope(
178
+ positions,
179
+ q_folded,
180
+ k_folded,
181
+ num_heads,
182
+ num_kv_heads,
183
+ head_dim,
184
+ rope_head_dim,
185
+ max_position,
186
+ base,
187
+ dtype,
188
+ )
189
+
190
+ q_gptj_perm = _apply_head_perm(q_gptj, num_heads, head_dim, idx)
191
+ k_gptj_perm = _apply_head_perm(k_gptj, num_kv_heads, head_dim, idx)
192
+
193
+ atol, rtol = 1e-2, 1e-2
194
+
195
+ # Two complementary checks. First, an element-wise check that the folded
196
+ # NeoX path equals the GPT-J path up to the channel permutation P
197
+ # (neox_rope(Pq, Pk) == P @ gptj_rope(q, k)), hence permuting q_gptj/k_gptj.
198
+ # Second, the scores check below: P cancels in q @ k.T.
199
+ torch.testing.assert_close(q_neox, q_gptj_perm, atol=atol, rtol=rtol)
200
+ torch.testing.assert_close(k_neox, k_gptj_perm, atol=atol, rtol=rtol)
201
+
202
+ def scores(qt, kt):
203
+ qh = qt.reshape(num_tokens, num_heads, head_dim)
204
+ kh = kt.reshape(num_tokens, num_kv_heads, head_dim)
205
+ g = num_heads // num_kv_heads
206
+ qh = qh.reshape(num_tokens, num_kv_heads, g, head_dim).float()
207
+ kh = kh.float()
208
+ return torch.einsum("tkgd,skd->tksg", qh, kh)
209
+
210
+ torch.testing.assert_close(
211
+ scores(q_neox, k_neox), scores(q_gptj, k_gptj), atol=atol, rtol=rtol
212
+ )
213
+
214
+
215
+@pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES)
216
+@pytest.mark.parametrize("num_heads", [8, 2])
217
+def test_qk_bias_fold_matches_permute(
218
+ head_dim, rope_head_dim, num_heads, default_vllm_config
219
+):
220
+ dtype = DTYPE
221
+ set_random_seed(0)
222
+ num_tokens = 17
223
+ hidden_size = 512
224
+ torch.set_default_dtype(dtype)
225
+
226
+ scale = hidden_size**-0.5
227
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=DEVICE)
228
+ w = (
229
+ torch.randn(num_heads * head_dim, hidden_size, dtype=dtype, device=DEVICE)
230
+ * scale
231
+ )
232
+ b = torch.randn(num_heads * head_dim, dtype=dtype, device=DEVICE)
233
+
234
+ idx = _rope_weight_perm(head_dim, rope_head_dim).to(DEVICE)
235
+
236
+ ref = _apply_head_perm(
237
+ torch.nn.functional.linear(x, w, b), num_heads, head_dim, idx
238
+ )
239
+ folded = torch.nn.functional.linear(
240
+ x,
241
+ fold_qk_proj_weight(w, head_dim, idx),
242
+ fold_qk_channels(b, head_dim, idx),
243
+ )
244
+
245
+ atol, rtol = 1e-2, 1e-2
246
+ torch.testing.assert_close(folded, ref, atol=atol, rtol=rtol)
247
+
248
+
249
+@pytest.mark.skipif(
250
+ not current_platform.is_cuda_alike(), reason="requires a GPU (CUDA/ROCm)"
251
+)
252
+@pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES)
253
+@pytest.mark.parametrize("num_heads", [8, 2])
254
+def test_qk_norm_scale_fold_matches_permute(
255
+ head_dim, rope_head_dim, num_heads, default_vllm_config
256
+):
257
+ dtype = DTYPE
258
+ set_random_seed(0)
259
+ num_tokens = 13
260
+ torch.set_default_dtype(dtype)
261
+
262
+ hidden = num_heads * head_dim
263
+ idx = _rope_weight_perm(head_dim, rope_head_dim).to(DEVICE)
264
+
265
+ x = torch.randn(num_tokens, hidden, dtype=dtype, device=DEVICE)
266
+
267
+ norm = K2HorizonRMSNorm(hidden_size=hidden, n_groups=num_heads).to(DEVICE)
268
+ with torch.no_grad():
269
+ norm.weight.copy_(torch.randn(hidden, dtype=dtype, device=DEVICE))
270
+
271
+ y_ref = _apply_head_perm(norm(x.clone()), num_heads, head_dim, idx)
272
+
273
+ folded_norm = K2HorizonRMSNorm(hidden_size=hidden, n_groups=num_heads).to(DEVICE)
274
+ with torch.no_grad():
275
+ folded_norm.weight.copy_(fold_qk_channels(norm.weight, head_dim, idx))
276
+ # Permute x to emulate the folded q_proj/k_proj: at runtime the norm sees the
277
+ # already-permuted projection output (P @ x), not the raw activations.
278
+ y_folded = folded_norm(_apply_head_perm(x.clone(), num_heads, head_dim, idx))
279
+
280
+ atol, rtol = 1e-2, 1e-2
281
+ torch.testing.assert_close(y_folded, y_ref, atol=atol, rtol=rtol)
282
+
283
+
284
+class _Fp8QKVOnlyConfig(Fp8Config):
285
+ def get_quant_method(self, layer, prefix):
286
+ if isinstance(layer, LinearBase) and not (
287
+ prefix.endswith("qkv_proj") or prefix.endswith("qk_proj")
288
+ ):
289
+ return UnquantizedLinearMethod()
290
+ return super().get_quant_method(layer, prefix)
291
+
292
+
293
+class _DummyAttention(nn.Module):
294
+ def __init__(self, *args, **kwargs) -> None:
295
+ super().__init__()
296
+
297
+ def forward(self, q, k, v):
298
+ return q
299
+
300
+
301
+@pytest.fixture
302
+def dist_init():
303
+ with set_current_vllm_config(VllmConfig()):
304
+ init_distributed_environment(
305
+ world_size=1,
306
+ rank=0,
307
+ distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}",
308
+ local_rank=0,
309
+ backend="gloo",
310
+ )
311
+ initialize_model_parallel(1, 1)
312
+ yield
313
+ cleanup_dist_env_and_memory()
314
+
315
+
316
+def _make_hf_config(
317
+ *,
318
+ hidden_size: int,
319
+ num_heads: int,
320
+ num_kv_heads: int,
321
+ head_dim: int,
322
+ rope_head_dim: int,
323
+ attention_bias: bool,
324
+ query_key_norm: bool,
325
+) -> SimpleNamespace:
326
+ return SimpleNamespace(
327
+ hidden_size=hidden_size,
328
+ max_position_embeddings=4096,
329
+ dual_chunk_attention_config=None,
330
+ mlp_only_layers=[],
331
+ num_experts=0, # dense layer: avoids the MoE path
332
+ decoder_sparse_step=1,
333
+ mova_num_experts=0,
334
+ num_attention_heads=num_heads,
335
+ num_key_value_heads=num_kv_heads,
336
+ rope_parameters={"rope_theta": 10000.0},
337
+ rope_head_dim=rope_head_dim,
338
+ query_key_norm=query_key_norm,
339
+ rms_norm_eps=1e-6,
340
+ attention_bias=attention_bias,
341
+ head_dim=head_dim,
342
+ attention_gate_func=None,
343
+ intermediate_size=2 * hidden_size,
344
+ hidden_act="silu",
345
+ layernorm_num_groups=1,
346
+ vocab_size=1000,
347
+ num_hidden_layers=1,
348
+ tie_word_embeddings=False,
349
+ )
350
+
351
+
352
+def _build_single_layer_model(
353
+ monkeypatch, hf_config: SimpleNamespace, *, quant_config=None
354
+) -> K2HorizonModel:
355
+ monkeypatch.setattr(k2_horizon, "Attention", _DummyAttention)
356
+
357
+ vllm_config = VllmConfig()
358
+ vllm_config.model_config = SimpleNamespace(
359
+ hf_text_config=hf_config, dtype=torch.get_default_dtype()
360
+ )
361
+ vllm_config.quant_config = quant_config
362
+
363
+ with set_current_vllm_config(vllm_config):
364
+ layer = K2HorizonDecoderLayer(vllm_config=vllm_config, prefix="model.layers.0")
365
+
366
+ model = K2HorizonModel.__new__(K2HorizonModel)
367
+ nn.Module.__init__(model)
368
+ model.config = hf_config
369
+ model.quant_config = quant_config
370
+ model.num_redundant_experts = 0
371
+ model.layers = nn.ModuleList([layer])
372
+ return model
373
+
374
+
375
+def _fold_rows(weight: torch.Tensor, head_dim: int, idx: torch.Tensor) -> torch.Tensor:
376
+ hidden = weight.shape[-1]
377
+ return weight.view(-1, head_dim, hidden)[:, idx, :].reshape(-1, hidden).contiguous()
378
+
379
+
380
+def _fold_channels(vec: torch.Tensor, head_dim: int, idx: torch.Tensor) -> torch.Tensor:
381
+ return vec.view(-1, head_dim)[:, idx].reshape(-1).contiguous()
382
+
383
+
384
+def _non_qk_weights(prefix, *, hidden_size, intermediate_size, q_size, device=None):
385
+ return [
386
+ (
387
+ f"{prefix}.self_attn.o_proj.weight",
388
+ torch.randn(hidden_size, q_size, device=device),
389
+ ),
390
+ (
391
+ f"{prefix}.mlp.gate_proj.weight",
392
+ torch.randn(intermediate_size, hidden_size, device=device),
393
+ ),
394
+ (
395
+ f"{prefix}.mlp.up_proj.weight",
396
+ torch.randn(intermediate_size, hidden_size, device=device),
397
+ ),
398
+ (
399
+ f"{prefix}.mlp.down_proj.weight",
400
+ torch.randn(hidden_size, intermediate_size, device=device),
401
+ ),
402
+ (
403
+ f"{prefix}.input_layernorm.weight",
404
+ torch.randn(hidden_size, device=device),
405
+ ),
406
+ (
407
+ f"{prefix}.post_attention_layernorm.weight",
408
+ torch.randn(hidden_size, device=device),
409
+ ),
410
+ ]
411
+
412
+
413
+@pytest.mark.usefixtures("dist_init")
414
+@pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES)
415
+@pytest.mark.parametrize("attention_bias", [False, True])
416
+def test_load_weights_dense_folds_qk(
417
+ head_dim,
418
+ rope_head_dim,
419
+ attention_bias,
420
+ monkeypatch,
421
+):
422
+ dtype = DTYPE
423
+ set_random_seed(0)
424
+ num_heads = 8
425
+ num_kv_heads = 2
426
+ hidden_size = 128
427
+ intermediate_size = 2 * hidden_size
428
+ q_size = num_heads * head_dim
429
+ kv_size = num_kv_heads * head_dim
430
+ prefix = "layers.0"
431
+
432
+ with set_default_torch_dtype(dtype):
433
+ hf_config = _make_hf_config(
434
+ hidden_size=hidden_size,
435
+ num_heads=num_heads,
436
+ num_kv_heads=num_kv_heads,
437
+ head_dim=head_dim,
438
+ rope_head_dim=rope_head_dim,
439
+ attention_bias=attention_bias,
440
+ query_key_norm=True,
441
+ )
442
+ model = _build_single_layer_model(monkeypatch, hf_config, quant_config=None)
443
+
444
+ assert model.layers[0].self_attn.fold_rope_weights is True
445
+
446
+ raw = {
447
+ "q_proj.weight": torch.randn(q_size, hidden_size),
448
+ "k_proj.weight": torch.randn(kv_size, hidden_size),
449
+ "v_proj.weight": torch.randn(kv_size, hidden_size),
450
+ "q_norm.weight": torch.randn(q_size),
451
+ "k_norm.weight": torch.randn(kv_size),
452
+ }
453
+ if attention_bias:
454
+ raw.update(
455
+ {
456
+ "q_proj.bias": torch.randn(q_size),
457
+ "k_proj.bias": torch.randn(kv_size),
458
+ "v_proj.bias": torch.randn(kv_size),
459
+ "o_proj.bias": torch.randn(hidden_size),
460
+ }
461
+ )
462
+
463
+ weights: list[tuple[str, torch.Tensor]] = [
464
+ (f"{prefix}.self_attn.{name}", tensor) for name, tensor in raw.items()
465
+ ]
466
+ weights += _non_qk_weights(
467
+ prefix,
468
+ hidden_size=hidden_size,
469
+ intermediate_size=intermediate_size,
470
+ q_size=q_size,
471
+ )
472
+
473
+ loaded = model.load_weights(iter(weights))
474
+
475
+ params = dict(model.named_parameters())
476
+ assert loaded == set(params)
477
+
478
+ idx = _rope_weight_perm(head_dim, rope_head_dim)
479
+
480
+ qkv = params[f"{prefix}.self_attn.qkv_proj.weight"]
481
+ q_part = qkv[:q_size]
482
+ k_part = qkv[q_size : q_size + kv_size]
483
+ v_part = qkv[q_size + kv_size :]
484
+
485
+ q_norm = params[f"{prefix}.self_attn.q_norm.weight"]
486
+ k_norm = params[f"{prefix}.self_attn.k_norm.weight"]
487
+
488
+ assert torch.equal(q_part, _fold_rows(raw["q_proj.weight"], head_dim, idx))
489
+ assert torch.equal(k_part, _fold_rows(raw["k_proj.weight"], head_dim, idx))
490
+ assert torch.equal(v_part, raw["v_proj.weight"]) # v never folded
491
+ assert torch.equal(q_norm, _fold_channels(raw["q_norm.weight"], head_dim, idx))
492
+ assert torch.equal(k_norm, _fold_channels(raw["k_norm.weight"], head_dim, idx))
493
+ assert qkv.dtype == dtype
494
+
495
+ if attention_bias:
496
+ qkv_bias = params[f"{prefix}.self_attn.qkv_proj.bias"]
497
+ q_bias = qkv_bias[:q_size]
498
+ k_bias = qkv_bias[q_size : q_size + kv_size]
499
+ v_bias = qkv_bias[q_size + kv_size :]
500
+ assert torch.equal(q_bias, _fold_channels(raw["q_proj.bias"], head_dim, idx))
501
+ assert torch.equal(k_bias, _fold_channels(raw["k_proj.bias"], head_dim, idx))
502
+ assert torch.equal(v_bias, raw["v_proj.bias"]) # v bias never folded
503
+
504
+
505
+@pytest.mark.skipif(
506
+ not is_quant_method_supported("fp8"),
507
+ reason="fp8 is not supported.",
508
+)
509
+@pytest.mark.usefixtures("dist_init")
510
+@pytest.mark.parametrize("head_dim,rope_head_dim", SHAPES)
511
+@pytest.mark.parametrize("attention_bias", [False, True])
512
+def test_load_weights_fp8_qk_no_fold(
513
+ head_dim,
514
+ rope_head_dim,
515
+ attention_bias,
516
+ monkeypatch,
517
+):
518
+ set_random_seed(0)
519
+ num_heads = 8
520
+ num_kv_heads = 2
521
+ hidden_size = 128
522
+ intermediate_size = 2 * hidden_size
523
+ q_size = num_heads * head_dim
524
+ kv_size = num_kv_heads * head_dim
525
+ prefix = "layers.0"
526
+
527
+ def rand_fp8(*shape):
528
+ return (torch.randn(*shape, device=DEVICE) * 32.0).to(torch.float8_e4m3fn)
529
+
530
+ def rand_scale():
531
+ return torch.rand(1, device=DEVICE, dtype=torch.float32) + 0.1
532
+
533
+ with set_default_torch_dtype(torch.bfloat16), torch.device(DEVICE):
534
+ hf_config = _make_hf_config(
535
+ hidden_size=hidden_size,
536
+ num_heads=num_heads,
537
+ num_kv_heads=num_kv_heads,
538
+ head_dim=head_dim,
539
+ rope_head_dim=rope_head_dim,
540
+ attention_bias=attention_bias,
541
+ query_key_norm=True,
542
+ )
543
+ quant_config = _Fp8QKVOnlyConfig(
544
+ is_checkpoint_fp8_serialized=True, activation_scheme="dynamic"
545
+ )
546
+ model = _build_single_layer_model(
547
+ monkeypatch, hf_config, quant_config=quant_config
548
+ )
549
+
550
+ qkv_proj = model.layers[0].self_attn.qkv_proj
551
+ assert isinstance(qkv_proj.quant_method, Fp8LinearMethod)
552
+ assert model.layers[0].self_attn.fold_rope_weights is False
553
+
554
+ raw_w = {
555
+ "q_proj.weight": rand_fp8(q_size, hidden_size),
556
+ "k_proj.weight": rand_fp8(kv_size, hidden_size),
557
+ "v_proj.weight": rand_fp8(kv_size, hidden_size),
558
+ }
559
+ raw_s = {
560
+ "q_proj.weight_scale": rand_scale(),
561
+ "k_proj.weight_scale": rand_scale(),
562
+ "v_proj.weight_scale": rand_scale(),
563
+ }
564
+ q_norm = torch.randn(q_size, device=DEVICE)
565
+ k_norm = torch.randn(kv_size, device=DEVICE)
566
+
567
+ weights: list[tuple[str, torch.Tensor]] = [
568
+ (f"{prefix}.self_attn.{name}", tensor)
569
+ for name, tensor in {**raw_w, **raw_s}.items()
570
+ ]
571
+ weights += [
572
+ (f"{prefix}.self_attn.q_norm.weight", q_norm),
573
+ (f"{prefix}.self_attn.k_norm.weight", k_norm),
574
+ ]
575
+ biases = {}
576
+ if attention_bias:
577
+ biases = {
578
+ "q_proj.bias": torch.randn(q_size, device=DEVICE),
579
+ "k_proj.bias": torch.randn(kv_size, device=DEVICE),
580
+ "v_proj.bias": torch.randn(kv_size, device=DEVICE),
581
+ "o_proj.bias": torch.randn(hidden_size, device=DEVICE),
582
+ }
583
+ weights += [
584
+ (f"{prefix}.self_attn.{name}", tensor)
585
+ for name, tensor in biases.items()
586
+ ]
587
+ weights += _non_qk_weights(
588
+ prefix,
589
+ hidden_size=hidden_size,
590
+ intermediate_size=intermediate_size,
591
+ q_size=q_size,
592
+ device=DEVICE,
593
+ )
594
+
595
+ loaded = model.load_weights(iter(weights))
596
+
597
+ params = dict(model.named_parameters())
598
+ assert loaded == set(params)
599
+
600
+ qkv = params[f"{prefix}.self_attn.qkv_proj.weight"]
601
+ weight_scale = params[f"{prefix}.self_attn.qkv_proj.weight_scale"]
602
+
603
+ assert qkv.dtype == torch.float8_e4m3fn
604
+ assert torch.equal(qkv[:q_size], raw_w["q_proj.weight"])
605
+ assert torch.equal(qkv[q_size : q_size + kv_size], raw_w["k_proj.weight"])
606
+ assert torch.equal(qkv[q_size + kv_size :], raw_w["v_proj.weight"])
607
+
608
+ assert torch.equal(
609
+ weight_scale,
610
+ torch.cat(
611
+ [
612
+ raw_s["q_proj.weight_scale"],
613
+ raw_s["k_proj.weight_scale"],
614
+ raw_s["v_proj.weight_scale"],
615
+ ]
616
+ ),
617
+ )
618
+
619
+ assert torch.equal(params[f"{prefix}.self_attn.q_norm.weight"], q_norm)
620
+ assert torch.equal(params[f"{prefix}.self_attn.k_norm.weight"], k_norm)
621
+
622
+ if attention_bias:
623
+ qkv_bias = params[f"{prefix}.self_attn.qkv_proj.bias"]
624
+ assert torch.equal(qkv_bias[:q_size], biases["q_proj.bias"])
625
+ assert torch.equal(qkv_bias[q_size : q_size + kv_size], biases["k_proj.bias"])
626
+ assert torch.equal(qkv_bias[q_size + kv_size :], biases["v_proj.bias"])
627
+
628
+
629
+def test_is_rope_weights_folding_supported(default_vllm_config):
630
+ dense = SimpleNamespace(quant_method=UnquantizedLinearMethod())
631
+ quantized = SimpleNamespace(quant_method=object())
632
+
633
+ assert is_rope_weights_folding_supported(dense, None) is True
634
+ assert is_rope_weights_folding_supported(quantized, None) is False
635
+ assert is_rope_weights_folding_supported(dense, {"chunk_size": 8}) is False
636
+ assert is_rope_weights_folding_supported(SimpleNamespace(), None) is False
vllm/model_executor/models/k2_horizon.py
+209
-78
@@ -71,6 +71,7 @@ from vllm.model_executor.layers.linear import (
71
QKVParallelLinear,
72
ReplicatedLinear,
73
RowParallelLinear,
74
+ UnquantizedLinearMethod,
75
)
76
from vllm.model_executor.layers.logits_processor import LogitsProcessor
77
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -119,6 +120,96 @@ def interleaved_to_split(x):
120
return x.reshape(*x.shape[:-1], -1, 2).transpose(-1, -2).reshape(*x.shape[:-1], -1)
121
122
123
+def is_rope_weights_folding_supported(
124
+ qk_proj: nn.Module,
125
+ dual_chunk_attention_config: dict[str, Any] | None,
126
+) -> bool:
127
+ """Whether the partial-RoPE weight-fold path is safe for this attention.
128
+
129
+ Folding permutes the q/k projection's output channels at load time, which
130
+ only produces correct numerics when:
131
+
132
+ * the q/k projection is dense/unquantized and
133
+ * dual-chunk attention is disabled -- the P.RoPE equivalence has not been
134
+ validated with dual-chunk attention.
135
+
136
+ When either condition fails we fall back to the runtime channel-permutation
137
+ RoPE path instead of folding.
138
+ """
139
+ is_dense_qk_proj = isinstance(
140
+ getattr(qk_proj, "quant_method", None), UnquantizedLinearMethod
141
+ )
142
+ return is_dense_qk_proj and dual_chunk_attention_config is None
143
+
144
+
145
+def apply_partial_rope(
146
+ rotary_emb,
147
+ positions: torch.Tensor,
148
+ q: torch.Tensor,
149
+ k: torch.Tensor,
150
+ num_heads: int,
151
+ num_kv_heads: int,
152
+ head_dim: int,
153
+ rope_head_dim: int,
154
+ fold_rope_weights: bool,
155
+) -> tuple[torch.Tensor, torch.Tensor]:
156
+ if fold_rope_weights or rope_head_dim == head_dim:
157
+ return rotary_emb(positions, q, k)
158
+
159
+ q = q.reshape(*q.shape[:-1], num_heads, head_dim)
160
+ k = k.reshape(*k.shape[:-1], num_kv_heads, head_dim)
161
+
162
+ q_rope, q_nope = torch.split(
163
+ split_to_interleaved(q),
164
+ split_size_or_sections=[rope_head_dim, head_dim - rope_head_dim],
165
+ dim=-1,
166
+ )
167
+ k_rope, k_nope = torch.split(
168
+ split_to_interleaved(k),
169
+ split_size_or_sections=[rope_head_dim, head_dim - rope_head_dim],
170
+ dim=-1,
171
+ )
172
+
173
+ q_rope, k_rope = rotary_emb(
174
+ positions, interleaved_to_split(q_rope), interleaved_to_split(k_rope)
175
+ )
176
+
177
+ q = interleaved_to_split(
178
+ torch.cat([split_to_interleaved(q_rope), q_nope], dim=-1)
179
+ ).reshape(*q.shape[:-2], -1)
180
+ k = interleaved_to_split(
181
+ torch.cat([split_to_interleaved(k_rope), k_nope], dim=-1)
182
+ ).reshape(*k.shape[:-2], -1)
183
+
184
+ return q, k
185
+
186
+
187
+def _rope_weight_perm(head_dim: int, rope_head_dim: int) -> torch.Tensor:
188
+ """Per-head gather index that folds K2Horizon's partial-RoPE channel reordering.
189
+
190
+ The K2Horizon checkpoint stores each head's channels in a GPT-J-interleaved
191
+ convention with rope/nope channels interleaved. Applying this fixed
192
+ (position-independent) permutation to the q/k weight rows lets the runtime
193
+ use vLLM's *native* NeoX partial-RoPE path (``head_size = head_dim``,
194
+ ``rotary_dim = rope_head_dim``) with no per-forward permutes/all-gather.
195
+ """
196
+ D, R = head_dim, rope_head_dim
197
+ assert R % 2 == 0 and R <= D, (
198
+ f"partial NeoX RoPE requires an even rope_head_dim <= head_dim, "
199
+ f"got R={R}, D={D}"
200
+ )
201
+ h = D // 2
202
+ idx = torch.cat(
203
+ [
204
+ torch.arange(0, R // 2), # rope first-half channels
205
+ torch.arange(h, h + R // 2), # rope second-half channels
206
+ torch.arange(R // 2, h), # nope remainder (first half)
207
+ torch.arange(h + R // 2, D), # nope remainder (second half)
208
+ ]
209
+ )
210
+ return idx
211
+
212
+
213
def calc_router_weights(
214
router_logits: torch.Tensor,
215
e_score_correction_bias: torch.Tensor | None,
@@ -529,12 +620,30 @@ class K2HorizonAttention(nn.Module):
620
)
621
622
self.rope_head_dim = rope_head_dim or self.head_dim
532
- self.rotary_emb = get_rope(
533
- self.rope_head_dim,
534
- max_position=max_position_embeddings,
535
- rope_parameters=rope_parameters,
536
- dual_chunk_attention_config=dual_chunk_attention_config,
623
+ # Fold the partial-RoPE channel permutation into the q/k weights only
624
+ # when supported (dense projections + no dual-chunk attention);
625
+ # otherwise keep the runtime channel-permutation path
626
+ # (see ``is_rope_weights_folding_supported``).
627
+ self.fold_rope_weights = is_rope_weights_folding_supported(
628
+ self.qkv_proj, dual_chunk_attention_config
629
)
630
+ if self.fold_rope_weights:
631
+ rope_parameters = dict(rope_parameters)
632
+ rope_parameters["rope_dim"] = self.rope_head_dim
633
+ self.rotary_emb = get_rope(
634
+ self.head_dim,
635
+ max_position=max_position_embeddings,
636
+ is_neox_style=True,
637
+ rope_parameters=rope_parameters,
638
+ dual_chunk_attention_config=dual_chunk_attention_config,
639
+ )
640
+ else:
641
+ self.rotary_emb = get_rope(
642
+ self.rope_head_dim,
643
+ max_position=max_position_embeddings,
644
+ rope_parameters=rope_parameters,
645
+ dual_chunk_attention_config=dual_chunk_attention_config,
646
+ )
647
self.attn = Attention(
648
self.num_heads,
649
self.head_dim,
@@ -584,40 +693,17 @@ class K2HorizonAttention(nn.Module):
693
q = self.q_norm(q)
694
k = self.k_norm(k)
695
587
- if self.rope_head_dim == self.head_dim:
588
- q, k = self.rotary_emb(positions, q, k)
589
- else:
590
- q = q.reshape(*q.shape[:-1], self.num_heads, self.head_dim)
591
- k = k.reshape(*k.shape[:-1], self.num_kv_heads, self.head_dim)
592
-
593
- q_rope, q_nope = torch.split(
594
- split_to_interleaved(q),
595
- split_size_or_sections=[
596
- self.rope_head_dim,
597
- self.head_dim - self.rope_head_dim,
598
- ],
599
- dim=-1,
600
- )
601
- k_rope, k_nope = torch.split(
602
- split_to_interleaved(k),
603
- split_size_or_sections=[
604
- self.rope_head_dim,
605
- self.head_dim - self.rope_head_dim,
606
- ],
607
- dim=-1,
608
- )
609
-
610
- q_rope, k_rope = self.rotary_emb(
611
- positions, interleaved_to_split(q_rope), interleaved_to_split(k_rope)
612
- )
613
-
614
- q = interleaved_to_split(
615
- torch.cat([split_to_interleaved(q_rope), q_nope], dim=-1)
616
- ).reshape(*q.shape[:-2], -1)
617
-
618
- k = interleaved_to_split(
619
- torch.cat([split_to_interleaved(k_rope), k_nope], dim=-1)
620
- ).reshape(*k.shape[:-2], -1)
696
+ q, k = apply_partial_rope(
697
+ self.rotary_emb,
698
+ positions,
699
+ q,
700
+ k,
701
+ self.num_heads,
702
+ self.num_kv_heads,
703
+ self.head_dim,
704
+ self.rope_head_dim,
705
+ self.fold_rope_weights,
706
+ )
707
708
attn_output = self.attn(q, k, v)
709
@@ -732,12 +818,30 @@ class K2HorizonMoVAAttention(nn.Module):
818
)
819
820
self.rope_head_dim = rope_head_dim or self.head_dim
735
- self.rotary_emb = get_rope(
736
- self.rope_head_dim,
737
- max_position=max_position_embeddings,
738
- rope_parameters=rope_parameters,
739
- dual_chunk_attention_config=dual_chunk_attention_config,
821
+ # Fold the partial-RoPE channel permutation into the q/k weights only
822
+ # when supported (dense projections + no dual-chunk attention);
823
+ # otherwise keep the runtime channel-permutation path
824
+ # (see ``is_rope_weights_folding_supported``).
825
+ self.fold_rope_weights = is_rope_weights_folding_supported(
826
+ self.qk_proj, dual_chunk_attention_config
827
)
828
+ if self.fold_rope_weights:
829
+ rope_parameters = dict(rope_parameters)
830
+ rope_parameters["rope_dim"] = self.rope_head_dim
831
+ self.rotary_emb = get_rope(
832
+ self.head_dim,
833
+ max_position=max_position_embeddings,
834
+ is_neox_style=True,
835
+ rope_parameters=rope_parameters,
836
+ dual_chunk_attention_config=dual_chunk_attention_config,
837
+ )
838
+ else:
839
+ self.rotary_emb = get_rope(
840
+ self.rope_head_dim,
841
+ max_position=max_position_embeddings,
842
+ rope_parameters=rope_parameters,
843
+ dual_chunk_attention_config=dual_chunk_attention_config,
844
+ )
845
self.attn = Attention(
846
self.num_heads,
847
self.head_dim,
@@ -838,40 +942,17 @@ class K2HorizonMoVAAttention(nn.Module):
942
q = self.q_norm(q)
943
k = self.k_norm(k)
944
841
- if self.rope_head_dim == self.head_dim:
842
- q, k = self.rotary_emb(positions, q, k)
843
- else:
844
- q = q.reshape(*q.shape[:-1], self.num_heads, self.head_dim)
845
- k = k.reshape(*k.shape[:-1], self.num_kv_heads, self.head_dim)
846
-
847
- q_rope, q_nope = torch.split(
848
- split_to_interleaved(q),
849
- split_size_or_sections=[
850
- self.rope_head_dim,
851
- self.head_dim - self.rope_head_dim,
852
- ],
853
- dim=-1,
854
- )
855
- k_rope, k_nope = torch.split(
856
- split_to_interleaved(k),
857
- split_size_or_sections=[
858
- self.rope_head_dim,
859
- self.head_dim - self.rope_head_dim,
860
- ],
861
- dim=-1,
862
- )
863
-
864
- q_rope, k_rope = self.rotary_emb(
865
- positions, interleaved_to_split(q_rope), interleaved_to_split(k_rope)
866
- )
867
-
868
- q = interleaved_to_split(
869
- torch.cat([split_to_interleaved(q_rope), q_nope], dim=-1)
870
- ).reshape(*q.shape[:-2], -1)
871
-
872
- k = interleaved_to_split(
873
- torch.cat([split_to_interleaved(k_rope), k_nope], dim=-1)
874
- ).reshape(*k.shape[:-2], -1)
945
+ q, k = apply_partial_rope(
946
+ self.rotary_emb,
947
+ positions,
948
+ q,
949
+ k,
950
+ self.num_heads,
951
+ self.num_kv_heads,
952
+ self.head_dim,
953
+ self.rope_head_dim,
954
+ self.fold_rope_weights,
955
+ )
956
957
attn_output = self.attn(q, k, v)
958
@@ -1109,6 +1190,27 @@ class K2HorizonModel(nn.Module, EagleModelMixin):
1190
("gate_up_proj", "up_proj", 1),
1191
]
1192
1193
+ # Per-head gather index used to fold K2Horizon's partial-RoPE channel
1194
+ # reordering into the q_proj/k_proj rows and the q_norm/k_norm scales,
1195
+ # so the runtime can use vLLM's native partial NeoX RoPE.
1196
+ head_dim = self.config.head_dim or (
1197
+ self.config.hidden_size // self.config.num_attention_heads
1198
+ )
1199
+ rope_head_dim = self.config.rope_head_dim or head_dim
1200
+ rope_perm_idx = _rope_weight_perm(head_dim, rope_head_dim)
1201
+
1202
+ # Attention modules whose q/k projections are dense: only these fold the
1203
+ # RoPE channel permutation into their q/k weights, biases and q/k norm
1204
+ # scales at load time. Attentions with packed/quantized q/k projections
1205
+ # keep the runtime channel-permutation path, so their weights must be
1206
+ # loaded verbatim (folding would misalign per-channel scales).
1207
+ rope_fold_prefixes = tuple(
1208
+ f"{module_name}."
1209
+ for module_name, module in self.named_modules()
1210
+ if isinstance(module, (K2HorizonAttention, K2HorizonMoVAAttention))
1211
+ and getattr(module, "fold_rope_weights", False)
1212
+ )
1213
+
1214
# Skip loading extra parameters for GPTQ/modelopt models.
1215
# ignore_suffixes = (
1216
# ".bias",
@@ -1139,6 +1241,14 @@ class K2HorizonModel(nn.Module, EagleModelMixin):
1241
if name not in params_dict:
1242
continue
1243
1244
+ # Fold the per-head RoPE channel permutation into the norm scales.
1245
+ if self.config.query_key_norm and name.startswith(rope_fold_prefixes):
1246
+ loaded_weight = (
1247
+ loaded_weight.view(-1, head_dim)[:, rope_perm_idx]
1248
+ .reshape(-1)
1249
+ .contiguous()
1250
+ )
1251
+
1252
param = params_dict[name]
1253
tp_rank = get_tensor_model_parallel_rank()
1254
tp_size = get_tensor_model_parallel_world_size()
@@ -1203,6 +1313,27 @@ class K2HorizonModel(nn.Module, EagleModelMixin):
1313
if ".self_attn." in name and param_name == "gate_up_proj":
1314
assert weight_name == "gate_proj"
1315
continue
1316
+
1317
+ # Fold the per-head RoPE channel permutation into the dense
1318
+ # q_proj / k_proj rows (W' = P W), and biases, before TP
1319
+ # sharding.
1320
+ if shard_id in ("q", "k") and name.startswith(rope_fold_prefixes):
1321
+ if name.endswith(".weight"):
1322
+ hidden = loaded_weight.shape[-1]
1323
+ loaded_weight = (
1324
+ loaded_weight.view(-1, head_dim, hidden)[
1325
+ :, rope_perm_idx, :
1326
+ ]
1327
+ .reshape(-1, hidden)
1328
+ .contiguous()
1329
+ )
1330
+ elif self.config.attention_bias and name.endswith(".bias"):
1331
+ loaded_weight = (
1332
+ loaded_weight.view(-1, head_dim)[:, rope_perm_idx]
1333
+ .reshape(-1)
1334
+ .contiguous()
1335
+ )
1336
+
1337
if (
1338
".self_attn." in name
1339
and param_name == "qkv_proj"