cf971b419a · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -3,13 +3,8 @@

3

4 import torch

5

6 -from vllm.platforms import current_platform

6 from vllm.triton_utils import tl, triton

7

9 -_IDX_BITS = tl.constexpr(20)

10 -_IDX_MASK = tl.constexpr((1 << 20) - 1)

11 -_NEG_INF = tl.constexpr(-float("inf"))

12 -

8

9 @triton.jit

10 def _max_with_nan(a, b):

@@ -143,252 +138,6 @@ def _mask_candidates_kernel(

138 )

139

140

146 -@triton.jit

147 -def _pack_key(score, index, valid):

148 - """Order-preserving int64 key: float order << _IDX_BITS | ~index.

149 -

150 - Flipping the sign bit (all bits when negative) maps IEEE order onto

151 - unsigned order. The index rides inverted in the low bits so ties resolve to

152 - the smaller index, matching torch.topk. NaN is forced above +inf because

153 - torch.topk ranks it highest, and _block_scores_kernel can produce it.

154 - Invalid lanes pack to 0, below every valid key.

155 - """

156 - bits = score.to(tl.uint32, bitcast=True)

157 - ordered = bits ^ tl.where(bits >> 31 != 0, 0xFFFFFFFF, 0x80000000)

158 - ordered = tl.where(score != score, 0xFFFFFFFF, ordered)

159 - tie = (_IDX_MASK - index).to(tl.int64)

160 - return tl.where(valid, (ordered.to(tl.int64) << _IDX_BITS) | tie, 0)

161 -

162 -

163 -@triton.jit

164 -def _scan_range_topk(scores, row, lo, hi, nblocks, K_PAD: tl.constexpr, TILE):

165 - """Running top-K of packed keys over blocks [lo, hi) of one row."""

166 - best = tl.zeros((K_PAD,), dtype=tl.int64)

167 - for start in range(lo, hi, TILE):

168 - cols = start + tl.arange(0, TILE)

169 - valid = cols < hi

170 - score = tl.load(scores + row * nblocks + cols, valid, other=_NEG_INF)

171 - # -inf marks padding; it must never take a slot.

172 - key = _pack_key(score, cols, valid & (score != _NEG_INF))

173 - best = tl.topk(tl.join(best, tl.topk(key, K_PAD)).reshape(2 * K_PAD), K_PAD)

174 - return best

175 -

176 -

177 -@triton.jit

178 -def _live_blocks(starts, ends, row, nblocks, BLOCK_SIZE, HAS_STARTS, ROW_REPEAT):

179 - """Blocks this row can populate, matching _block_scores_kernel's extent.

180 -

181 - The logits workspace is padded to the full width, so past a row's end every

182 - block scores -inf and can never win a slot. Bounding the scan here is what

183 - torch.topk cannot do: it has no per-row extent.

184 - """

185 - bound = row // ROW_REPEAT

186 - start = tl.load(starts + bound) if HAS_STARTS else 0

187 - end = tl.load(ends + bound)

188 - live = (end - start + BLOCK_SIZE - 1) // BLOCK_SIZE

189 - return tl.minimum(tl.maximum(live, 1), nblocks).to(tl.int32)

190 -

191 -

192 -@triton.jit

193 -def _active_chunks(live, NUM_CHUNKS: tl.constexpr, TILE: tl.constexpr):

194 - """Chunks worth running: each should own a full tile, else short contexts

195 - pay NUM_CHUNKS tile-wide top-ks to look at almost nothing."""

196 - return tl.minimum(tl.maximum((live + TILE - 1) // TILE, 1), NUM_CHUNKS)

197 -

198 -

199 -@triton.jit

200 -def _write_candidates(

201 - best,

202 - out,

203 - row,

204 - stride_row,

205 - stride_col,

206 - OUT_K: tl.constexpr,

207 - K_PAD: tl.constexpr,

208 - TAIL: tl.constexpr,

209 -):

210 - slots = tl.arange(0, K_PAD)

211 - index = (_IDX_MASK - (best & _IDX_MASK)).to(tl.int32)

212 - tl.store(

213 - out + row * stride_row + slots * stride_col,

214 - tl.where(best != 0, index, -1),

215 - slots < OUT_K,

216 - )

217 - if TAIL > 0:

218 - # Fewer blocks than requested candidates: pad the unused slots.

219 - rest = K_PAD + tl.arange(0, TAIL)

220 - tl.store(out + row * stride_row + rest * stride_col, -1, rest < OUT_K)

221 -

222 -

223 -@triton.jit(do_not_specialize=["nblocks"])

224 -def _row_topk_kernel(

225 - scores,

226 - starts,

227 - ends,

228 - out,

229 - nblocks,

230 - stride_row,

231 - stride_col,

232 - BLOCK_SIZE: tl.constexpr,

233 - HAS_STARTS: tl.constexpr,

234 - ROW_REPEAT: tl.constexpr,

235 - OUT_K: tl.constexpr,

236 - K_PAD: tl.constexpr,

237 - TAIL: tl.constexpr,

238 - TILE: tl.constexpr,

239 -):

240 - row = tl.program_id(0).to(tl.int64)

241 - hi = _live_blocks(starts, ends, row, nblocks, BLOCK_SIZE, HAS_STARTS, ROW_REPEAT)

242 - best = _scan_range_topk(scores, row, 0, hi, nblocks, K_PAD, TILE)

243 - _write_candidates(best, out, row, stride_row, stride_col, OUT_K, K_PAD, TAIL)

244 -

245 -

246 -@triton.jit(do_not_specialize=["nblocks"])

247 -def _row_topk_partial_kernel(

248 - scores,

249 - starts,

250 - ends,

251 - partial,

252 - nblocks,

253 - BLOCK_SIZE: tl.constexpr,

254 - HAS_STARTS: tl.constexpr,

255 - ROW_REPEAT: tl.constexpr,

256 - NUM_CHUNKS: tl.constexpr,

257 - K_PAD: tl.constexpr,

258 - TILE: tl.constexpr,

259 -):

260 - row = tl.program_id(0).to(tl.int64)

261 - chunk = tl.program_id(1)

262 - live = _live_blocks(starts, ends, row, nblocks, BLOCK_SIZE, HAS_STARTS, ROW_REPEAT)

263 - # Split on device from this row's live extent: a host-side bound would cost

264 - # a sync, and the padded width says nothing about the context.

265 - active = _active_chunks(live, NUM_CHUNKS, TILE)

266 - if chunk < active:

267 - split = (live + active - 1) // active

268 - lo = chunk * split

269 - hi = tl.minimum(lo + split, live)

270 - best = _scan_range_topk(scores, row, lo, hi, nblocks, K_PAD, TILE)

271 - slots = tl.arange(0, K_PAD)

272 - tl.store(partial + row * (NUM_CHUNKS * K_PAD) + chunk * K_PAD + slots, best)

273 -

274 -

275 -@triton.jit(do_not_specialize=["nblocks"])

276 -def _row_topk_merge_kernel(

277 - partial,

278 - starts,

279 - ends,

280 - out,

281 - nblocks,

282 - stride_row,

283 - stride_col,

284 - BLOCK_SIZE: tl.constexpr,

285 - HAS_STARTS: tl.constexpr,

286 - ROW_REPEAT: tl.constexpr,

287 - NUM_CHUNKS: tl.constexpr,

288 - OUT_K: tl.constexpr,

289 - K_PAD: tl.constexpr,

290 - TAIL: tl.constexpr,

291 - TILE: tl.constexpr,

292 -):

293 - row = tl.program_id(0).to(tl.int64)

294 - live = _live_blocks(starts, ends, row, nblocks, BLOCK_SIZE, HAS_STARTS, ROW_REPEAT)

295 - active = _active_chunks(live, NUM_CHUNKS, TILE)

296 - lanes = tl.arange(0, NUM_CHUNKS * K_PAD)

297 - # Chunks the partial pass skipped hold stale keys; mask them out.

298 - keys = tl.load(

299 - partial + row * (NUM_CHUNKS * K_PAD) + lanes,

300 - (lanes // K_PAD) < active,

301 - other=0,

302 - )

303 - _write_candidates(

304 - tl.topk(keys, K_PAD), out, row, stride_row, stride_col, OUT_K, K_PAD, TAIL

305 - )

306 -

307 -

308 -# One wave of workgroups on MI355X; below this a program-per-row leaves most

309 -# CUs idle and splitting the row pays for the extra merge.

310 -_TARGET_PROGRAMS = 256

311 -# The merge top-k spans NUM_CHUNKS*K_PAD lanes whether or not those chunks ran,

312 -# so more chunks buy parallelism at a cost the low-row cases cannot absorb.

313 -_MAX_CHUNKS = 8

314 -_TILE = 512

315 -

316 -

317 -def _rocm_select_topk(

318 - scores: torch.Tensor,

319 - row_ks: torch.Tensor | None,

320 - row_ke: torch.Tensor,

321 - nblocks: int,

322 - topk_blocks: int,

323 - out: torch.Tensor,

324 - row_repeat: int,

325 - block_size: int,

326 -) -> None:

327 - """Fused block top-k for ROCm, replacing topk + _store_candidates_kernel.

328 -

329 - torch.topk runs a multi-pass radix select that sorts values it does not

330 - need -- the caller only consumes the *set* of block ids -- and scans the

331 - padding because it cannot see per-row extents.

332 - """

333 - rows = scores.shape[0]

334 - k_pad = triton.next_power_of_2(min(topk_blocks, nblocks))

335 - # _scan_range_topk does a running top-K_PAD merge one TILE-wide slice at a

336 - # time; tl.topk silently corrupts its output if asked for more elements

337 - # than the slice holds, so the tile must never be smaller than K_PAD.

338 - tile = max(_TILE, k_pad)

339 -

340 - num_chunks = 1

341 - if rows < _TARGET_PROGRAMS:

342 - # Sized only to fill the machine; how many actually run is decided per

343 - # row on device from its live extent.

344 - want = min(

345 - triton.cdiv(_TARGET_PROGRAMS, max(rows, 1)),

346 - _MAX_CHUNKS,

347 - max(1, nblocks // tile),

348 - )

349 - num_chunks = max(1, triton.next_power_of_2(want + 1) // 2)

350 -

351 - common = dict(

352 - BLOCK_SIZE=block_size,

353 - HAS_STARTS=row_ks is not None,

354 - ROW_REPEAT=row_repeat,

355 - OUT_K=topk_blocks,

356 - K_PAD=k_pad,

357 - TAIL=triton.next_power_of_2(topk_blocks - k_pad) if topk_blocks > k_pad else 0,

358 - TILE=tile,

359 - )

360 - if num_chunks == 1:

361 - _row_topk_kernel[(rows,)](

362 - scores, row_ks, row_ke, out, nblocks, *out.stride(), **common

363 - )

364 - return

365 -

366 - partial = scores.new_empty((rows, num_chunks * k_pad), dtype=torch.int64)

367 - _row_topk_partial_kernel[(rows, num_chunks)](

368 - scores,

369 - row_ks,

370 - row_ke,

371 - partial,

372 - nblocks,

373 - BLOCK_SIZE=block_size,

374 - HAS_STARTS=row_ks is not None,

375 - ROW_REPEAT=row_repeat,

376 - NUM_CHUNKS=num_chunks,

377 - K_PAD=k_pad,

378 - TILE=tile,

379 - )

380 - _row_topk_merge_kernel[(rows,)](

381 - partial,

382 - row_ks,

383 - row_ke,

384 - out,

385 - nblocks,

386 - *out.stride(),

387 - NUM_CHUNKS=num_chunks,

388 - **common,

389 - )

390 -

391 -

141 def select_candidate_blocks(

142 logits: torch.Tensor,

143 row_ks: torch.Tensor | None,

@@ -427,11 +176,6 @@ def select_candidate_blocks(

176 row_repeat,

177 128,

178 )

430 - if current_platform.is_rocm():

431 - _rocm_select_topk(

432 - scores, row_ks, row_ke, nblocks, topk_blocks, out, row_repeat, block_size

433 - )

434 - return

179 # Keep the existing top-k tie behavior.

180 top = scores.topk(min(topk_blocks, nblocks), dim=-1)

181 _store_candidates_kernel[(rows, triton.cdiv(topk_blocks, 256))](

[Bugfix] Revert "[ROCm][Perf] Replace torch.topk in DSA candidate block selection" (#59125)