@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 172 lines 6.21 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """CPU regression for the packed-DSV4 KV zeroer geometry (upstream #50276).
4
5 Four DSV4 fp8_ds_mla layers share one cross-layer allocation: each block row packs
6 the four alignment-padded pages side by side (layout BLHNC). The zeroer must step by
7 the full packed row per block while zeroing only each layer's meaningful page — never
8 the alignment padding, never an adjacent layer's page.
9
10 CPU only: the zeroer's ``__init__`` precomputes the segment tables; no kernel is
11 launched.
12 """
13
14 from types import SimpleNamespace
15
16 import pytest
17 import torch
18
19 from tests.v1.attention.utils import dense_kv_cache_views
20 from vllm.v1.kv_cache_interface import KVCacheLayout, MLAAttentionSpec
21 from vllm.v1.worker.utils import (
22 AttentionGroup,
23 KVBlockZeroer,
24 allocate_kv_cache,
25 )
26
27 pytestmark = pytest.mark.cpu_test
28
29 NUM_BLOCKS = 100
30 NUM_LAYERS = 4
31 ALIGNMENT = 576
32
33
34 def test_packed_dsv4_zeroer_zeroes_only_each_layers_page():
35 spec = MLAAttentionSpec(
36 block_size=256,
37 num_kv_heads=1,
38 head_size=512,
39 dtype=torch.uint8,
40 cache_dtype_str="fp8_ds_mla",
41 model_version="deepseek_v4",
42 tokens_per_state=4,
43 alignment=ALIGNMENT,
44 # DeepseekV4 fp8_ds_mla: 584B per token, published by the layer.
45 state_content_bytes=584,
46 )
47 # DSV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per stored state.
48 assert spec.state_content_size_bytes == 584
49 assert spec.num_states == 64
50 unpadded_page = spec.unpadded_page_size_bytes
51 padded_page = spec.page_size_bytes
52 assert unpadded_page == 64 * 584
53 assert padded_page > unpadded_page, "alignment must pad the page"
54
55 raw = torch.zeros(NUM_BLOCKS * NUM_LAYERS * padded_page, dtype=torch.int8)
56 views = dense_kv_cache_views(raw, spec, NUM_BLOCKS, NUM_LAYERS, KVCacheLayout.BLHNC)
57 base = raw.data_ptr()
58 block_row = NUM_LAYERS * padded_page
59 for i, view in enumerate(views):
60 assert view.data_ptr() - base == i * padded_page
61 assert view.stride(0) * view.element_size() == block_row
62
63 zeroer = KVBlockZeroer(
64 torch.device("cpu"),
65 attn_groups_iter=iter(
66 [
67 AttentionGroup(
68 backend=None,
69 layer_names=[f"layer.{i}" for i in range(NUM_LAYERS)],
70 kv_cache_spec=spec,
71 kv_cache_group_id=0,
72 )
73 ]
74 ),
75 kernel_block_sizes=[spec.block_size],
76 static_forward_context={
77 f"layer.{i}": SimpleNamespace(kv_cache=views[i]) for i in range(NUM_LAYERS)
78 },
79 num_blocks=NUM_BLOCKS,
80 )
81 seg_addrs, seg_block_strides, seg_page_sizes, _, _, n_segs = zeroer._meta
82
83 assert n_segs == NUM_LAYERS
84 # Segments step by the full packed row per block...
85 assert (seg_block_strides * 4 == block_row).all()
86 # ...but zero only the layer's meaningful page: no alignment padding, no
87 # adjacent layer's page.
88 assert (seg_page_sizes * 4 == unpadded_page).all()
89 assert sorted(a - base for a in seg_addrs.tolist()) == [
90 i * padded_page for i in range(NUM_LAYERS)
91 ]
92 # Block 99 of the highest-offset layer stays within the packed backing.
93 last_end = (
94 max(seg_addrs.tolist())
95 + (NUM_BLOCKS - 1) * int(seg_block_strides[0]) * 4
96 + int(seg_page_sizes[0]) * 4
97 )
98 assert last_end <= base + raw.numel()
99
100
101 def test_overlaid_zeroer_dedups_segments_with_max_span():
102 """Two groups overlay one allocation; the zeroer must emit one segment per distinct
103 byte offset, spanning the widest overlaid page, so a newly allocated block is fully
104 zeroed no matter which group owns it."""
105 from unittest.mock import MagicMock
106
107 from vllm.v1.core.kv_cache_utils import get_kv_cache_config_from_groups
108 from vllm.v1.kv_cache_interface import (
109 KVCacheGroupSpec,
110 KVCacheLayout,
111 UniformTypeKVCacheSpecs,
112 )
113
114 def make_spec(head_size):
115 return MLAAttentionSpec(
116 block_size=64, num_kv_heads=1, head_size=head_size, dtype=torch.uint8
117 )
118
119 g1_specs = {"g1.big": make_spec(512), "g1.small": make_spec(128)}
120 g2_specs = {"g2.huge": make_spec(1024)}
121 groups = [
122 KVCacheGroupSpec(
123 list(g1_specs),
124 UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=g1_specs),
125 ),
126 KVCacheGroupSpec(
127 list(g2_specs),
128 UniformTypeKVCacheSpecs(block_size=64, kv_cache_specs=g2_specs),
129 ),
130 ]
131 from vllm.config import CacheConfig
132
133 vllm_config = MagicMock()
134 vllm_config.cache_config = CacheConfig()
135 vllm_config.cache_config.num_gpu_blocks_override = None
136 vllm_config.cache_config.kv_cache_layout = "BLHNC"
137 vllm_config.attention_config.hisparse_config = None
138 config = get_kv_cache_config_from_groups(vllm_config, groups, 8 * 1024 * 1024)
139 views = allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout.BLHNC, None)
140 buf_ptr = views["g1.big"].data_ptr()
141
142 attn_groups = [
143 AttentionGroup(
144 backend=None,
145 layer_names=list(specs),
146 kv_cache_spec=next(iter(specs.values())),
147 kv_cache_group_id=gid,
148 )
149 for gid, specs in enumerate((g1_specs, g2_specs))
150 ]
151 zeroer = KVBlockZeroer(
152 torch.device("cpu"),
153 attn_groups_iter=iter(attn_groups),
154 kernel_block_sizes=[64, 64],
155 static_forward_context={
156 name: SimpleNamespace(kv_cache=views[name]) for name in views
157 },
158 num_blocks=config.num_blocks,
159 )
160 seg_addrs, seg_block_strides, seg_page_sizes, _, _, n_segs = zeroer._meta
161
162 # g1.big and g2.huge overlay at offset 0 -> one segment with g2's wider
163 # span; g1.small keeps its own segment.
164 assert n_segs == 2
165 pages = {n: s.page_size_bytes for n, s in (g1_specs | g2_specs).items()}
166 by_offset = {
167 a - buf_ptr: p * 4 for a, p in zip(seg_addrs.tolist(), seg_page_sizes.tolist())
168 }
169 assert by_offset[0] == max(pages["g1.big"], pages["g2.huge"])
170 assert by_offset[pages["g1.big"]] == pages["g1.small"]
171 packed_block_stride = max(sum(pages[n] for n in g) for g in (g1_specs, g2_specs))
172 assert (seg_block_strides * 4 == packed_block_stride).all()