@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 272 lines 9.35 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """FlashMLA mega attention vs the split-KV decode pipeline (DeepSeek V4.1).
4
5 Times the full set of per-layer ops each path runs in a decode step, captured
6 in a CUDA graph -- eager timing is dominated by launch overhead and flatters
7 whichever path launches fewer kernels.
8
9 mega: fused_qnorm_rope_kv_rope_quant_insert (Q pad + KV insert)
10 + mega kernel
11 split-KV: fused_qnorm_rope_kv_rope_quant_insert (Q RoPE + pad + KV insert)
12 + flash_mla_with_kvcache + fused_inv_rope_fp8_quant
13
14 Both arms fuse their Q preparation into the KV insert, so that op is charged
15 to both; the mega arm's Q side is only a zero-pad (nothing at all when the
16 shard is already at the kernel head count) because its kernel does the Q RoPE,
17 the inverse RoPE and the FP8 cast.
18
19 ``--local-heads`` is the variable that decides the outcome: the mega kernel
20 absorbs work proportional to the live head count while paying for the padded
21 one, so it wins at 64 live heads and is a wash at 16 (TP4 on a 64-head model).
22
23 Both arms use the same KV records so the comparison is kernel-for-kernel:
24 a V4.1 fp8 (528 B) sliding-window cache and a compressed cache that is either
25 V4.1 fp8 or V4.1 NVFP4 (288 B) -- FlashMLA's SM100 sparse decode reads the
26 NVFP4 compressed record in both the fused and unfused kernels.
27
28 Run: .venv/bin/python benchmarks/kernels/benchmark_dsv41_mega_attn.py
29 """
30
31 import argparse
32
33 import torch
34
35 import vllm.v1.attention.ops.flashmla as fm
36 from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
37 fused_inv_rope_fp8_quant,
38 )
39 from vllm.models.deepseek_v41.common.ops import quantize_and_insert_k_cache
40 from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import (
41 rope_quant_insert,
42 )
43 from vllm.models.deepseek_v41.common.ops.fused_layout import permute_q_to_fused
44 from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import (
45 alloc_mega_attn_output,
46 is_flashmla_mega_attn_supported,
47 )
48 from vllm.utils.math_utils import round_up
49
50 HEAD_DIM, ROPE_DIM = 512, 64
51 V41_BYTES, V41_FP4_BYTES = 528, 288
52 SWA_BLOCK, COMPRESSED_BLOCK = 32, 128
53
54
55 def make_cos_sin_cache(max_pos, device):
56 inv_freq = 1.0 / (
57 10000 ** (torch.arange(0, ROPE_DIM, 2, device=device).float() / ROPE_DIM)
58 )
59 freqs = torch.outer(torch.arange(max_pos, device=device).float(), inv_freq)
60 return torch.cat([freqs.cos(), freqs.sin()], -1)
61
62
63 def empty_paged_cache(num_rows, block_size, bytes_per_token, device):
64 """A zeroed paged cache view [num_blocks, block_size, bytes_per_token]."""
65 num_blocks = (num_rows + block_size - 1) // block_size + 1
66 page = round_up(block_size * bytes_per_token, 512)
67 backing = torch.zeros(num_blocks, page, dtype=torch.uint8, device=device)
68 return backing, backing.as_strided(
69 (num_blocks, block_size, bytes_per_token), (page, bytes_per_token, 1)
70 )
71
72
73 def fill_cache(view, bytes_per_token, cos_sin, device):
74 n = view.shape[0] * view.shape[1]
75 rows = torch.randn(n, HEAD_DIM, device=device, dtype=torch.bfloat16)
76 slots = torch.arange(n, dtype=torch.int64, device=device)
77 if bytes_per_token == V41_FP4_BYTES:
78 # A page-aligned cache has more slots than the cos_sin table has rows;
79 # wrap so the insert never indexes past it.
80 positions = slots % cos_sin.shape[0]
81 rope_quant_insert(rows, positions, cos_sin, view, slots, 1)
82 else:
83 quantize_and_insert_k_cache(
84 rows,
85 view.reshape(view.shape[0], -1),
86 slots,
87 block_size=view.shape[1],
88 bytes_per_token=bytes_per_token,
89 )
90
91
92 def time_graph(fn, iters=100, warmup=5):
93 """Capture ``fn`` in a CUDA graph and time its replay."""
94 s = torch.cuda.Stream()
95 s.wait_stream(torch.cuda.current_stream())
96 with torch.cuda.stream(s):
97 for _ in range(3):
98 fn()
99 torch.cuda.current_stream().wait_stream(s)
100 g = torch.cuda.CUDAGraph()
101 with torch.cuda.graph(g):
102 fn()
103 for _ in range(warmup):
104 g.replay()
105 torch.accelerator.synchronize()
106 start = torch.cuda.Event(enable_timing=True)
107 end = torch.cuda.Event(enable_timing=True)
108 start.record()
109 for _ in range(iters):
110 g.replay()
111 end.record()
112 torch.accelerator.synchronize()
113 return start.elapsed_time(end) * 1000 / iters
114
115
116 def bench_decode(s_q, device, topk_extra, extra_bytes, local_heads):
117 padded_heads, scale = 64, HEAD_DIM**-0.5
118 n_wv_group = padded_heads // 8
119 cos_sin = make_cos_sin_cache(65536, device)
120 positions = torch.randint(0, 65536, (s_q,), device=device, dtype=torch.int64)
121 pos32 = positions.to(torch.int32)
122
123 # Per-step inputs: the wq_b output (local heads) and this step's KV row.
124 q_local = torch.randn(
125 s_q, local_heads, HEAD_DIM, device=device, dtype=torch.bfloat16
126 )
127 q_fused = permute_q_to_fused(q_local)
128 kv = torch.randn(s_q, HEAD_DIM, device=device, dtype=torch.bfloat16)
129
130 n_swa = s_q * 128 + 128
131 swa_backing, swa = empty_paged_cache(n_swa, SWA_BLOCK, V41_BYTES, device)
132 fill_cache(swa, V41_BYTES, cos_sin, device)
133 _, extra = empty_paged_cache(65536, COMPRESSED_BLOCK, extra_bytes, device)
134 fill_cache(extra, extra_bytes, cos_sin, device)
135
136 slot_mapping = torch.arange(s_q, dtype=torch.int64, device=device)
137 swa_idx = torch.randint(0, n_swa, (s_q, 128), device=device, dtype=torch.int32)
138 swa_len = torch.full((s_q,), 128, device=device, dtype=torch.int32)
139 ex_idx = torch.randint(
140 0, 65536, (s_q, topk_extra), device=device, dtype=torch.int32
141 )
142 ex_len = torch.full((s_q,), topk_extra, device=device, dtype=torch.int32)
143 sink = torch.zeros(padded_heads, device=device)
144 sched = fm.FlashMLASchedMeta()
145 out = alloc_mega_attn_output(s_q, n_wv_group, device)
146 swa4 = swa.unsqueeze(-2)
147 extra4 = extra.unsqueeze(-2)
148 swa_2d = swa_backing
149
150 def mega():
151 # One launch zero-pads the fused-layout Q and inserts the SWA KV.
152 q_pad = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
153 q_fused,
154 kv,
155 swa_2d,
156 slot_mapping,
157 positions,
158 cos_sin,
159 0 if padded_heads == local_heads else padded_heads,
160 1e-6,
161 SWA_BLOCK,
162 False, # apply_q_norm
163 True, # kv_mxfp8
164 False, # apply_q_rope: the mega kernel rotates Q itself
165 True, # is_q_interleaved
166 )
167 torch.ops._flashmla_C.fused_norm_rope_attn_rope_cast_decode(
168 q_fused if padded_heads == local_heads else q_pad,
169 swa4,
170 swa_idx,
171 scale,
172 HEAD_DIM,
173 sink,
174 swa_len,
175 extra4,
176 ex_idx,
177 ex_len,
178 False,
179 0.0,
180 pos32,
181 False,
182 64,
183 cos_sin,
184 n_wv_group,
185 32,
186 True,
187 True,
188 True,
189 out.data,
190 out.scale,
191 )
192
193 def split_kv():
194 q_pad = torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
195 q_local,
196 kv,
197 swa_2d,
198 slot_mapping,
199 positions,
200 cos_sin,
201 padded_heads,
202 1e-6,
203 SWA_BLOCK,
204 False,
205 True,
206 )
207 o = fm.flash_mla_with_kvcache(
208 q=q_pad.unsqueeze(1),
209 k_cache=swa4,
210 block_table=None,
211 head_dim_v=HEAD_DIM,
212 tile_scheduler_metadata=sched,
213 cache_seqlens=None,
214 is_fp8_kvcache=True,
215 indices=swa_idx.view(s_q, 1, -1),
216 topk_length=swa_len,
217 softmax_scale=scale,
218 attn_sink=sink,
219 extra_k_cache=extra4,
220 extra_indices_in_kvcache=ex_idx.view(s_q, 1, -1),
221 extra_topk_length=ex_len,
222 )[0]
223 fused_inv_rope_fp8_quant(
224 o.squeeze(1)[:, :local_heads],
225 positions,
226 cos_sin,
227 n_groups=local_heads // 8,
228 heads_per_group=8,
229 quant_group_size=32,
230 tma_aligned_scales=True,
231 )
232
233 return time_graph(mega), time_graph(split_kv)
234
235
236 def main():
237 parser = argparse.ArgumentParser()
238 parser.add_argument(
239 "--decode-s-q",
240 nargs="+",
241 type=int,
242 default=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512],
243 )
244 parser.add_argument("--topk-extra", type=int, default=512)
245 parser.add_argument("--local-heads", type=int, default=16)
246 parser.add_argument(
247 "--extra-bytes",
248 type=int,
249 choices=[V41_BYTES, V41_FP4_BYTES],
250 default=V41_BYTES,
251 help="compressed-cache record: 528 (V4.1 fp8) or 288 (V4.1 NVFP4)",
252 )
253 args = parser.parse_args()
254 ok, reason = is_flashmla_mega_attn_supported()
255 if not ok:
256 raise SystemExit(reason)
257 device = torch.device("cuda")
258 torch.manual_seed(0)
259 print(
260 f"decode, cudagraph (topk_swa=128, topk_extra={args.topk_extra}, "
261 f"local_heads={args.local_heads}, extra={args.extra_bytes}B)"
262 )
263 print(f"{'s_q':>6} {'mega_us':>9} {'split_kv_us':>12} {'speedup':>8}")
264 for s_q in args.decode_s_q:
265 m, s = bench_decode(
266 s_q, device, args.topk_extra, args.extra_bytes, args.local_heads
267 )
268 print(f"{s_q:>6d} {m:9.1f} {s:12.1f} {s / m:7.2f}x")
269
270
271 if __name__ == "__main__":
272 main()