@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 255 lines 9.94 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Indexer K production for DeepSeek V4.1 kv-source layers.
4
5 In v4.1 the index key is derived from the *main* compressor's latent:
6 ``k = k_norm(wk(latent))`` (reference model.py Indexer.forward), then RoPE'd
7 at the group's first-token position and MXFP4/FP8-quantized into the paged
8 indexer K cache. The ``wk(latent)`` GEMM runs in torch; this kernel fuses the
9 remaining k_norm → RoPE → quant → paged store, one program per token.
10
11 Unlike the legacy (v4.0) indexer path there is no per-token pooling from a
12 compressor state cache: the latent already stands for a whole group, so only
13 group-boundary tokens ``(position + 1) % compress_ratio == 0`` produce a key.
14 """
15
16 import torch
17
18 from vllm.platforms import current_platform
19 from vllm.triton_utils import tl, triton
20
21 from . import MXFP4_BLOCK_SIZE, _fp32x2_to_fp4x2
22
23 # ROCm tiled ("SHUFFLE") indexer K value layout, mirroring
24 # indexer_k_quant_and_cache_triton's defaults: 16 positions x 16 bytes.
25 _BLOCK_TILE_SIZE = 16
26 _HEAD_TILE_SIZE = 16
27
28
29 def indexer_k_norm_rope_store(
30 k_pre: torch.Tensor,
31 positions: torch.Tensor,
32 cos_sin_cache: torch.Tensor,
33 rms_norm_weight: torch.Tensor,
34 rms_norm_eps: float,
35 k_cache: torch.Tensor,
36 kv_slot_mapping: torch.Tensor,
37 compress_ratio: int,
38 use_fp4_cache: bool,
39 ) -> None:
40 """k_norm → RoPE → quant → paged store for indexer keys.
41
42 Args:
43 k_pre: [num_tokens, 128] bf16, the ``wk(latent)`` projection. Only
44 group-boundary rows are read.
45 positions: [num_tokens] int64 token positions.
46 cos_sin_cache: [max_pos, rope_head_dim] GPT-J layout (cos half, then
47 sin half), from the layer's compress-RoPE instance.
48 rms_norm_weight: [128] k_norm weight.
49 rms_norm_eps: Epsilon of the k_norm RMSNorm.
50 k_cache: uint8 paged indexer cache [num_blocks, block_size, row_bytes].
51 kv_slot_mapping: [num_tokens] slots in the indexer cache (-1 = skip).
52 compress_ratio: group size; keys are emitted at group boundaries.
53 use_fp4_cache: MXFP4 (2 nibbles/byte + ue8m0 per 32) when True, else
54 per-token FP8 with a single fp32 scale.
55
56 """
57 num_tokens = kv_slot_mapping.numel()
58 assert k_pre.ndim == 2 and k_pre.shape[1] == 128
59 assert k_pre.dtype == torch.bfloat16 and k_pre.stride(1) == 1
60 assert num_tokens <= k_pre.shape[0] and num_tokens <= positions.numel()
61 assert compress_ratio in (1, 2)
62 if num_tokens == 0:
63 return
64
65 head_dim = k_pre.shape[1]
66 if use_fp4_cache:
67 token_stride = head_dim // 2
68 scale_dim = head_dim // MXFP4_BLOCK_SIZE
69 else:
70 token_stride = head_dim
71 scale_dim = 4 # single float32 scale
72
73 # ROCm reads this cache back with the 16x16-tiled ("SHUFFLE") value
74 # layout: cp_gather_indexer_k_quant_cache_triton on prefill and aiter's
75 # deepgemm_fp8_paged_mqa_logits(Preshuffle=True) on decode both select it
76 # from ``block_size > 1``, matching the v4.0 writer
77 # (indexer_k_quant_and_cache_triton). Writing row-major here would hand
78 # both readers permuted key bytes.
79 block_size = k_cache.shape[1]
80 shuffle = current_platform.is_rocm() and block_size > 1
81 if shuffle:
82 if use_fp4_cache:
83 raise NotImplementedError(
84 "MXFP4 indexer K cache has no tiled ROCm layout; "
85 "the ROCm readers only implement the FP8 one."
86 )
87 if block_size % _BLOCK_TILE_SIZE != 0 or head_dim % _HEAD_TILE_SIZE != 0:
88 raise ValueError(
89 f"ROCm tiled indexer K cache needs block_size "
90 f"({block_size}) % {_BLOCK_TILE_SIZE} == 0 and head_dim "
91 f"({head_dim}) % {_HEAD_TILE_SIZE} == 0."
92 )
93
94 launch_kwargs = {"launch_pdl": False} if current_platform.is_cuda() else {}
95 _indexer_k_norm_rope_quant_store_kernel[(num_tokens,)](
96 k_pre,
97 k_pre.stride(0),
98 positions,
99 rms_norm_weight,
100 rms_norm_eps,
101 cos_sin_cache,
102 cos_sin_cache.stride(0),
103 k_cache,
104 kv_slot_mapping,
105 k_cache.shape[1],
106 HEAD_SIZE=head_dim,
107 ROPE_HEAD_DIM=64,
108 COMPRESS_RATIO=compress_ratio,
109 TOKEN_STRIDE=token_stride,
110 SCALE_DIM=scale_dim,
111 KV_BLOCK_STRIDE=k_cache.stride(0),
112 FP8_MAX=448.0,
113 USE_FP4=use_fp4_cache,
114 SHUFFLE=shuffle,
115 BLOCK_TILE_SIZE=_BLOCK_TILE_SIZE,
116 HEAD_TILE_SIZE=_HEAD_TILE_SIZE,
117 num_warps=1,
118 **launch_kwargs,
119 )
120
121
122 @triton.jit
123 def _indexer_k_norm_rope_quant_store_kernel(
124 k_pre_ptr,
125 k_pre_stride,
126 positions_ptr,
127 rms_norm_weight_ptr,
128 rms_norm_eps,
129 cos_sin_cache_ptr,
130 cos_sin_stride,
131 k_cache_ptr,
132 kv_slot_mapping_ptr,
133 kv_cache_block_size,
134 HEAD_SIZE: tl.constexpr,
135 ROPE_HEAD_DIM: tl.constexpr,
136 COMPRESS_RATIO: tl.constexpr,
137 TOKEN_STRIDE: tl.constexpr,
138 SCALE_DIM: tl.constexpr,
139 KV_BLOCK_STRIDE: tl.constexpr,
140 FP8_MAX: tl.constexpr,
141 USE_FP4: tl.constexpr,
142 SHUFFLE: tl.constexpr,
143 BLOCK_TILE_SIZE: tl.constexpr,
144 HEAD_TILE_SIZE: tl.constexpr,
145 ):
146 token_idx = tl.program_id(0)
147
148 kv_slot_idx = tl.load(kv_slot_mapping_ptr + token_idx)
149 if kv_slot_idx < 0:
150 return
151 position = tl.load(positions_ptr + token_idx)
152 # Only the last token of a group publishes that group's index key.
153 if (position + 1) % COMPRESS_RATIO != 0:
154 return
155
156 block = tl.arange(0, HEAD_SIZE)
157
158 # ── k_norm (fp32 throughout, bf16 roundtrip like the reference) ────
159 k = tl.load(k_pre_ptr + token_idx * k_pre_stride + block).to(tl.float32)
160 rms_w = tl.load(rms_norm_weight_ptr + block).to(tl.float32)
161 variance = tl.sum(k * k, axis=0) / HEAD_SIZE
162 k = (k * tl.rsqrt(variance + rms_norm_eps) * rms_w).to(tl.bfloat16)
163 k = k.to(tl.float32)
164
165 # ── Register-based GPT-J forward RoPE in fp32 ─────────────────────
166 # A latent stands for the first token of its group, so group j takes
167 # position j * compress_ratio.
168 NUM_PAIRS: tl.constexpr = HEAD_SIZE // 2
169 NOPE_HEAD_DIM: tl.constexpr = HEAD_SIZE - ROPE_HEAD_DIM
170 NOPE_PAIRS: tl.constexpr = NOPE_HEAD_DIM // 2
171 HALF_ROPE: tl.constexpr = ROPE_HEAD_DIM // 2
172
173 even, odd = tl.split(tl.reshape(k, (NUM_PAIRS, 2))) # each [NUM_PAIRS]
174 pair_idx = tl.arange(0, NUM_PAIRS)
175 rope_pair_local = pair_idx - NOPE_PAIRS
176 is_rope_pair = rope_pair_local >= 0
177 cs_idx = tl.maximum(rope_pair_local, 0)
178
179 compressed_pos = (position // COMPRESS_RATIO) * COMPRESS_RATIO
180 cache_base = cos_sin_cache_ptr + compressed_pos * cos_sin_stride
181 cos_v = tl.load(cache_base + cs_idx, mask=is_rope_pair, other=1.0)
182 sin_v = tl.load(cache_base + HALF_ROPE + cs_idx, mask=is_rope_pair, other=0.0)
183
184 new_even = even * cos_v - odd * sin_v
185 new_odd = odd * cos_v + even * sin_v
186
187 # bf16 roundtrip for parity with the reference / Q-side kernel numerics.
188 new_even = new_even.to(tl.bfloat16).to(tl.float32)
189 new_odd = new_odd.to(tl.bfloat16).to(tl.float32)
190
191 # ── Paged cache pointers (segregated: values first, then scales) ──
192 kv_block_idx = kv_slot_idx // kv_cache_block_size
193 kv_pos_in_block = kv_slot_idx % kv_cache_block_size
194 cache_block_ptr = k_cache_ptr + kv_block_idx.to(tl.int64) * KV_BLOCK_STRIDE
195 if SHUFFLE:
196 # [block_size // BT, HEAD_SIZE // HT, BT, HT] value tiles.
197 val_ptr = (
198 cache_block_ptr
199 + (kv_pos_in_block // BLOCK_TILE_SIZE) * BLOCK_TILE_SIZE * TOKEN_STRIDE
200 + (kv_pos_in_block % BLOCK_TILE_SIZE) * HEAD_TILE_SIZE
201 )
202 else:
203 val_ptr = cache_block_ptr + kv_pos_in_block * TOKEN_STRIDE
204 scale_ptr = (
205 cache_block_ptr
206 + kv_cache_block_size * TOKEN_STRIDE
207 + kv_pos_in_block * SCALE_DIM
208 )
209
210 if USE_FP4:
211 # MXFP4: each 32-element block = 16 consecutive even/odd pairs, so
212 # tiling the halves into (N_BLOCKS, 16) lands one block per row.
213 N_QUANT_BLOCKS: tl.constexpr = HEAD_SIZE // 32
214 HALF_BLOCK: tl.constexpr = 16
215 even_2d = tl.reshape(new_even, (N_QUANT_BLOCKS, HALF_BLOCK))
216 odd_2d = tl.reshape(new_odd, (N_QUANT_BLOCKS, HALF_BLOCK))
217
218 amax = tl.maximum(
219 tl.max(tl.abs(even_2d), axis=1),
220 tl.max(tl.abs(odd_2d), axis=1),
221 )
222 amax = tl.maximum(amax, 6.0 * (2**-126))
223 # ue8m0 block scale: 2^ceil(log2(amax / 6.0)), stored (exp + 127).
224 log2_ratio = tl.ceil(tl.log2(amax * (1.0 / 6.0)))
225 log2_ratio = tl.minimum(tl.maximum(log2_ratio, -127.0), 127.0)
226 inv_scale = tl.exp2(-log2_ratio)
227 ue8m0 = (log2_ratio + 127.0).to(tl.uint8) # [N_QUANT_BLOCKS]
228
229 inv_scale_col = tl.reshape(inv_scale, (N_QUANT_BLOCKS, 1))
230 packed = _fp32x2_to_fp4x2(
231 even_2d * inv_scale_col, odd_2d * inv_scale_col
232 ) # (N_BLOCKS, HALF_BLOCK) uint8
233 packed_flat = tl.reshape(packed, (TOKEN_STRIDE,))
234
235 tl.store(val_ptr + tl.arange(0, TOKEN_STRIDE), packed_flat)
236 tl.store(scale_ptr + tl.arange(0, SCALE_DIM), ue8m0)
237 else:
238 # Per-token FP8 (single 128-wide block) with one float32 scale.
239 result = tl.interleave(new_even, new_odd) # [HEAD_SIZE] fp32
240 result_bf16 = result.to(tl.bfloat16).to(tl.float32)
241 INV_FP8_MAX: tl.constexpr = 1.0 / FP8_MAX
242 absmax = tl.maximum(tl.max(tl.abs(result_bf16), axis=0), 1e-4)
243 exponent = tl.ceil(tl.log2(absmax * INV_FP8_MAX))
244 inv_scale = tl.exp2(-exponent)
245 x_clamped = tl.clamp(result_bf16 * inv_scale, -FP8_MAX, FP8_MAX)
246 x_uint8 = x_clamped.to(tl.float8e4nv).to(tl.uint8, bitcast=True)
247 if SHUFFLE:
248 tiled = (
249 block // HEAD_TILE_SIZE * BLOCK_TILE_SIZE * HEAD_TILE_SIZE
250 + block % HEAD_TILE_SIZE
251 )
252 tl.store(val_ptr + tiled, x_uint8)
253 else:
254 tl.store(val_ptr + block, x_uint8)
255 tl.store(scale_ptr.to(tl.pointer_type(tl.float32)), tl.exp2(exponent))