@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 268 lines 7.98 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 """Simplified batch specification grammar for attention benchmarks.
5
6 Grammar (underscore-separated segments):
7 Format: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
8
9 - count: Number of identical requests (optional, default=1)
10 - q_len: Query length (number of new tokens)
11 - seq_len: Total sequence length (optional, defaults to q_len for prefill)
12 - 'k' suffix: Multiplies value by 1024
13
14 Common patterns:
15 - Prefill: q_len == seq_len (e.g., "q2k" → 2048 new tokens, 2048 seq)
16 - Decode: q_len == 1 (e.g., "q1s1k" → 1 token, 1024 seq length)
17 - Extend: q_len < seq_len (e.g., "q4s1k" → 4 tokens, 1024 seq length)
18
19 Examples:
20 q2k -> [(2048, 2048)] # Prefill: 2048 tokens
21 q1s1k -> [(1, 1024)] # Decode: 1 token, 1K sequence
22 8q1s1k -> [(1, 1024)] * 8 # 8 decode requests
23 q4s1k -> [(4, 1024)] # 4-token extend (spec decode)
24 2q1k_32q1s1k -> [(1024, 1024)] * 2 + [(1, 1024)] * 32 # Mixed batch
25 16q4s1k -> [(4, 1024)] * 16 # 16 spec decode requests
26
27 """
28
29 from collections import Counter
30 from dataclasses import dataclass
31
32 import regex as re
33
34
35 @dataclass
36 class BatchRequest:
37 """Represents a single request in a batch."""
38
39 q_len: int # Query length (number of new tokens)
40 kv_len: int # Total KV cache length
41
42 @property
43 def is_decode(self) -> bool:
44 """True if this is a decode request (q_len == 1)."""
45 return self.q_len == 1
46
47 @property
48 def is_prefill(self) -> bool:
49 """True if this is a pure prefill (q_len == kv_len)."""
50 return self.q_len == self.kv_len
51
52 @property
53 def is_extend(self) -> bool:
54 """True if this is context extension (q_len > 1, kv_len > q_len)."""
55 return self.q_len > 1 and self.kv_len > self.q_len
56
57 @property
58 def context_len(self) -> int:
59 """Context length (KV cache - query)."""
60 return self.kv_len - self.q_len
61
62 def as_tuple(self) -> tuple[int, int]:
63 """Return as (q_len, kv_len) tuple for compatibility."""
64 return (self.q_len, self.kv_len)
65
66
67 def _parse_size(size_str: str, k_suffix: str) -> int:
68 """Parse size string with optional 'k' suffix."""
69 size = int(size_str)
70 return size * 1024 if k_suffix == "k" else size
71
72
73 def parse_batch_spec(spec: str) -> list[BatchRequest]:
74 """Parse batch specification string into list of BatchRequest objects.
75
76 Grammar: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
77
78 Args:
79 spec: Batch specification string (see module docstring for grammar)
80
81 Returns:
82 List of BatchRequest objects
83
84 Raises:
85 ValueError: If spec format is invalid
86
87 """
88 requests = []
89
90 for seg in spec.split("_"):
91 # Unified pattern: (<count>?) q<q_len>(k?) (s<seq_len>(k?))?
92 m = re.match(r"^(?:(\d+))?q(\d+)(k?)(?:s(\d+)(k?))?$", seg)
93 if m:
94 cnt = int(m.group(1)) if m.group(1) else 1
95 q_len = _parse_size(m.group(2), m.group(3))
96 kv_len = _parse_size(m.group(4), m.group(5)) if m.group(4) else q_len
97 requests.extend([BatchRequest(q_len=q_len, kv_len=kv_len)] * cnt)
98 continue
99
100 raise ValueError(f"Invalid batch spec segment: '{seg}'")
101
102 return requests
103
104
105 def format_batch_spec(requests: list[BatchRequest]) -> str:
106 """Format list of BatchRequest into human-readable string.
107
108 Groups requests by type and provides counts and sizes.
109
110 Args:
111 requests: List of BatchRequest objects
112
113 Returns:
114 Formatted string describing the batch
115
116 """
117 kinds = {
118 "prefill": [],
119 "extend": [],
120 "decode": [],
121 }
122
123 for req in requests:
124 tup = (req.q_len, req.kv_len)
125 if req.is_prefill:
126 kinds["prefill"].append(tup)
127 elif req.is_extend:
128 kinds["extend"].append(tup)
129 elif req.is_decode:
130 kinds["decode"].append(tup)
131
132 parts = []
133 for kind in ["prefill", "extend", "decode"]:
134 lst = kinds[kind]
135 if not lst:
136 continue
137
138 cnt_total = len(lst)
139 ctr = Counter(lst)
140 inner = []
141
142 for (q, kv), cnt in ctr.items():
143 if kind == "prefill":
144 size = f"{q // 1024}k" if q % 1024 == 0 else str(q)
145 inner.append(f"{cnt}x{size}")
146 elif kind == "decode":
147 size = f"{kv // 1024}k" if kv % 1024 == 0 else str(kv)
148 inner.append(f"{cnt}x{size}")
149 else: # extend
150 qstr = f"{q // 1024}k" if q % 1024 == 0 else str(q)
151 kstr = f"{kv // 1024}k" if kv % 1024 == 0 else str(kv)
152 inner.append(f"{cnt}xq{qstr}kv{kstr}")
153
154 parts.append(f"{cnt_total} {kind} ({', '.join(inner)})")
155
156 return ", ".join(parts)
157
158
159 def reorder_for_flashinfer(requests: list[BatchRequest]) -> list[BatchRequest]:
160 """Reorder requests for FlashInfer: decode first, then prefill.
161
162 FlashInfer expects decode requests before prefill requests for
163 optimal performance.
164
165 Args:
166 requests: Original list of BatchRequest
167
168 Returns:
169 Reordered list with decode requests first
170
171 """
172 decodes = [r for r in requests if r.is_decode]
173 non_decodes = [r for r in requests if not r.is_decode]
174 return decodes + non_decodes
175
176
177 def split_by_type(
178 requests: list[BatchRequest],
179 ) -> dict[str, list[BatchRequest]]:
180 """Split requests by type for analysis.
181
182 Args:
183 requests: List of BatchRequest
184
185 Returns:
186 Dict with keys: 'decode', 'prefill', 'extend'
187
188 """
189 result = {
190 "decode": [],
191 "prefill": [],
192 "extend": [],
193 }
194
195 for req in requests:
196 if req.is_decode:
197 result["decode"].append(req)
198 elif req.is_prefill:
199 result["prefill"].append(req)
200 elif req.is_extend:
201 result["extend"].append(req)
202
203 return result
204
205
206 def get_batch_stats(requests: list[BatchRequest]) -> dict:
207 """Compute statistics about a batch.
208
209 Args:
210 requests: List of BatchRequest
211
212 Returns:
213 Dict with batch statistics
214
215 """
216 by_type = split_by_type(requests)
217
218 return {
219 "total_requests": len(requests),
220 "num_decode": len(by_type["decode"]),
221 "num_prefill": len(by_type["prefill"]),
222 "num_extend": len(by_type["extend"]),
223 "total_tokens": sum(r.q_len for r in requests),
224 "total_kv_cache": sum(r.kv_len for r in requests),
225 "max_q_len": max((r.q_len for r in requests), default=0),
226 "max_kv_len": max((r.kv_len for r in requests), default=0),
227 "avg_q_len": sum(r.q_len for r in requests) / len(requests) if requests else 0,
228 "avg_kv_len": (
229 sum(r.kv_len for r in requests) / len(requests) if requests else 0
230 ),
231 }
232
233
234 def get_batch_type(batch_spec: str, spec_decode_threshold: int = 8) -> str:
235 """Classify a batch spec into a type string.
236
237 Args:
238 batch_spec: Batch specification string (e.g., "q2k", "8q1s1k", "2q2k_8q1s1k")
239 spec_decode_threshold: Max q_len to be considered spec-decode vs extend
240
241 Returns:
242 Type string: "prefill", "decode", "spec-decode", "extend", or "mixed (types...)"
243
244 """
245 requests = parse_batch_spec(batch_spec)
246
247 # Classify each request
248 types_present = set()
249 for req in requests:
250 if req.is_decode:
251 types_present.add("decode")
252 elif req.is_prefill:
253 types_present.add("prefill")
254 elif req.is_extend:
255 # Distinguish spec-decode (small q_len) from extend (chunked prefill)
256 if req.q_len <= spec_decode_threshold:
257 types_present.add("spec-decode")
258 else:
259 types_present.add("extend")
260
261 if len(types_present) == 1:
262 return types_present.pop()
263 elif len(types_present) > 1:
264 # Sort for consistent output
265 sorted_types = sorted(types_present)
266 return f"mixed ({'+'.join(sorted_types)})"
267 else:
268 return "unknown"