@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d9dabfa351

[Kimi-K3][Kernel] Optimize the low-M fused latent MoE tail (#54168)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Yongye Zhu committed Aug 28, 2026 at 10:45 UTC d9dabfa351acdb4ceaad18c9677bc7a08b6c04ea

6 files changed +653 -194

tests/models/kimi_k3/test_latent_moe_tail.py

+107 -2

@@ -153,6 +153,57 @@ def _make_deferred_routed_output(

153 )

154

155

156 +def _make_bf16_top16_deferred_output(

157 + num_tokens: int,

158 + device: torch.device,

159 + *,

160 + drop_last_route: bool = False,

161 +) -> tuple[UnfinalizedMoEOutput, torch.Tensor]:

162 + top_k = 16

163 + num_routes = num_tokens * top_k

164 + num_permuted_rows = num_routes + 7

165 + expanded_output = torch.randn(

166 + num_routes,

167 + LATENT_SIZE,

168 + device=device,

169 + dtype=torch.bfloat16,

170 + ).mul_(0.01)

171 + expert_weights = torch.rand(

172 + num_tokens,

173 + top_k,

174 + device=device,

175 + dtype=torch.bfloat16,

176 + )

177 + expert_weights.div_(expert_weights.sum(dim=-1, keepdim=True))

178 + expanded_idx = torch.randperm(num_permuted_rows, device=device)[:num_routes]

179 + gemm2_permuted = torch.empty(

180 + num_permuted_rows,

181 + LATENT_SIZE,

182 + device=device,

183 + dtype=torch.bfloat16,

184 + )

185 + gemm2_permuted[expanded_idx] = expanded_output

186 +

187 + expanded_output = expanded_output.view(num_tokens, top_k, LATENT_SIZE)

188 + expanded_idx = expanded_idx.view(num_tokens, top_k)

189 + if drop_last_route:

190 + expanded_output[:, -1].zero_()

191 + expanded_idx[:, -1] = -1

192 + finalized = (

193 + (expanded_output.float() * expert_weights[:, :, None].float())

194 + .sum(1)

195 + .to(torch.bfloat16)

196 + )

197 + return (

198 + UnfinalizedMoEOutput(

199 + gemm2_permuted=gemm2_permuted,

200 + expert_weights=expert_weights,

201 + expanded_idx_to_permuted_idx=expanded_idx.to(torch.int32),

202 + ),

203 + finalized,

204 + )

205 +

206 +

207 @ray.remote(num_gpus=1, max_calls=1)

208 def _test_latent_moe_tail_worker(

209 monkeypatch: pytest.MonkeyPatch,

@@ -197,7 +248,8 @@ def _test_latent_moe_tail_worker(

248 )

249 cutedsl_warmup()

250

200 - for iteration, num_tokens in enumerate((1, 5, 8, 16, 5)):

251 + # M=65 crosses two 32-CTA token waves and reuses the first DSM slot.

252 + for iteration, num_tokens in enumerate((1, 5, 8, 16, 33, 65, 5)):

253 torch.manual_seed(100 * iteration + rank + 1)

254 routed_output = torch.randn(

255 num_tokens,

@@ -310,9 +362,17 @@ def _test_deferred_finalize_parity_worker(

362 rms_eps=EPS,

363 experts_per_token=TOP_K,

364 )

365 + bf16_deferred_op = KimiK3LatentMoETailOp.initialize(

366 + hidden_size=HIDDEN_SIZE,

367 + latent_size=LATENT_SIZE,

368 + dtype=torch.bfloat16,

369 + device=device,

370 + rms_eps=EPS,

371 + experts_per_token=16,

372 + )

373 cutedsl_warmup()

374

315 - for iteration, num_tokens in enumerate((1, 5, 16)):

375 + for iteration, num_tokens in enumerate((1, 5, 16, 33, 65)):

376 torch.manual_seed(2000 + 100 * iteration + rank)

377 deferred_output, finalized_output = _make_deferred_routed_output(

378 num_tokens,

@@ -340,6 +400,51 @@ def _test_deferred_finalize_parity_worker(

400

401 torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)

402

403 + for iteration, num_tokens in enumerate((1, 2, 4, 5, 16)):

404 + torch.manual_seed(3000 + 100 * iteration + rank)

405 + deferred_output, finalized_output = _make_bf16_top16_deferred_output(

406 + num_tokens,

407 + device,

408 + drop_last_route=num_tokens == 5,

409 + )

410 + shared_output = torch.randn(

411 + num_tokens,

412 + HIDDEN_SIZE,

413 + device=device,

414 + dtype=torch.bfloat16,

415 + )

416 +

417 + expected = finalized_op(

418 + finalized_output,

419 + shared_output,

420 + rms_weight,

421 + up_weight,

422 + )

423 + actual = bf16_deferred_op(

424 + deferred_output,

425 + shared_output,

426 + rms_weight,

427 + up_weight,

428 + )

429 +

430 + torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)

431 + if num_tokens == 4:

432 + graph = torch.cuda.CUDAGraph()

433 + with torch.cuda.graph(graph):

434 + graph_actual = bf16_deferred_op(

435 + deferred_output,

436 + shared_output,

437 + rms_weight,

438 + up_weight,

439 + )

440 + graph.replay()

441 + torch.testing.assert_close(

442 + graph_actual,

443 + expected,

444 + atol=8e-2,

445 + rtol=3e-2,

446 + )

447 +

448

449 def _run_deferred_finalize_parity_test(

450 monkeypatch: pytest.MonkeyPatch,

vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/allreduce_rmsnorm_reduce_scatter_early_exit.py

+327 -161

@@ -23,21 +23,27 @@ from .primitives import (

23 PACKED_BYTES,

24 VEC_BF16,

25 bf16x8_to_packed_u32x4,

26 - block_sum_specialized,

26 + finalize_top16_bf16,

27 fragment_is_dirty,

28 load_global_u32x4,

29 + load_shared_f32x2,

30 + load_shared_f32x4,

31 load_volatile_u32,

32 map_shared_to_peer,

33 packed_u32x4_to_bf16x8,

34 red_async_release_gpu_add_u32,

35 sanitize_negative_zero,

36 + stmc_bf16x8,

37 store_global_u32x4,

38 store_lamport_sentinel_128,

39 store_shared_cluster_f32,

40 to_cute,

41 to_cute_dynamic_m,

42 + warp_sum_specialized,

43 )

45 +_SEVEN_CTA_MAX_M = 4

46 +

48 def _mapping(

49 tp_size: int,

@@ -92,7 +98,7 @@ def _select_routed_schedule(

100 class AllReduceRMSNormWithReduceScatterEarlyExit:

95 - """One routed role plus one ReduceScatter role per destination group."""

101 + """One routed role plus one compact ReduceScatter role."""

102

103 def __init__(

104 self,

@@ -127,22 +133,44 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

133 mapped_threads,

134 shared_roles,

135 ) = _mapping(tp_size, latent_dim, hidden_dim)

136 + seven_cta_geometry = (

137 + include_routed

138 + and max_m <= _SEVEN_CTA_MAX_M

139 + and max_token_ctas == max_m

140 + and (tp_size, latent_dim, hidden_dim) == (8, 3584, 7168)

141 + )

142 + self.seven_cta_geometry = seven_cta_geometry

143 + self.single_token_geometry = seven_cta_geometry and max_m == 1

144 if include_reduce_scatter:

131 - self.cluster_ctas = mapped_cluster

132 - self.threads = mapped_threads

145 + if seven_cta_geometry:

146 + self.cluster_ctas = 7

147 + self.threads = 64

148 + shared_roles = (tp_size + self.cluster_ctas - 1) // self.cluster_ctas

149 + else:

150 + self.cluster_ctas = mapped_cluster

151 + self.threads = mapped_threads

152 else:

134 - self.threads, self.cluster_ctas = _select_routed_schedule(

135 - tp_size, latent_dim, hidden_dim, max_m

136 - )

153 + if seven_cta_geometry:

154 + self.threads, self.cluster_ctas = 64, 7

155 + else:

156 + self.threads, self.cluster_ctas = _select_routed_schedule(

157 + tp_size, latent_dim, hidden_dim, max_m

158 + )

159 + self.shared_roles = 1 if include_reduce_scatter else shared_roles

160 + self.shared_destination_stride = self.shared_roles * self.cluster_ctas

161 + self.shared_destinations_per_cta = (

162 + self.tp_size + self.shared_destination_stride - 1

163 + ) // self.shared_destination_stride

164 + self.shard_vectors = self.shard_dim // VEC_BF16

165 # Diagnostic specialization: a one-role grid executes only the routed

166 # AllReduce/RMSNorm path. Keep this compile-time so the production

167 # fused path is unchanged when include_reduce_scatter=True.

168 if include_routed and include_reduce_scatter:

141 - self.roles = 1 + shared_roles

169 + self.roles = 1 + self.shared_roles

170 elif include_routed:

171 self.roles = 1

172 else:

145 - self.roles = shared_roles

173 + self.roles = self.shared_roles

174 self.warps = (self.threads + 31) // 32

175 self.last_warp_lanes = self.threads - (self.warps - 1) * 32

176 self.last_warp_mask = (1 << self.last_warp_lanes) - 1

@@ -178,6 +206,7 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

206 expert_weights: cute.Tensor,

207 expanded_idx_to_permuted_idx: cute.Tensor,

208 ):

209 + grid_x = m if cutlass.const_expr(self.seven_cta_geometry) else self.token_ctas

210 self.kernel(

211 latent_source,

212 gamma,

@@ -195,10 +224,10 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

224 expert_weights,

225 expanded_idx_to_permuted_idx,

226 ).launch(

198 - grid=(self.token_ctas, self.cluster_ctas, self.roles),

227 + grid=(grid_x, self.cluster_ctas, self.roles),

228 block=(self.threads, 1, 1),

229 cluster=(1, self.cluster_ctas, 1),

201 - smem=(self.warps + self.cluster_ctas) * 4,

230 + smem=2 * self.cluster_ctas * self.warps * 4,

231 stream=stream,

232 use_pdl=True,

233 )

@@ -233,6 +262,7 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

262

263 cute.arch.griddepcontrol_wait()

264 token = token_cta

265 + parity = Int32(0)

266 while token < m:

267 self._token_device(

268 latent_source,

@@ -249,6 +279,7 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

279 m,

280 epsilon,

281 token,

282 + parity,

283 token_cta,

284 cta_y,

285 logical_role,

@@ -258,6 +289,7 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

289 expanded_idx_to_permuted_idx,

290 )

291 token = token + self.token_ctas

292 + parity = parity ^ Int32(1)

293

294 @cute.jit

295 def _token_device(

@@ -276,6 +308,7 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

308 m: Int32,

309 epsilon: Float32,

310 token: Int32,

311 + parity: Int32,

312 token_cta: Int32,

313 cta_y: Int32,

314 role: Int32,

@@ -308,36 +341,72 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

341

342 if cutlass.const_expr(self.top_k > 0):

343 # finalize topk reduction

311 - local_values = cute.make_rmem_tensor(

312 - cute.make_layout((VEC_BF16,)), BFloat16

313 - )

314 - for element in cutlass.range_constexpr(VEC_BF16):

315 - local_values[element] = BFloat16(0.0)

316 - for slot in cutlass.range_constexpr(self.top_k):

317 - permuted_idx = expanded_idx_to_permuted_idx[token, slot]

318 - if permuted_idx >= Int32(0):

319 - permuted_element = (

320 - Int64(permuted_idx) * self.latent_dim

321 - + Int64(packed_idx) * VEC_BF16

322 - )

323 - permuted_ptr = cute.make_ptr(

324 - BFloat16,

325 - (latent_source.iterator + permuted_element).llvm_ptr,

326 - cute.AddressSpace.gmem,

327 - assumed_align=16,

328 - )

329 - values = packed_u32x4_to_bf16x8(

330 - load_global_u32x4(permuted_ptr, volatile=False)

344 + if cutlass.const_expr(

345 + self.top_k == 16

346 + and self.latent_dim == 3584

347 + and expert_weights.element_type == BFloat16

348 + ):

349 + gemm2_vector = cute.make_ptr(

350 + BFloat16,

351 + (

352 + latent_source.iterator + Int64(packed_idx) * VEC_BF16

353 + ).llvm_ptr,

354 + cute.AddressSpace.gmem,

355 + assumed_align=16,

356 + )

357 + route_indices = cute.make_ptr(

358 + Int32,

359 + (

360 + expanded_idx_to_permuted_idx.iterator

361 + + Int64(token) * self.top_k

362 + ).llvm_ptr,

363 + cute.AddressSpace.gmem,

364 + assumed_align=16,

365 + )

366 + route_weights = cute.make_ptr(

367 + BFloat16,

368 + (expert_weights.iterator + Int64(token) * self.top_k).llvm_ptr,

369 + cute.AddressSpace.gmem,

370 + assumed_align=16,

371 + )

372 + local_packed = sanitize_negative_zero(

373 + finalize_top16_bf16(

374 + gemm2_vector,

375 + route_indices,

376 + route_weights,

377 )

332 - weight = expert_weights[token, slot].to(Float32)

333 - for element in cutlass.range_constexpr(VEC_BF16):

334 - scaled = (values[element].to(Float32) * weight).to(BFloat16)

335 - local_values[element] = (

336 - local_values[element].to(Float32) + scaled.to(Float32)

337 - ).to(BFloat16)

338 - local_packed = sanitize_negative_zero(

339 - bf16x8_to_packed_u32x4(local_values.load())

340 - )

378 + )

379 + else:

380 + local_values = cute.make_rmem_tensor(

381 + cute.make_layout((VEC_BF16,)), Float32

382 + )

383 + for element in cutlass.range_constexpr(VEC_BF16):

384 + local_values[element] = Float32(0.0)

385 + for slot in cutlass.range_constexpr(self.top_k):

386 + permuted_idx = expanded_idx_to_permuted_idx[token, slot]

387 + if permuted_idx >= Int32(0):

388 + permuted_element = (

389 + Int64(permuted_idx) * self.latent_dim

390 + + Int64(packed_idx) * VEC_BF16

391 + )

392 + permuted_ptr = cute.make_ptr(

393 + BFloat16,

394 + (latent_source.iterator + permuted_element).llvm_ptr,

395 + cute.AddressSpace.gmem,

396 + assumed_align=16,

397 + )

398 + values = packed_u32x4_to_bf16x8(

399 + load_global_u32x4(permuted_ptr, volatile=False)

400 + )

401 + weight = expert_weights[token, slot].to(Float32)

402 + for element in cutlass.range_constexpr(VEC_BF16):

403 + local_values[element] = (

404 + local_values[element]

405 + + values[element].to(Float32) * weight

406 + )

407 + local_packed = sanitize_negative_zero(

408 + bf16x8_to_packed_u32x4(local_values.load().to(BFloat16))

409 + )

410 else:

411 local_ptr = cute.make_ptr(

412 BFloat16,

@@ -356,17 +425,20 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

425 )

426 * 2

427 )

359 - store_global_u32x4(

428 + stmc_bf16x8(

429 latent_multicast_ptr + multicast_offset,

430 local_packed,

362 - volatile=False,

431 )

432 + cute.arch.griddepcontrol_launch_dependents()

433

365 - cute.arch.cluster_arrive()

366 - if cluster_rank == 0 and tidx < 32:

367 - cute.arch.cluster_wait()

368 - if tidx == 0:

369 - red_async_release_gpu_add_u32(latent_flags.iterator + 8, Uint32(1))

434 + if cutlass.const_expr(not self.single_token_geometry):

435 + cute.arch.cluster_arrive()

436 + if cluster_rank == 0 and tidx < 32:

437 + cute.arch.cluster_wait()

438 + if tidx == 0:

439 + red_async_release_gpu_add_u32(

440 + latent_flags.iterator + 8, Uint32(1)

441 + )

442

443 global_tid = (

444 Int64(token) * self.cluster_ctas + Int64(cta_y)

@@ -425,9 +497,6 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

497 for element in cutlass.range_constexpr(VEC_BF16):

498 accum[element] = accum[element] + values[element]

499

428 - # Preserve the original early PDL point before RMSNorm.

429 - cute.arch.griddepcontrol_launch_dependents()

430 -

500 if cutlass.const_expr(self.fp32_internal):

501 # High-precision fused mode: retain the rank reduction in

502 # FP32 through the RMS square and row reduction.

@@ -445,32 +514,53 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

514 reduction_profile=0,

515 )

516 smem = cutlass.utils.SmemAllocator()

448 - warp_sums = smem.allocate_tensor(

449 - Float32, cute.make_layout((self.warps,)), byte_alignment=4

450 - )

517 cluster_sums = smem.allocate_tensor(

518 Float32,

453 - cute.make_layout((self.cluster_ctas,)),

454 - byte_alignment=4,

519 + cute.make_layout((2 * self.cluster_ctas * self.warps,)),

520 + byte_alignment=16,

521 )

456 - block_sum = block_sum_specialized(

522 + lane = cute.arch.lane_idx()

523 + warp_idx = cute.arch.warp_idx()

524 + warp_sum = warp_sum_specialized(

525 thread_sum,

458 - warp_sums,

459 - tidx,

526 + warp_idx,

527 + lane,

528 self.warps,

529 self.last_warp_lanes,

530 self.last_warp_mask,

531 )

464 - if tidx < self.cluster_ctas:

465 - local_slot = cluster_sums.iterator + cluster_rank

466 - remote_slot = map_shared_to_peer(local_slot, Int32(tidx))

467 - store_shared_cluster_f32(remote_slot, block_sum)

532 + # Alternate DSM slots until the next cluster synchronization so

533 + # peers may safely begin publishing the following token wave.

534 + parity_offset = parity * Int32(self.cluster_ctas * self.warps)

535 + if lane < self.cluster_ctas:

536 + local_slot = (

537 + cluster_sums.iterator

538 + + parity_offset

539 + + cluster_rank * self.warps

540 + + warp_idx

541 + )

542 + remote_slot = map_shared_to_peer(local_slot, lane)

543 + store_shared_cluster_f32(remote_slot, warp_sum)

544 cute.arch.cluster_arrive()

545 cute.arch.cluster_wait()

546

547 full_sum = Float32(0.0)

548 for peer in cutlass.range_constexpr(self.cluster_ctas):

473 - full_sum = full_sum + cluster_sums[peer]

549 + peer_slot = cluster_sums.iterator + parity_offset + peer * self.warps

550 + if cutlass.const_expr(self.warps == 4):

551 + sum0, sum1, sum2, sum3 = load_shared_f32x4(peer_slot)

552 + full_sum = full_sum + sum0 + sum1 + sum2 + sum3

553 + elif cutlass.const_expr(self.warps == 2):

554 + sum0, sum1 = load_shared_f32x2(peer_slot)

555 + full_sum = full_sum + sum0 + sum1

556 + else:

557 + for peer_warp in cutlass.range_constexpr(self.warps):

558 + full_sum = (

559 + full_sum

560 + + cluster_sums[

561 + parity_offset + peer * self.warps + peer_warp

562 + ]

563 + )

564 inv_rms = cute.math.rsqrt(

565 full_sum / Float32(self.latent_dim) + epsilon, fastmath=True

566 )

@@ -490,9 +580,9 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

580 volatile=False,

581 )

582

493 - # The x=0 CTA rotates only after reaching its final token wave.

494 - # Waiting for all M arrivals then guarantees every token-wave CTA

495 - # loaded the current generation before the metadata is advanced.

583 + # The general schedule waits until every token cluster has loaded

584 + # this generation. The M=1 schedule has only this cluster, whose

585 + # DSM barrier above already covers all routed CTAs.

586 if (

587 token_cta == 0

588 and token + self.token_ctas >= m

@@ -500,9 +590,10 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

590 and tidx == 0

591 ):

592 access_counter = latent_flags.iterator + 8

503 - arrived = load_volatile_u32(access_counter)

504 - while arrived < Uint32(m):

593 + if cutlass.const_expr(not self.single_token_geometry):

594 arrived = load_volatile_u32(access_counter)

595 + while arrived < Uint32(m):

596 + arrived = load_volatile_u32(access_counter)

597 next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS)

598 actual_bytes = Uint32(m) * Uint32(self.tp_size * self.latent_dim * 2)

599 cute.arch.store((latent_flags.iterator + 0).llvm_ptr, next_index)

@@ -540,111 +631,136 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

631 )

632 dirty_elements = Int64(dirty_index) * (Int64(bytes_per_buffer) // Int64(2))

633

543 - source_element = (

544 - Int64(token) * self.hidden_dim

545 - + Int64(destination) * self.shard_dim

546 - + Int64(tidx) * VEC_BF16

547 - )

548 - source_ptr = cute.make_ptr(

549 - BFloat16,

550 - (shared_source.iterator + source_element).llvm_ptr,

551 - cute.AddressSpace.gmem,

552 - assumed_align=16,

553 - )

554 - local_packed = sanitize_negative_zero(

555 - load_global_u32x4(source_ptr, volatile=False)

556 - )

557 - peer_base = cute.arch.load(

558 - (shared_peer_ptrs.iterator + destination).llvm_ptr,

559 - Int64,

560 - )

561 - destination_element = current_elements + (

562 - (Int64(token) * self.tp_size + self.rank) * self.shard_dim

563 - + Int64(tidx) * VEC_BF16

564 - )

565 - store_global_u32x4(

566 - peer_base + destination_element * 2,

567 - local_packed,

568 - volatile=False,

569 - )

634 + for destination_round in cutlass.range_constexpr(

635 + self.shared_destinations_per_cta

636 + ):

637 + round_destination = destination + Int32(

638 + destination_round * self.shared_destination_stride

639 + )

640 + if round_destination < Int32(self.tp_size):

641 + peer_base = cute.arch.load(

642 + (shared_peer_ptrs.iterator + round_destination).llvm_ptr,

643 + Int64,

644 + )

645 + vector = tidx

646 + while vector < self.shard_vectors:

647 + source_element = (

648 + Int64(token) * self.hidden_dim

649 + + Int64(round_destination) * self.shard_dim

650 + + Int64(vector) * VEC_BF16

651 + )

652 + source_ptr = cute.make_ptr(

653 + BFloat16,

654 + (shared_source.iterator + source_element).llvm_ptr,

655 + cute.AddressSpace.gmem,

656 + assumed_align=16,

657 + )

658 + local_packed = sanitize_negative_zero(

659 + load_global_u32x4(source_ptr, volatile=False)

660 + )

661 + destination_element = current_elements + (

662 + (Int64(token) * self.tp_size + self.rank) * self.shard_dim

663 + + Int64(vector) * VEC_BF16

664 + )

665 + store_global_u32x4(

666 + peer_base + destination_element * 2,

667 + local_packed,

668 + volatile=False,

669 + )

670 + vector = vector + self.threads

671 +

672 + cute.arch.griddepcontrol_launch_dependents()

673

571 - # One arrival per shared destination group and token.

674 + # One arrival per shared cluster and token.

675 cute.arch.cluster_arrive()

676 if cluster_rank == 0 and tidx < 32:

677 cute.arch.cluster_wait()

678 if tidx == 0:

679 red_async_release_gpu_add_u32(shared_flags.iterator + 8, Uint32(1))

680

578 - global_tid = (

579 - Int64(token) * self.tp_size + Int64(destination)

580 - ) * self.threads + Int64(tidx)

681 total_threads = Int64(m) * self.tp_size * self.threads

682 clear_fragments = (Int64(bytes_to_clear) + PACKED_BYTES - 1) // PACKED_BYTES

583 - clear_idx = global_tid

584 - if dirty_num_stages > Uint32(0):

585 - while clear_idx < clear_fragments:

586 - clear_ptr = cute.make_ptr(

587 - BFloat16,

588 - (

589 - shared_workspace.iterator

590 - + dirty_elements

591 - + clear_idx * VEC_BF16

592 - ).llvm_ptr,

593 - cute.AddressSpace.gmem,

594 - assumed_align=16,

595 - )

596 - store_lamport_sentinel_128(clear_ptr)

597 - clear_idx = clear_idx + total_threads

598 -

599 - if destination == self.rank:

600 - rank_words = cute.make_rmem_tensor(

601 - cute.make_layout((self.tp_size, 4), stride=(4, 1)), Uint32

683 + for destination_round in cutlass.range_constexpr(

684 + self.shared_destinations_per_cta

685 + ):

686 + round_destination = destination + Int32(

687 + destination_round * self.shared_destination_stride

688 )

603 - valid = False

604 - while not valid:

605 - valid = True

606 - for source_rank in cutlass.range_constexpr(self.tp_size):

607 - remote_element = current_elements + (

608 - (Int64(token) * self.tp_size + source_rank) * self.shard_dim

609 - + Int64(tidx) * VEC_BF16

610 - )

611 - remote_ptr = cute.make_ptr(

689 + global_tid = (

690 + Int64(token) * self.tp_size + Int64(round_destination)

691 + ) * self.threads + Int64(tidx)

692 + clear_idx = global_tid

693 + if round_destination < Int32(

694 + self.tp_size

695 + ) and dirty_num_stages > Uint32(0):

696 + while clear_idx < clear_fragments:

697 + clear_ptr = cute.make_ptr(

698 BFloat16,

613 - (shared_workspace.iterator + remote_element).llvm_ptr,

699 + (

700 + shared_workspace.iterator

701 + + dirty_elements

702 + + clear_idx * VEC_BF16

703 + ).llvm_ptr,

704 cute.AddressSpace.gmem,

705 assumed_align=16,

706 )

617 - remote = load_global_u32x4(remote_ptr, volatile=True)

618 - for word in cutlass.range_constexpr(4):

619 - rank_words[source_rank, word] = remote[word]

620 - valid = valid & (not fragment_is_dirty(remote))

621 -

622 - accum = cute.make_rmem_tensor(cute.make_layout((VEC_BF16,)), Float32)

623 - for element in cutlass.range_constexpr(VEC_BF16):

624 - accum[element] = Float32(0.0)

625 - for source_rank in cutlass.range_constexpr(self.tp_size):

626 - values = packed_u32x4_to_bf16x8(

627 - rank_words[source_rank, None].load()

628 - ).to(Float32)

629 - for element in cutlass.range_constexpr(VEC_BF16):

630 - accum[element] = accum[element] + values[element]

631 - result = accum.load().to(BFloat16)

632 - output_element = (

633 - Int64(token) * self.hidden_dim

634 - + self.rank * self.shard_dim

635 - + Int64(tidx) * VEC_BF16

636 - )

637 - store_global_u32x4(

638 - Int64((shared_output.iterator + output_element).toint()),

639 - bf16x8_to_packed_u32x4(result),

640 - volatile=False,

641 - )

642 -

643 - if destination == self.rank:

644 - cute.arch.barrier()

645 -

646 - cute.arch.griddepcontrol_launch_dependents()

707 + store_lamport_sentinel_128(clear_ptr)

708 + clear_idx = clear_idx + total_threads

709 +

710 + if round_destination == self.rank:

711 + vector = tidx

712 + while vector < self.shard_vectors:

713 + rank_words = cute.make_rmem_tensor(

714 + cute.make_layout((self.tp_size, 4), stride=(4, 1)),

715 + Uint32,

716 + )

717 + valid = False

718 + while not valid:

719 + valid = True

720 + for source_rank in cutlass.range_constexpr(self.tp_size):

721 + remote_element = current_elements + (

722 + (Int64(token) * self.tp_size + source_rank)

723 + * self.shard_dim

724 + + Int64(vector) * VEC_BF16

725 + )

726 + remote_ptr = cute.make_ptr(

727 + BFloat16,

728 + (

729 + shared_workspace.iterator + remote_element

730 + ).llvm_ptr,

731 + cute.AddressSpace.gmem,

732 + assumed_align=16,

733 + )

734 + remote = load_global_u32x4(remote_ptr, volatile=True)

735 + for word in cutlass.range_constexpr(4):

736 + rank_words[source_rank, word] = remote[word]

737 + valid = valid & (not fragment_is_dirty(remote))

738 +

739 + accum = cute.make_rmem_tensor(

740 + cute.make_layout((VEC_BF16,)), Float32

741 + )

742 + for element in cutlass.range_constexpr(VEC_BF16):

743 + accum[element] = Float32(0.0)

744 + for source_rank in cutlass.range_constexpr(self.tp_size):

745 + values = packed_u32x4_to_bf16x8(

746 + rank_words[source_rank, None].load()

747 + ).to(Float32)

748 + for element in cutlass.range_constexpr(VEC_BF16):

749 + accum[element] = accum[element] + values[element]

750 + result = accum.load().to(BFloat16)

751 + output_element = (

752 + Int64(token) * self.hidden_dim

753 + + self.rank * self.shard_dim

754 + + Int64(vector) * VEC_BF16

755 + )

756 + store_global_u32x4(

757 + Int64((shared_output.iterator + output_element).toint()),

758 + bf16x8_to_packed_u32x4(result),

759 + volatile=False,

760 + )

761 + vector = vector + self.threads

762

763 + cute.arch.barrier()

764 if (

765 token_cta == 0

766 and token + self.token_ctas >= m

@@ -654,7 +770,7 @@ class AllReduceRMSNormWithReduceScatterEarlyExit:

770 ):

771 access_counter = shared_flags.iterator + 8

772 arrived = load_volatile_u32(access_counter)

657 - target = Uint32(m) * Uint32(self.tp_size // self.cluster_ctas)

773 + target = Uint32(m) * Uint32(self.shared_roles)

774 while arrived < target:

775 arrived = load_volatile_u32(access_counter)

776 next_index = (current_index + Uint32(1)) % Uint32(NUM_LAMPORT_BUFFERS)

@@ -948,6 +1064,14 @@ class CollectiveKernel:

1064 self._dummy_expanded_idx = torch.empty(

1065 (max_m, max(top_k, 1)), dtype=torch.int32, device=device

1066 )

1067 + self._seven_cta_max_m = (

1068 + min(max_m, _SEVEN_CTA_MAX_M)

1069 + if (tp_size, latent_dim, hidden_dim) == (8, 3584, 7168)

1070 + and top_k == 16

1071 + and self._dummy_expert_weights.dtype == torch.bfloat16

1072 + and torch.cuda.get_device_capability(device)[0] == 10

1073 + else 0

1074 + )

1075

1076 bytes_per_routed_buffer = max_m * tp_size * latent_dim * 2

1077 routed_bytes = NUM_LAMPORT_BUFFERS * bytes_per_routed_buffer

@@ -1009,6 +1133,39 @@ class CollectiveKernel:

1133

1134 torch.accelerator.synchronize(device)

1135 dist.barrier(group=group, device_ids=[device.index])

1136 + if self._seven_cta_max_m:

1137 + specializations = (

1138 + (1, 1),

1139 + (self._seven_cta_max_m, self._seven_cta_max_m),

1140 + )

1141 + for owner in range(tp_size):

1142 + if rank == owner:

1143 + for compile_max_m, compile_token_ctas in specializations:

1144 + if (compile_max_m, compile_token_ctas) == (

1145 + max_m,

1146 + max_token_ctas,

1147 + ):

1148 + continue

1149 + compile_kernel(

1150 + rank=rank,

1151 + tp_size=tp_size,

1152 + latent_dim=latent_dim,

1153 + hidden_dim=hidden_dim,

1154 + max_m=compile_max_m,

1155 + max_token_ctas=compile_token_ctas,

1156 + latent_output=self._latent_output,

1157 + routed_workspace=self._routed_workspace,

1158 + routed_flags=self._routed_flags,

1159 + routed_multicast_ptr=self._routed_multicast_ptr,

1160 + shared_output=self._shared_output,

1161 + shared_workspace=self._shared_workspace,

1162 + shared_flags=self._shared_flags,

1163 + shared_peer_ptrs=self._shared_peer_ptrs,

1164 + rms_eps=self.rms_eps,

1165 + fp32_internal=fp32_internal,

1166 + top_k=top_k,

1167 + )

1168 + dist.barrier(group=group, device_ids=[device.index])

1169 for owner in range(tp_size):

1170 if rank == owner:

1171 compile_kernel(

@@ -1110,6 +1267,15 @@ class CollectiveKernel:

1267 if not 1 <= m <= self.max_m:

1268 raise ValueError(f"runtime M={m} must be in [1, {self.max_m}]")

1269

1270 + launch_max_m = self.max_m

1271 + launch_token_ctas = self.max_token_ctas

1272 + if self._seven_cta_max_m and m <= self._seven_cta_max_m:

1273 + if m == 1:

1274 + launch_max_m = 1

1275 + launch_token_ctas = 1

1276 + else:

1277 + launch_max_m = self._seven_cta_max_m

1278 + launch_token_ctas = self._seven_cta_max_m

1279 with torch.accelerator.device_index(device.index):

1280 launch(

1281 gemm2_permuted,

@@ -1128,8 +1294,8 @@ class CollectiveKernel:

1294 tp_size=self.tp_size,

1295 latent_dim=self.latent_dim,

1296 hidden_dim=self.hidden_dim,

1131 - max_m=self.max_m,

1132 - max_token_ctas=self.max_token_ctas,

1297 + max_m=launch_max_m,

1298 + max_token_ctas=launch_token_ctas,

1299 fp32_internal=self.fp32_internal,

1300 top_k=self.top_k,

1301 expert_weights=expert_weights,

vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/fused_add_multicast_skinny_gemm.py

+24 -10

@@ -19,6 +19,7 @@ from .primitives import (

19 bf16x2_to_u32,

20 bf16x4_to_packed_u32x2,

21 bf16x8_to_packed_u32x4,

22 + fma_f32_bf16,

23 sanitize_negative_zero,

24 sanitize_negative_zero_u32,

25 sanitize_negative_zero_u32x2,

@@ -38,6 +39,13 @@ class SkinnyConfig:

41 def config_for_m(num_rows: int, shard_dim: int = 896) -> SkinnyConfig:

42 + if num_rows <= 5 and shard_dim in (448, 896):

43 + return SkinnyConfig(

44 + block_size=224,

45 + outputs_per_block=2,

46 + k_unroll=1,

47 + vector_width=16,

48 + )

49 if shard_dim == 448:

50 if num_rows >= 6:

51 return SkinnyConfig(

@@ -59,7 +67,7 @@ def config_for_m(num_rows: int, shard_dim: int = 896) -> SkinnyConfig:

67 def _as_cute(tensor: torch.Tensor):

68 return from_dlpack(

69 CUDAGraphCompatibleWrapper(tensor.detach()),

62 - assumed_align=16,

70 + assumed_align=32,

71 )

@@ -194,9 +202,11 @@ class FusedAddMulticastSkinnyGemm:

202 for vi in cutlass.range_constexpr(vector_width):

203 for mi in cutlass.range_constexpr(num_rows):

204 for ni in cutlass.range_constexpr(outputs_per_block):

197 - acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(Float32) * b_regs[

198 - ni, vi

199 - ].to(Float32)

205 + acc[mi, ni] = fma_f32_bf16(

206 + a_regs[mi, vi],

207 + b_regs[ni, vi],

208 + acc[mi, ni],

209 + )

210

211 for k_tile in cutlass.range(1, num_k_tiles, unroll=self.k_unroll):

212 for mi in cutlass.range_constexpr(num_rows):

@@ -211,9 +221,11 @@ class FusedAddMulticastSkinnyGemm:

221 for vi in cutlass.range_constexpr(vector_width):

222 for mi in cutlass.range_constexpr(num_rows):

223 for ni in cutlass.range_constexpr(outputs_per_block):

214 - acc[mi, ni] = acc[mi, ni] + a_regs[mi, vi].to(Float32) * b_regs[

215 - ni, vi

216 - ].to(Float32)

224 + acc[mi, ni] = fma_f32_bf16(

225 + a_regs[mi, vi],

226 + b_regs[ni, vi],

227 + acc[mi, ni],

228 + )

229

230 for mi in cutlass.range_constexpr(num_rows):

231 for ni in cutlass.range_constexpr(outputs_per_block):

@@ -315,19 +327,19 @@ def compile_kernel(

327 BFloat16,

328 (num_rows, latent_dim),

329 stride=(latent_dim, 1),

318 - assumed_align=16,

330 + assumed_align=32,

331 )

332 b = make_fake_tensor(

333 BFloat16,

334 (shard_dim, latent_dim),

335 stride=(latent_dim, 1),

324 - assumed_align=16,

336 + assumed_align=32,

337 )

338 shared = make_fake_tensor(

339 BFloat16,

340 (num_rows, shard_dim),

341 stride=(hidden_dim, 1),

330 - assumed_align=16,

342 + assumed_align=32,

343 )

344 compiled = cute.compile(

345 FusedAddMulticastSkinnyGemm(

@@ -424,6 +436,8 @@ class FusedAddMulticastSkinnyGemmKernel:

436 or not mailbox.is_contiguous()

437 ):

438 raise ValueError("skinny up-projection inputs have unsupported strides")

439 + if any(tensor.data_ptr() % 32 for tensor in (latent, weight, shared_shard)):

440 + raise ValueError("skinny up-projection inputs must be 32-byte aligned")

441 if mailbox.shape[1] < self.num_rows:

442 raise ValueError("mailbox capacity is smaller than runtime M")

443

vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/lamport_copy.py

+11 -20

@@ -40,8 +40,13 @@ class LamportCopy:

40 m: cutlass.Int32,

41 stream: cuda.CUstream,

42 ):

43 + fragments = m * cutlass.Int32(self.hidden_dim // VEC_BF16)

44 + grid_ctas = cutlass.min(

45 + cutlass.Int32(self.ctas),

46 + cute.ceil_div(fragments, self.threads),

47 + )

48 self.kernel(symmetric_mailbox, local_output, m).launch(

44 - grid=(self.ctas, 1, 1),

49 + grid=(grid_ctas, 1, 1),

50 block=(self.threads, 1, 1),

51 stream=stream,

52 use_pdl=True,

@@ -54,14 +59,15 @@ class LamportCopy:

59 local_output: cute.Tensor,

60 m: cutlass.Int32,

61 ):

57 - # The CTA may be scheduled early, but mailbox inspection must not pass

58 - # the producer GEMM's programmatic completion point.

59 - cute.arch.griddepcontrol_wait()

62 + # The Lamport marker carries producer readiness, so polling can begin

63 + # before the producer grid reaches ordinary completion.

64 + cute.arch.griddepcontrol_launch_dependents()

66 tidx, _, _ = cute.arch.thread_idx()

67 block, _, _ = cute.arch.block_idx()

68 + grid_x, _, _ = cute.arch.grid_dim()

69 thread = cutlass.Int64(block * self.threads + tidx)

64 - stride = cutlass.Int64(self.ctas * self.threads)

70 + stride = cutlass.Int64(grid_x * self.threads)

71 fragments = cutlass.Int64(m) * cutlass.Int64(self.hidden_dim // VEC_BF16)

73 fragment = thread

@@ -79,21 +85,6 @@ class LamportCopy:

86 destination = cutlass.Int64((local_output.iterator + element).toint())

87 store_global_u32x4(destination, packed, volatile=False)

82 - fragment = fragment + stride

83 -

84 - # The returned ordinary tensor is complete. A same-stream successor

85 - # may overlap the mailbox cleanup below.

86 - cute.arch.griddepcontrol_launch_dependents()

87 -

88 - fragment = thread

89 - while fragment < fragments:

90 - element = fragment * VEC_BF16

91 - source = cute.make_ptr(

92 - cutlass.BFloat16,

93 - (symmetric_mailbox.iterator + element).llvm_ptr,

94 - cute.AddressSpace.gmem,

95 - assumed_align=16,

96 - )

88 store_lamport_sentinel_128(source)

89 fragment = fragment + stride

vllm/models/kimi_k3/nvidia/ops/cute_dsl/latent_moe_tail/primitives.py

+183

@@ -59,6 +59,31 @@ def to_cute_dynamic_m(

59 )

62 +@dsl_user_op

63 +def fma_f32_bf16(

64 + a: BFloat16,

65 + b: BFloat16,

66 + acc: Float32,

67 + *,

68 + loc=None,

69 + ip=None,

70 +) -> Float32:

71 + a_bits = llvm.bitcast(T.i16(), a.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)

72 + b_bits = llvm.bitcast(T.i16(), b.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)

73 + result = llvm.inline_asm(

74 + T.f32(),

75 + [a_bits, b_bits, acc.ir_value(loc=loc, ip=ip)],

76 + "fma.rn.f32.bf16 $0, $1, $2, $3;",

77 + "=f,h,h,f",

78 + has_side_effects=False,

79 + is_align_stack=False,

80 + asm_dialect=llvm.AsmDialect.AD_ATT,

81 + loc=loc,

82 + ip=ip,

83 + )

84 + return Float32(result)

85 +

86 +

87 @dsl_user_op

88 def load_global_u32x4(

89 pointer: cute.Pointer,

@@ -96,6 +121,104 @@ def load_global_u32x4(

121 return cute.TensorSSA(packed, 4, Uint32)

122

123

124 +def _make_top16_bf16_finalize_asm() -> str:

125 + lines = [

126 + "{",

127 + ".reg .pred _valid<16>;",

128 + ".reg .s32 _index<16>;",

129 + ".reg .b32 _weight_words<8>;",

130 + ".reg .b16 _weights<16>;",

131 + ".reg .b32 _data<64>;",

132 + ".reg .b16 _lo, _hi;",

133 + ".reg .f32 _acc<8>;",

134 + ".reg .u64 _address<16>;",

135 + "ld.global.v4.u32 {_index0, _index1, _index2, _index3}, [$5];",

136 + "ld.global.v4.u32 {_index4, _index5, _index6, _index7}, [$5+16];",

137 + "ld.global.v4.u32 {_index8, _index9, _index10, _index11}, [$5+32];",

138 + "ld.global.v4.u32 {_index12, _index13, _index14, _index15}, [$5+48];",

139 + "ld.global.v4.u32 "

140 + "{_weight_words0, _weight_words1, _weight_words2, "

141 + "_weight_words3}, [$6];",

142 + "ld.global.v4.u32 "

143 + "{_weight_words4, _weight_words5, _weight_words6, "

144 + "_weight_words7}, [$6+16];",

145 + ]

146 + for pair in range(8):

147 + lines.append(

148 + f"mov.b32 {{_weights{2 * pair}, _weights{2 * pair + 1}}}, "

149 + f"_weight_words{pair};"

150 + )

151 + for element in range(8):

152 + lines.append(f"mov.f32 _acc{element}, 0f00000000;")

153 + for route in range(16):

154 + data = 4 * route

155 + lines.extend(

156 + [

157 + f"setp.ge.s32 _valid{route}, _index{route}, 0;",

158 + f"@_valid{route} mad.wide.s32 _address{route}, "

159 + f"_index{route}, 7168, $4;",

160 + f"@_valid{route} ld.global.v4.u32 "

161 + f"{{_data{data}, _data{data + 1}, _data{data + 2}, "

162 + f"_data{data + 3}}}, [_address{route}];",

163 + ]

164 + )

165 + for route in range(16):

166 + for pair in range(4):

167 + data = 4 * route + pair

168 + element = 2 * pair

169 + lines.extend(

170 + [

171 + f"@_valid{route} mov.b32 {{_lo, _hi}}, _data{data};",

172 + f"@_valid{route} fma.rn.f32.bf16 _acc{element}, _lo, "

173 + f"_weights{route}, _acc{element};",

174 + f"@_valid{route} fma.rn.f32.bf16 _acc{element + 1}, _hi, "

175 + f"_weights{route}, _acc{element + 1};",

176 + ]

177 + )

178 + for pair in range(4):

179 + lines.append(f"cvt.rn.bf16x2.f32 ${pair}, _acc{2 * pair + 1}, _acc{2 * pair};")

180 + lines.append("}")

181 + return "\n".join(lines)

182 +

183 +

184 +_TOP16_BF16_FINALIZE_ASM = _make_top16_bf16_finalize_asm()

185 +

186 +

187 +@dsl_user_op

188 +def finalize_top16_bf16(

189 + gemm2_vector: cute.Pointer,

190 + route_indices: cute.Pointer,

191 + route_weights: cute.Pointer,

192 + *,

193 + loc=None,

194 + ip=None,

195 +):

196 + """Finalize one Kimi K3 top-16 vector with wide metadata loads."""

197 +

198 + addresses = [

199 + pointer.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip)

200 + for pointer in (gemm2_vector, route_indices, route_weights)

201 + ]

202 + out = llvm.inline_asm(

203 + llvm.StructType.get_literal([T.i32()] * 4),

204 + addresses,

205 + _TOP16_BF16_FINALIZE_ASM,

206 + "=r,=r,=r,=r,l,l,l",

207 + has_side_effects=False,

208 + is_align_stack=False,

209 + asm_dialect=llvm.AsmDialect.AD_ATT,

210 + loc=loc,

211 + ip=ip,

212 + )

213 + packed = vector.from_elements(

214 + ir.VectorType.get([4], T.i32(), loc=loc),

215 + [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(4)],

216 + loc=loc,

217 + ip=ip,

218 + )

219 + return cute.TensorSSA(packed, 4, Uint32)

220 +

221 +

222 @dsl_user_op

223 def store_global_u32x4(

224 address: Int64,

@@ -168,6 +291,24 @@ def store_global_u32(

291 )

292

293

294 +@dsl_user_op

295 +def stmc_bf16x8(address: Int64, packed, *, loc=None, ip=None) -> None:

296 + """Publish eight BF16 values through an NVLS multicast mapping."""

297 +

298 + words = [packed[i].ir_value(loc=loc, ip=ip) for i in range(4)]

299 + llvm.inline_asm(

300 + None,

301 + [address.ir_value(loc=loc, ip=ip), *words],

302 + "multimem.st.relaxed.sys.global.v4.bf16x2 [$0], {$1, $2, $3, $4};",

303 + "l,r,r,r,r",

304 + has_side_effects=True,

305 + is_align_stack=False,

306 + asm_dialect=llvm.AsmDialect.AD_ATT,

307 + loc=loc,

308 + ip=ip,

309 + )

310 +

311 +

312 @dsl_user_op

313 def store_lamport_sentinel_128(pointer: cute.Pointer, *, loc=None, ip=None) -> None:

314 """Reset one Lamport fragment to four FP32 negative-zero bit patterns."""

@@ -278,6 +419,48 @@ def store_shared_cluster_f32(

419 )

420

421

422 +@dsl_user_op

423 +def load_shared_f32x2(pointer: cute.Pointer, *, loc=None, ip=None):

424 + """Load two aligned FP32 DSM partials from local shared memory."""

425 +

426 + address = Int32(pointer.toint(loc=loc, ip=ip))

427 + out = llvm.inline_asm(

428 + llvm.StructType.get_literal([T.f32()] * 2),

429 + [address.ir_value(loc=loc, ip=ip)],

430 + "ld.shared.v2.f32 {$0, $1}, [$2];",

431 + "=f,=f,r",

432 + has_side_effects=True,

433 + is_align_stack=False,

434 + asm_dialect=llvm.AsmDialect.AD_ATT,

435 + loc=loc,

436 + ip=ip,

437 + )

438 + return tuple(

439 + Float32(llvm.extractvalue(T.f32(), out, [i], loc=loc, ip=ip)) for i in range(2)

440 + )

441 +

442 +

443 +@dsl_user_op

444 +def load_shared_f32x4(pointer: cute.Pointer, *, loc=None, ip=None):

445 + """Load four aligned FP32 DSM partials from local shared memory."""

446 +

447 + address = Int32(pointer.toint(loc=loc, ip=ip))

448 + out = llvm.inline_asm(

449 + llvm.StructType.get_literal([T.f32()] * 4),

450 + [address.ir_value(loc=loc, ip=ip)],

451 + "ld.shared.v4.f32 {$0, $1, $2, $3}, [$4];",

452 + "=f,=f,=f,=f,r",

453 + has_side_effects=True,

454 + is_align_stack=False,

455 + asm_dialect=llvm.AsmDialect.AD_ATT,

456 + loc=loc,

457 + ip=ip,

458 + )

459 + return tuple(

460 + Float32(llvm.extractvalue(T.f32(), out, [i], loc=loc, ip=ip)) for i in range(4)

461 + )

462 +

463 +

464 @dsl_user_op

465 def packed_u32x4_to_bf16x8(packed, *, loc=None, ip=None):

466 target_type = ir.VectorType.get([VEC_BF16], BFloat16.mlir_type, loc=loc)

vllm/models/kimi_k3/nvidia/ops/latent_moe_tail.py

+1 -1

@@ -22,7 +22,7 @@ _GEMM_CLUSTER_MN = (1, 8)

22 _B_PRIME_STAGES = 2

23 _COLLECTIVE_TOKEN_CTAS = 32

24 _LAMPORT_COPY_CTAS = 32

25 -_LAMPORT_COPY_THREADS = 224

25 +_LAMPORT_COPY_THREADS = 128

26 _SUPPORTED_TP_SIZES = (8, 16)