6a1570711c · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -220,6 +220,241 @@ def _patch_make_bitmatrix_metadata() -> None:

220 _bm.make_bitmatrix_metadata = _make_bitmatrix_metadata_pow2_safe

221

222

223 +def _patch_legacy_routing_for_nonpow2_topk() -> None:

224 + """Monkey-patch the legacy (v3.5.1) triton_kernels routing path to support

225 + non-power-of-2 top_k (e.g. DeepSeek-V4 top_k=6).

226 +

227 + The bundled ``_routing_compute_indx`` does ``tl.arange(0, N_EXPTS_ACT *

228 + BLOCK_M)``, which fails to compile when ``N_EXPTS_ACT`` (top_k) is not a

229 + power of 2 (6 * 32 = 192). This installs a pow2-safe variant that pads the

230 + ``tl.arange`` to the next power of 2, strides by the real per-block size,

231 + and masks the padded tail so it neither loads the next block's gates nor

232 + writes any output. For power-of-2 top_k it is identical to the original.

233 +

234 + A matching ``sort_tokens`` is installed that threads the padded size into

235 + the patched kernel. Only needed on the legacy path; the v3.6+ SparseMatrix

236 + path is handled by ``_patch_make_bitmatrix_metadata``.

237 + """

238 + import triton

239 + import triton.language as tl

240 +

241 + # Import via the `triton_kernels` alias (set up by has_triton_kernels) so

242 + # we patch the SAME module object that `make_routing_data` consumes. The

243 + # `vllm.third_party.triton_kernels.routing` path is a *different* module

244 + # object under the import alias, so patching it would have no effect.

245 + try:

246 + import triton_kernels.routing as _routing

247 + from triton_kernels.routing_details import _routing_compute as _rc

248 + except ImportError:

249 + return

250 +

251 + _keyed_add = _rc._keyed_add

252 + _expt_data_compute = _rc._expt_data_compute

253 +

254 + @triton.jit

255 + def _routing_compute_indx_pow2(

256 + pid_m,

257 + GatherIndx,

258 + ScatterIndx,

259 + GateScal,

260 + ExptScal,

261 + ExptIndx,

262 + PartialOffs,

263 + stride_pm,

264 + stride_pn,

265 + TokensStart,

266 + n_tokens,

267 + BLOCK_M: tl.constexpr,

268 + N_EXPTS_ACT: tl.constexpr,

269 + BLOCK_SIZE_PADDED: tl.constexpr,

270 + ):

271 + if isinstance(n_tokens, tl.tensor) and n_tokens.dtype.is_ptr():

272 + n_tokens = tl.load(n_tokens)

273 + n_gates = n_tokens * N_EXPTS_ACT

274 + BLOCK_SIZE: tl.constexpr = N_EXPTS_ACT * BLOCK_M

275 + tl.static_assert(BLOCK_SIZE_PADDED <= 32768)

276 + local_offs = tl.arange(0, BLOCK_SIZE_PADDED)

277 + offs = pid_m * BLOCK_SIZE + local_offs

278 + expert = tl.load(

279 + ExptIndx + offs,

280 + mask=(local_offs < BLOCK_SIZE) & (offs < n_gates),

281 + other=-1,

282 + ).to(tl.uint32)

283 + kv_pairs = ((expert << 16) | local_offs).to(tl.uint32)

284 + kv_pairs = tl.sort(kv_pairs, 0)

285 + expert = kv_pairs >> 16

286 + offs = pid_m * BLOCK_SIZE + (kv_pairs & 0xFFFF)

287 + mask = expert != 0xFFFF

288 + gate_scal = tl.load(ExptScal + offs, mask=mask)

289 + x = kv_pairs & 0xFFFF0000 | 0x00000001

290 + run_lengths = tl.associative_scan(x, 0, _keyed_add)

291 + exclusive_run_lengths = (run_lengths - 1) & 0xFFFF

292 + gates = tl.load(PartialOffs + pid_m * stride_pm + expert * stride_pn, mask=mask)

293 + gates += tl.load(TokensStart + expert, mask=mask)

294 + gates += exclusive_run_lengths

295 + tl.store(ScatterIndx + offs, gates, mask=mask)

296 + tl.store(GatherIndx + gates, offs, mask=mask)

297 + tl.store(GateScal + gates, gate_scal, mask=mask)

298 +

299 + @triton.jit

300 + def _combined_routing_compute_pow2(

301 + GatherIndx,

302 + ScatterIndx,

303 + GateScal,

304 + ExptScal,

305 + ExptIndx,

306 + PartialOffs,

307 + stride_pm,

308 + stride_pn,

309 + TokensStart,

310 + n_tokens,

311 + BLOCK_M: tl.constexpr,

312 + N_EXPTS_ACT: tl.constexpr,

313 + Hist,

314 + MDTileStarts,

315 + tile_starts_stridem,

316 + MDTileInfo,

317 + tile_info_stridem,

318 + first_tile_dim_log2,

319 + SIZES: tl.constexpr,

320 + BLOCK: tl.constexpr,

321 + blocks2a,

322 + BLOCK_SIZE_PADDED: tl.constexpr,

323 + ):

324 + pid = tl.program_id(0)

325 + if pid < blocks2a:

326 + _expt_data_compute(

327 + Hist,

328 + MDTileStarts,

329 + tile_starts_stridem,

330 + MDTileInfo,

331 + tile_info_stridem,

332 + first_tile_dim_log2,

333 + SIZES,

334 + BLOCK,

335 + )

336 + else:

337 + pid -= blocks2a

338 + _routing_compute_indx_pow2(

339 + pid,

340 + GatherIndx,

341 + ScatterIndx,

342 + GateScal,

343 + ExptScal,

344 + ExptIndx,

345 + PartialOffs,

346 + stride_pm,

347 + stride_pn,

348 + TokensStart,

349 + n_tokens,

350 + BLOCK_M,

351 + N_EXPTS_ACT,

352 + BLOCK_SIZE_PADDED,

353 + )

354 +

355 + def _sort_tokens_pow2(expt_scal, expt_indx, n_expts_tot, bitmatrix):

356 + import torch

357 +

358 + HIST_BLOCK_M = 32

359 + INDX_OFFS_BLOCK_M = 512

360 + MEMSET_BLOCK = 1024

361 + cdiv = triton.cdiv

362 + device = expt_scal.device

363 + dtype = expt_scal.dtype

364 + n_tokens_raw, _ = bitmatrix.shape

365 + n_tokens_pad, n_expts_act = expt_scal.shape

366 + n_gates_pad = n_tokens_pad * n_expts_act

367 + # pad per-block gate count (HIST_BLOCK_M * top_k) up to a pow2.

368 + block_size_padded = triton.next_power_of_2(HIST_BLOCK_M * n_expts_act)

369 +

370 + hist, partial_hist = bitmatrix.sum(partials_block_size=HIST_BLOCK_M)

371 + hist = hist[:n_expts_tot]

372 + expt_offs = torch.empty(n_expts_tot, dtype=torch.int32, device=device)

373 + combined_indx = torch.empty(n_gates_pad * 2, dtype=torch.int32, device=device)

374 + topk_indx = combined_indx[:n_gates_pad]

375 + gate_indx = combined_indx[n_gates_pad:]

376 + gate_scal = torch.empty(n_gates_pad, dtype=dtype, device=device)

377 +

378 + (

379 + token_offs_combined,

380 + token_offs_raw,

381 + token_offs_pad,

382 + block_pid_map,

383 + blocks1a,

384 + blocks2a,

385 + MEMSET_BLOCK_A,

386 + HIST2_BLOCK_M,

387 + block_m_log2_start,

388 + block_m_num,

389 + ) = _routing._compute_expt_data_internal(hist, n_expts_tot, n_gates_pad)

390 +

391 + blocks1b = cdiv(n_gates_pad * 2, MEMSET_BLOCK) + n_expts_tot + 1

392 + blocks2b = cdiv(n_tokens_pad, HIST_BLOCK_M)

393 +

394 + _rc._combined_routing_memset[(blocks1a + blocks1b,)](

395 + combined_indx,

396 + n_gates_pad * 2,

397 + -1,

398 + MEMSET_BLOCK,

399 + hist,

400 + expt_offs,

401 + hist.shape[0],

402 + n_expts_tot,

403 + partial_hist,

404 + partial_hist.shape[0],

405 + partial_hist.stride(0),

406 + partial_hist.stride(1),

407 + token_offs_combined,

408 + token_offs_combined.stride(0),

409 + blocks1a,

410 + block_pid_map,

411 + block_m_log2_start,

412 + SIZES=block_m_num,

413 + BLOCK_A=MEMSET_BLOCK_A,

414 + BLOCK_N=512,

415 + BLOCK_M=INDX_OFFS_BLOCK_M,

416 + )

417 +

418 + indx_offs = partial_hist

419 + _combined_routing_compute_pow2[(blocks2a + blocks2b,)](

420 + topk_indx,

421 + gate_indx,

422 + gate_scal,

423 + expt_scal,

424 + expt_indx,

425 + indx_offs,

426 + indx_offs.stride(0),

427 + indx_offs.stride(1),

428 + expt_offs,

429 + n_tokens_raw,

430 + HIST_BLOCK_M,

431 + n_expts_act,

432 + hist,

433 + token_offs_pad,

434 + token_offs_pad.stride(0),

435 + block_pid_map,

436 + block_pid_map.stride(0),

437 + block_m_log2_start,

438 + block_m_num,

439 + HIST2_BLOCK_M,

440 + blocks2a,

441 + block_size_padded,

442 + )

443 + return (

444 + hist,

445 + topk_indx,

446 + gate_indx,

447 + gate_scal,

448 + token_offs_raw,

449 + token_offs_pad,

450 + block_pid_map,

451 + )

452 +

453 + # `routing_from_bitmatrix` looks up `sort_tokens` via the routing module

454 + # global, so replacing it here redirects the legacy path to the pow2 kernel.

455 + _routing.sort_tokens = _sort_tokens_pow2

456 +

457 +

458 # Two API generations of triton_kernels are supported:

459 # - v3.5.1 (the version bundled with vLLM): exposes `routing()` and

460 # `routing_from_bitmatrix()` in triton_kernels.routing; the `Bitmatrix`

@@ -260,6 +495,9 @@ if has_triton_kernels():

495 use_legacy_triton_kernels = True

496 if not use_legacy_triton_kernels:

497 _patch_make_bitmatrix_metadata()

498 + else:

499 + # Legacy routing fails to compile for non-pow2 top_k (DeepSeek-V4).

500 + _patch_legacy_routing_for_nonpow2_topk()

501 except (AttributeError, ImportError) as e:

502 logger.error(

503 "Failed to import Triton kernels. Please make sure your triton "

[Bugfix] Support non-power-of-2 top_k in legacy triton_kernels routing (#46406)