A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 14:45 UTC
[Hybrid Allocator] Better layer padding strategy for gpt-oss eagle (#29303)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Nov 24, 2025 at 14:28 UTC
71df2a57effc15b5f67cdbf55f3d1e1b71f90e86
2 files changed
+69
-1
tests/v1/core/test_kv_cache_utils.py
+59
@@ -1436,6 +1436,65 @@ def test_get_kv_cache_config_one_worker():
1436
],
1437
)
1438
1439
+ # 6 full + 5 sliding, pad to 6 full + 6 sliding. This is a typical case for gpt-oss
1440
+ # eagle where there is only one more full attention layer than sliding window layers
1441
+ kv_cache_specs_hybrid = {
1442
+ "layer_1": new_kv_cache_spec(),
1443
+ "layer_2": new_kv_cache_spec(),
1444
+ "layer_3": new_kv_cache_spec(),
1445
+ "layer_4": new_kv_cache_spec(),
1446
+ "layer_5": new_kv_cache_spec(),
1447
+ "layer_6": new_kv_cache_spec(),
1448
+ "layer_7": new_sliding_window_spec(),
1449
+ "layer_8": new_sliding_window_spec(),
1450
+ "layer_9": new_sliding_window_spec(),
1451
+ "layer_10": new_sliding_window_spec(),
1452
+ "layer_11": new_sliding_window_spec(),
1453
+ }
1454
+
1455
+ kv_cache_config_hybrid = get_kv_cache_configs(
1456
+ vllm_config, [kv_cache_specs_hybrid], [mem_per_block_per_layer * 6 * 32]
1457
+ )[0]
1458
+ print(kv_cache_config_hybrid)
1459
+ assert kv_cache_config_hybrid == KVCacheConfig(
1460
+ num_blocks=32,
1461
+ kv_cache_tensors=[
1462
+ KVCacheTensor(
1463
+ size=mem_per_block_per_layer * 32,
1464
+ shared_by=["layer_1", "layer_7"],
1465
+ ),
1466
+ KVCacheTensor(
1467
+ size=mem_per_block_per_layer * 32,
1468
+ shared_by=["layer_2", "layer_8"],
1469
+ ),
1470
+ KVCacheTensor(
1471
+ size=mem_per_block_per_layer * 32,
1472
+ shared_by=["layer_3", "layer_9"],
1473
+ ),
1474
+ KVCacheTensor(
1475
+ size=mem_per_block_per_layer * 32,
1476
+ shared_by=["layer_4", "layer_10"],
1477
+ ),
1478
+ KVCacheTensor(
1479
+ size=mem_per_block_per_layer * 32,
1480
+ shared_by=["layer_5", "layer_11"],
1481
+ ),
1482
+ KVCacheTensor(
1483
+ size=mem_per_block_per_layer * 32,
1484
+ shared_by=["layer_6"],
1485
+ ),
1486
+ ],
1487
+ kv_cache_groups=[
1488
+ KVCacheGroupSpec(
1489
+ ["layer_1", "layer_2", "layer_3", "layer_4", "layer_5", "layer_6"],
1490
+ new_kv_cache_spec(),
1491
+ ),
1492
+ KVCacheGroupSpec(
1493
+ ["layer_7", "layer_8", "layer_9", "layer_10", "layer_11"],
1494
+ new_sliding_window_spec(),
1495
+ ),
1496
+ ],
1497
+ )
1498
# different hidden size
1499
kv_cache_specs_hybrid = {
1500
"layer_1": new_kv_cache_spec(head_size=128),
vllm/v1/core/kv_cache_utils.py
+10
-1
@@ -971,7 +971,16 @@ def _get_kv_cache_groups_uniform_page_size(
971
# is the minimum number of layers among all attention types. Need a better
972
# strategy if we want to support more complex patterns (e.g., 20 full + 30
973
# sw, where the group size should be 10).
974
- group_size = min([len(layers) for layers in same_type_layers.values()])
974
+ min_num_layers = min([len(layers) for layers in same_type_layers.values()])
975
+ group_size = min_num_layers
976
+ max_num_layers = max([len(layers) for layers in same_type_layers.values()])
977
+ if max_num_layers < min_num_layers * 1.25:
978
+ # If the number of layers is not much larger than the minimum number of layers,
979
+ # use the maximum number of layers as the group size to avoid too many padding
980
+ # layers. A typical example is gpt-oss-20b + eagle, with 12 sw + 13 full. We
981
+ # pad it to (13 sw, 13 full) instead of (12 sw, 24 full). 1.25 is just a
982
+ # magic number to avoid too many padding layers.
983
+ group_size = max_num_layers
984
grouped_layers = []
985
for layers in same_type_layers.values():
986
num_padding_layers = group_size - len(layers) % group_size