0df620d429 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -0,0 +1,195 @@

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.

4 +

5 +Constructor-crossing regression for the packed-KV zeroing stride fix (upstream

6 +PR #50276): every number below comes from instantiating the REAL production

7 +code (no arithmetic is reimplemented here):

8 +

9 + - ``vllm.v1.core.kv_cache_utils._get_packed_kv_cache_layout`` (packed

10 + offsets / block stride)

11 + - ``vllm.v1.worker.gpu.attn_utils._reshape_attention_kv_cache``

12 + (per-layer column views)

13 + - ``vllm.models.deepseek_v4.sparse_mla.DeepseekV4FlashMLABackend``

14 + (``get_kv_cache_shape`` / ``get_kv_cache_block_dim``)

15 + - ``vllm.v1.worker.utils.AttentionGroup`` / ``KVBlockZeroer``

16 + (``KVBlockZeroer.__init__`` precomputes the segment tables)

17 +

18 +CPU only: no Triton launch, no CUDA tensor, no CUDA call is made. The test

19 +asserts the packed-DSV4 constructor metadata:

20 +

21 + * block stride = the full packed row (4 aligned pages)

22 + * per-layer zero = the meaningful page (the spec's unpadded

23 + ``real_page_size_bytes``), so adjacent packed layers are

24 + never wiped

25 + * 4 packed layer views; block 99 of the highest-offset layer is the last

26 + written region and stays within the packed backing.

27 +

28 +The pre-fix constructor exposed a five-field ``_meta`` with no block stride, so

29 +its zero width equals the FULL packed row; the extraction below records that

30 +legacy defect and then asserts the fixed semantics, so failure on the old code

31 +is a semantic one (wrong zero width / out-of-bounds block 99), not a

32 +schema-shape accident.

33 +"""

34 +

35 +from types import SimpleNamespace

36 +

37 +import pytest

38 +import torch

39 +

40 +from vllm.models.deepseek_v4.sparse_mla import DeepseekV4FlashMLABackend

41 +from vllm.v1.core.kv_cache_utils import _get_packed_kv_cache_layout

42 +from vllm.v1.kv_cache_interface import MLAAttentionSpec, UniformTypeKVCacheSpecs

43 +from vllm.v1.worker.gpu.attn_utils import _reshape_attention_kv_cache

44 +from vllm.v1.worker.utils import AttentionGroup, KVBlockZeroer

45 +

46 +pytestmark = pytest.mark.cpu_test

47 +

48 +NUM_BLOCKS = 100

49 +BLOCK_SIZE_CFG = 256 # vllm_config.cache_config.block_size for DSV4 (group 0)

50 +COMPRESS = 4 # main-MLA compress_ratio

51 +STORAGE_BS = BLOCK_SIZE_CFG // COMPRESS # 64

52 +KERNEL_BS = 256 # DeepseekV4FlashMLABackend supports [256]

53 +LAST_BLOCK = NUM_BLOCKS - 1 # block 99

54 +

55 +

56 +def _extract_segments(meta, base):

57 + """Read (offset_bytes, block_stride_bytes, zero_span_bytes) per segment.

58 +

59 + Accepts both the pre-fix five-field ``_meta`` (no block stride; the zero

60 + width is the FULL packed row -- the defect) and the fixed six-field ``_meta``.

61 + All values come from the real constructor output; nothing is recomputed. Returns

62 + ``(segs, legacy_why)`` where ``legacy_why`` describes the pre-fix defect

63 + when detected.

64 + """

65 + legacy_why = None

66 + if len(meta) == 6:

67 + seg_addrs, seg_block_strides, seg_page_sizes, _max_chunks, _blk, n_segs = meta

68 + else:

69 + seg_addrs, seg_page_sizes, _max_chunks, _blk, n_segs = meta

70 + seg_block_strides = seg_page_sizes

71 + legacy_why = (

72 + f"legacy {len(meta)}-field _meta (no block stride): zero width is "

73 + f"the FULL packed row ({int(seg_page_sizes[0]) * 4} B) -- the defect"

74 + )

75 + segs = [

76 + (int(a) - base, int(bs) * 4, int(ps) * 4)

77 + for a, bs, ps in zip(seg_addrs, seg_block_strides, seg_page_sizes)

78 + ]

79 + return segs, legacy_why, n_segs

80 +

81 +

82 +def test_dsv4_packed_zeroer_geometry():

83 + """Real DSV4 constructor metadata separates block stride from zero span."""

84 + layer_names = [f"model.layers.{i}.self_attn" for i in range(4)]

85 +

86 + specs = [

87 + MLAAttentionSpec(

88 + block_size=BLOCK_SIZE_CFG,

89 + num_kv_heads=1,

90 + head_size=512,

91 + dtype=torch.uint8,

92 + compress_ratio=COMPRESS,

93 + cache_dtype_str="fp8_ds_mla",

94 + alignment=576,

95 + model_version="deepseek_v4",

96 + )

97 + for _ in layer_names

98 + ]

99 + uniform = UniformTypeKVCacheSpecs.from_specs(

100 + {ln: sp for ln, sp in zip(layer_names, specs)}

101 + )

102 + assert uniform is not None

103 +

104 + # --- Real packed layout planner -------------------------------------------

105 + block_stride, layers_by_offset = _get_packed_kv_cache_layout(

106 + [

107 + AttentionGroup(

108 + backend=DeepseekV4FlashMLABackend,

109 + layer_names=layer_names,

110 + kv_cache_spec=uniform,

111 + kv_cache_group_id=0,

112 + )

113 + ]

114 + )

115 + page_bytes = specs[0].page_size_bytes # DSV4 fp8_ds_mla page after 576B alignment

116 + real_page_bytes = specs[0].real_page_size_bytes # unpadded meaningful page

117 + alignment_gap = page_bytes - real_page_bytes # trailing planner padding

118 + total_size = block_stride * NUM_BLOCKS

119 + offsets = sorted(layers_by_offset)

120 + assert block_stride == 4 * page_bytes, (block_stride, 4 * page_bytes)

121 + assert offsets == [i * page_bytes for i in range(4)]

122 +

123 + # --- Real per-layer views (packed branch) --------------------------------

124 + backing = torch.zeros(total_size, dtype=torch.uint8)

125 + base = backing.data_ptr()

126 + views = {}

127 + for ln, off in zip(layer_names, offsets):

128 + packing = (off, block_stride)

129 + kv_shape = DeepseekV4FlashMLABackend.get_kv_cache_shape(

130 + NUM_BLOCKS, STORAGE_BS, 1, 512, cache_dtype_str="fp8_ds_mla"

131 + )

132 + stride_order = tuple(range(len(kv_shape)))

133 + view = _reshape_attention_kv_cache(

134 + backing, specs[0], kv_shape, stride_order, NUM_BLOCKS, packing

135 + )

136 + views[ln] = view

137 + assert view.data_ptr() - base == off, "data_ptr must be base+offset"

138 + assert view.stride(0) * view.element_size() == block_stride, (

139 + "row stride == block_stride (logical block stride)"

140 + )

141 +

142 + # --- Real KVBlockZeroer.__init__ (production arithmetic, CPU) ----------

143 + sctx = {ln: SimpleNamespace(kv_cache=v) for ln, v in views.items()}

144 + zeroer = KVBlockZeroer(

145 + torch.device("cpu"),

146 + attn_groups_iter=[

147 + AttentionGroup(

148 + backend=DeepseekV4FlashMLABackend,

149 + layer_names=layer_names,

150 + kv_cache_spec=specs[0],

151 + kv_cache_group_id=0,

152 + )

153 + ],

154 + kernel_block_sizes=[KERNEL_BS],

155 + cache_dtype="fp8_ds_mla",

156 + static_forward_context=sctx,

157 + )

158 + segs, legacy_why, n_segs = _extract_segments(zeroer._meta, base)

159 +

160 + # --- Fixed semantics (proven by metadata values, not by schema shape) ----

161 + assert n_segs == 4, "all 4 packed layers must register as segments"

162 + assert [off for off, _, _ in segs] == offsets, "one segment per layer offset"

163 + assert all(bs == block_stride for _, bs, _ in segs), (

164 + "block stride must be the full packed row (4 aligned pages)"

165 + )

166 + assert all(zs == real_page_bytes for _, _, zs in segs), (

167 + f"zero span must be the meaningful page ({real_page_bytes} B), got "

168 + f"{[zs for _, _, zs in segs]} B" + (f"; {legacy_why}" if legacy_why else "")

169 + )

170 +

171 + # --- Block 99 (the real final block) bounds, via constructor metadata ---

172 + # Production write for block b of a segment is

173 + # [off + b*block_stride_bytes, off + b*block_stride_bytes + zero_span_bytes).

174 + worst = 0

175 + for off, bs, zs in segs:

176 + start = off + LAST_BLOCK * bs

177 + end = start + zs

178 + oob = end - total_size

179 + worst = max(worst, oob)

180 + assert end <= total_size, (

181 + f"segment@off={off} block {LAST_BLOCK} writes [{start}, {end}) "

182 + f"-> {oob} bytes past backing ({total_size})"

183 + )

184 + assert worst == 0, f"highest-offset final block goes {worst} bytes OOB"

185 +

186 + # Highest-offset segment's block 99 is the last written region and ends exactly at

187 + # total_size - alignment_gap: the trailing planner padding stays untouched.

188 + off_last, bs_last, zs_last = segs[-1]

189 + assert off_last + LAST_BLOCK * bs_last + zs_last == total_size - alignment_gap, (

190 + "highest-offset final block must end at total_size - alignment_gap "

191 + f"({alignment_gap}-byte gap intact)"

192 + )

193 +

194 + # Zero span < block stride proves adjacent packed layers are never wiped.

195 + assert max(zs for _, _, zs in segs) < block_stride

[Test] Add packed DeepSeek-V4 KV zeroer geometry regression (#51288)