From cea96cac251b05aacb115835aa1704d33ca05fc1 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Thu, 17 Sep 2026 07:47:34 +0000 Subject: [PATCH 01/11] support mxfp8 Signed-off-by: Thien Tran --- tests/kernels/test_kimi_k3_gemm_rs_ar.py | 112 +++++- vllm/cute_utils/_tcgen05.py | 71 ++++ .../kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py | 354 +++++++++++++++--- 3 files changed, 478 insertions(+), 59 deletions(-) diff --git a/tests/kernels/test_kimi_k3_gemm_rs_ar.py b/tests/kernels/test_kimi_k3_gemm_rs_ar.py index 7a887c501f10..dcfd024b7029 100644 --- a/tests/kernels/test_kimi_k3_gemm_rs_ar.py +++ b/tests/kernels/test_kimi_k3_gemm_rs_ar.py @@ -32,15 +32,16 @@ def _reference( world_size: int, group: dist.ProcessGroup, all_reduce: bool, + out_dtype: torch.dtype | None = None, ) -> torch.Tensor: M = x.shape[0] padded_M = (M + world_size - 1) // world_size * world_size partial = torch.empty( (padded_M, weight.shape[0]), - dtype=x.dtype, + dtype=out_dtype or x.dtype, device=x.device, ) - torch.mm(x, weight.T, out=partial[:M]) + partial[:M] = torch.mm(x, weight.T) if padded_M > M: partial[M:].zero_() @@ -85,6 +86,7 @@ def _run_mode( rank: int, world_size: int, weights: dict[int, torch.Tensor], + weights_mxfp8: dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]], ) -> None: # cute_dsl is unavailable off CUDA, so import it only inside the GPU worker. from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import GemmRsAr @@ -96,8 +98,8 @@ def _run_mode( ) input_generator = torch.Generator(device=device) - # Alternating shapes exercise producer-flag reuse across different grids, - # CTA-group choices, and both BN=128 and BN=256 dispatches. + # Alternating shapes exercise producer-flag reuse across different grids + # and CTA-group choices. for M, K in (*_SHAPES, *_SHAPES[::-1]): input_generator.manual_seed(2000 + M + K) x = torch.randn( @@ -112,6 +114,34 @@ def _run_mode( torch.accelerator.synchronize(device) _assert_valid_rows_close(actual, expected, M, rank, all_reduce) + # MXFP8 dispatch over the same shapes. The wrapper quantizes activations + # internally; mirror that quantization for the reference. FP8 values times + # a power-of-2 scale are exact in BF16, so the reference only differs from + # the kernel in accumulation order and output rounding. + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + dequant_mxfp8_to_bf16, + mxfp8_e4m3_quantize, + ) + + for M, K in (*_SHAPES, *_SHAPES[::-1]): + input_generator.manual_seed(4000 + M + K) + x = torch.randn( + M, + K, + dtype=torch.bfloat16, + device=device, + generator=input_generator, + ) + x_q, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=False) + x_ref = dequant_mxfp8_to_bf16(x_q, x_sf) + w_q, w_sf, w_ref = weights_mxfp8[K] + expected = _reference( + x_ref, w_ref, world_size, group, all_reduce, out_dtype=torch.bfloat16 + ) + actual = gemm_rs_ar(x, w_q, w_sf) + torch.accelerator.synchronize(device) + _assert_valid_rows_close(actual, expected, M, rank, all_reduce) + if all_reduce: # AR outputs must remain valid after the symmetric workspace is reused. lifetime_M, lifetime_K = 257, 768 @@ -173,8 +203,60 @@ def _run_mode( all_reduce, ) + # MXFP8 under CUDA graph capture: the wrapper runs the activation-quant + # kernel inside the captured region. + w_q, w_sf, w_ref = weights_mxfp8[graph_K] + input_generator.manual_seed(4500) + graph_x_mxfp8 = torch.randn( + graph_M, + graph_K, + dtype=torch.bfloat16, + device=device, + generator=input_generator, + ) + x_q, x_sf = mxfp8_e4m3_quantize(graph_x_mxfp8, is_sf_swizzled_layout=False) + graph_expected_mxfp8 = _reference( + dequant_mxfp8_to_bf16(x_q, x_sf), + w_ref, + world_size, + group, + all_reduce, + out_dtype=torch.bfloat16, + ) + graph_output_mxfp8 = torch.empty_like(graph_expected_mxfp8) + + with torch.cuda.stream(capture_stream): + for _ in range(3): + graph_output_mxfp8.copy_(gemm_rs_ar(graph_x_mxfp8, w_q, w_sf)) + capture_stream.synchronize() + dist.barrier(group=group) + + graph_mxfp8 = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph_mxfp8, stream=capture_stream): + graph_output_mxfp8.copy_(gemm_rs_ar(graph_x_mxfp8, w_q, w_sf)) + torch.cuda.current_stream().wait_stream(capture_stream) + dist.barrier(group=group) + + for _ in range(3): + graph_mxfp8.replay() + torch.accelerator.synchronize(device) + _assert_valid_rows_close( + graph_output_mxfp8, + graph_expected_mxfp8, + graph_M, + rank, + all_reduce, + ) + dist.barrier(group=group) - del gemm_rs_ar, capture_stream, graph, graph_output + del ( + gemm_rs_ar, + capture_stream, + graph, + graph_output, + graph_mxfp8, + graph_output_mxfp8, + ) def _worker(local_rank: int, world_size: int, master_port: int) -> None: @@ -203,6 +285,13 @@ def _worker(local_rank: int, world_size: int, master_port: int) -> None: weight_generator = torch.Generator(device=device) weights = {} + weights_mxfp8 = {} + + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + dequant_mxfp8_to_bf16, + mxfp8_e4m3_quantize, + ) + for K in {K for _, K in _SHAPES}: weight_generator.manual_seed(1000 + rank * 10 + K) weights[K] = torch.randn( @@ -213,6 +302,18 @@ def _worker(local_rank: int, world_size: int, master_port: int) -> None: generator=weight_generator, ) + weight_generator.manual_seed(5000 + rank * 10 + K) + w = torch.randn( + _N, + K, + dtype=torch.bfloat16, + device=device, + generator=weight_generator, + ) + w_q, w_sf = mxfp8_e4m3_quantize(w, is_sf_swizzled_layout=True) + _, w_sf_2d = mxfp8_e4m3_quantize(w, is_sf_swizzled_layout=False) + weights_mxfp8[K] = (w_q, w_sf, dequant_mxfp8_to_bf16(w_q, w_sf_2d)) + # Production binds one mode per worker; exercise both mode-bound instances # sequentially in this test without implying that both are initialized. for all_reduce in (False, True): @@ -223,6 +324,7 @@ def _worker(local_rank: int, world_size: int, master_port: int) -> None: rank=rank, world_size=world_size, weights=weights, + weights_mxfp8=weights_mxfp8, ) cleanup_dist_env_and_memory() diff --git a/vllm/cute_utils/_tcgen05.py b/vllm/cute_utils/_tcgen05.py index 325c30498c76..4ab75cc0cc3b 100644 --- a/vllm/cute_utils/_tcgen05.py +++ b/vllm/cute_utils/_tcgen05.py @@ -78,6 +78,10 @@ def make_sdesc_128B_swizzle(LBO: int): return Uint64((LBO >> 4 << 16) | (SBO >> 4 << 32) | (1 << 46) | (2 << 61)) +def make_mxfp8_idesc(MMA_M: int, MMA_N: int): + return Uint32((MMA_N >> 3 << 17) | (1 << 23) | (MMA_M >> 7 << 27)) + + @dsl_user_op def mma_f16( d_tmem, @@ -130,6 +134,73 @@ def mma_ts_f16( ) +@dsl_user_op +def mma_mxfp8( + d_tmem, + a_desc, + b_desc, + idesc, + sfa_tmem, + sfb_tmem, + enable_input_d, + cta_group: int = 1, + *, + loc=None, + ip=None, +) -> None: + with cute.arch.elect_one(): + nvvm.tcgen05_mma_block_scale( + nvvm.Tcgen05MMAKind.MXF8F6F4, + NVVM_CTA_GROUP_MAP[cta_group], + _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip), + Uint64(a_desc).ir_value(loc=loc, ip=ip), + Uint64(b_desc).ir_value(loc=loc, ip=ip), + Int32(idesc).ir_value(loc=loc, ip=ip), + Boolean(enable_input_d).ir_value(loc=loc, ip=ip), + _make_tmem_llvm_ptr(sfa_tmem, loc=loc, ip=ip), + _make_tmem_llvm_ptr(sfb_tmem, loc=loc, ip=ip), + scale_vec_size=nvvm.Tcgen05MMAScaleVecSize.X1, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cp( + tmem, + sdesc, + shape: str, + mcast: str, + cta_group: int = 1, + *, + loc=None, + ip=None, +) -> None: + SHAPE_MAP = { + "128x256b": nvvm.Tcgen05CpShape.SHAPE_128x256b, + "4x256b": nvvm.Tcgen05CpShape.SHAPE_4x256b, + "128x128b": nvvm.Tcgen05CpShape.SHAPE_128x128b, + "64x128b": nvvm.Tcgen05CpShape.SHAPE_64x128b, + "32x128b": nvvm.Tcgen05CpShape.SHAPE_32x128b, + } + MCAST_MAP = { + "none": nvvm.Tcgen05CpMulticast.NONE, + "warpx2::02_13": nvvm.Tcgen05CpMulticast.WARPX2_02_13, + "warpx2::01_23": nvvm.Tcgen05CpMulticast.WARPX2_01_23, + "warpx4": nvvm.Tcgen05CpMulticast.WARPX4, + } + with cute.arch.elect_one(): + nvvm.tcgen05_cp( + SHAPE_MAP[shape], + _make_tmem_llvm_ptr(tmem, loc=loc, ip=ip), + Uint64(sdesc).ir_value(loc=loc, ip=ip), + group=NVVM_CTA_GROUP_MAP[cta_group], + multicast=MCAST_MAP[mcast], + loc=loc, + ip=ip, + ) + + @dsl_user_op def commit(mbar, cta_mask=None, cta_group: int = 1, *, loc=None, ip=None): mbar_llvm = mbar.to_llvm_ptr(loc=loc, ip=ip) diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py index f8a90cc1abce..64e8c71579b9 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""SM100 BF16 GEMM with fused tensor-parallel reduce-scatter/all-reduce.""" +"""SM100 BF16/MXFP8 GEMM with fused tensor-parallel reduce-scatter/all-reduce.""" # Based on CUTLASS's Blackwell distributed GEMM-RS example at dcf215a. # See https://github.com/NVIDIA/cutlass/issues/3117 for memory semantics. +# The MXFP8 path follows the block-scaled SM100 GEMM template from +# https://github.com/gau-nernst/gn-kernels (sm100_mm_mxfp8.py). from functools import cache @@ -12,7 +14,17 @@ import torch import torch.distributed._symmetric_memory as symm_mem from cuda.bindings.driver import CUstream -from cutlass import BFloat16, Int32, Int64, Uint16, cute, utils +from cutlass import ( + BFloat16, + Float8E4M3FN, + Float8E8M0FNU, + Int32, + Int64, + Uint16, + Uint64, + cute, + utils, +) from cutlass._mlir import ir from cutlass._mlir.dialects import llvm, nvvm, vector from cutlass.cute.nvgpu import cpasync, tcgen05 @@ -88,24 +100,31 @@ def multimem_st_16B(dst: cute.Tensor, value: cute.Tensor, *, loc=None, ip=None) ) -class Sm100GemmRsArBF16: +class Sm100GemmRsAr: def __init__( self, rank: int, num_ranks: int, - BN: int = 128, cta_group: int = 1, all_reduce: bool = False, + dtype: str = "bf16", ) -> None: + assert dtype in ("bf16", "mxfp8") self.rank = rank self.num_ranks = num_ranks - BM, BK = 128, 64 + BM = 128 + BN = 128 + BK = {"bf16": 64, "mxfp8": 128}[dtype] # 128B self.cta_tile = (BM, BN, BK) self.cta_group = cta_group self.all_reduce = all_reduce + self.dtype = dtype smem_bytes = get_smem_capacity_in_bytes() - self.stage_size = (BM + (BN // cta_group)) * BK * 2 + if dtype == "bf16": + self.stage_size = (BM + (BN // cta_group)) * BK * 2 + elif dtype == "mxfp8": + self.stage_size = (BM + (BN // cta_group)) * BK + (BM + BN) * (BK // 32) self.num_stages = smem_bytes // self.stage_size @cute.jit @@ -124,11 +143,36 @@ def prepare_tma( layout = cute.make_composed_layout(swizzle_128b, 0, layout) return cpasync.make_tiled_tma_atom(tma_op, tensor, layout, (BM, BK)) + @cute.jit + def prepare_sf_tma( + self, + SF: cute.Tensor, + R: Int32, + Ksf: Int32, + BR: cutlass.Constexpr, + BKsf: cutlass.Constexpr, + ) -> cpasync.TmaInfo: + tma_group = ( + tcgen05.CtaGroup.TWO if self.cta_group == 2 else tcgen05.CtaGroup.ONE + ) + tma_op = cpasync.CopyBulkTensorTileG2SOp(cta_group=tma_group) + # NVIDIA SF gmem layout: an [R, Ksf] scale grid padded to R % 128 == 0 + # and permuted as [R/128, Ksf/4, 32, 4, 4]. The trailing [32, 4, 4] + # atom is the opaque 512B block tcgen05 expects, so treat it as bytes. + # The Int64 recast keeps the TMA boxDim <= 256. + g_layout = cute.make_layout((Ksf * 128, cute.ceil_div(R, 128))) + SF_i64 = cute.recast_tensor(cute.make_tensor(SF.iterator, g_layout), Int64) + tiler = (128 * BKsf // 8, BR // 128) + layout = cute.make_layout((*tiler, self.num_stages)) + return cpasync.make_tiled_tma_atom(tma_op, SF_i64, layout, tiler) + @cute.jit def __call__( self, A: cute.Tensor, B: cute.Tensor, + SFA: cute.Tensor, + SFB: cute.Tensor, partial_uc: cute.Tensor, partial_mc_ptr: cute.Pointer, output: cute.Tensor, @@ -138,10 +182,17 @@ def __call__( grid_size: Int32, stream: CUstream, ) -> None: + M, K = A.shape N = B.shape[0] BM, BN, BK = self.cta_tile A_tma = self.prepare_tma(A, BM, BK) B_tma = self.prepare_tma(B, BN // self.cta_group, BK) + if cutlass.const_expr(self.dtype == "mxfp8"): + SFA_tma = self.prepare_sf_tma(SFA, M, K // 32, BM, BK // 32) + SFB_tma = self.prepare_sf_tma(SFB, N, K // 32, BN, BK // 32) + else: + SFA_tma = None + SFB_tma = None padded_M = partial_uc.shape[0] partial_mc = cute.make_tensor( partial_mc_ptr, @@ -158,6 +209,8 @@ def __call__( self.kernel( A_tma, B_tma, + SFA_tma, + SFB_tma, partial_uc, partial_mc, output, @@ -171,6 +224,8 @@ def kernel( self, A_tma: cpasync.TmaInfo, B_tma: cpasync.TmaInfo, + SFA_tma: cpasync.TmaInfo, + SFB_tma: cpasync.TmaInfo, partial_uc: cute.Tensor, partial_mc: cute.Tensor, output: cute.Tensor, @@ -190,25 +245,33 @@ def kernel( is_2cta = cta_group == 2 cta_rank = raw_bid % self.cta_group - num_tmem_stages = 512 // BN + + # MXFP8 reserves the last 16 tmem columns (496-512) for the A/B scale + # factors. + num_tmem_stages = (496 if self.dtype == "mxfp8" else 512) // BN smem = utils.SmemAllocator() sA = smem.allocate_tensor( - BFloat16, + A_tma.atom.value_type, A_tma.smem_layout.outer, byte_alignment=128, swizzle=A_tma.smem_layout.inner, ) sB = smem.allocate_tensor( - BFloat16, + B_tma.atom.value_type, B_tma.smem_layout.outer, byte_alignment=128, swizzle=B_tma.smem_layout.inner, ) + if cutlass.const_expr(self.dtype == "mxfp8"): + sSFA = smem.allocate_tensor(Int64, SFA_tma.smem_layout, byte_alignment=128) + sSFB = smem.allocate_tensor(Int64, SFB_tma.smem_layout, byte_alignment=128) + tma_full_mbar = smem.allocate_array(Int64, num_stages) tma_empty_mbar = smem.allocate_array(Int64, num_stages) tmem_full_mbar = smem.allocate_array(Int64, num_tmem_stages) tmem_empty_mbar = smem.allocate_array(Int64, num_tmem_stages) + taddr = smem.allocate(Int32, 4) # Named barriers @@ -238,6 +301,9 @@ def kernel( elif warp_id == 1: cpasync.prefetch_descriptor(A_tma.atom) cpasync.prefetch_descriptor(B_tma.atom) + if cutlass.const_expr(self.dtype == "mxfp8"): + cpasync.prefetch_descriptor(SFA_tma.atom) + cpasync.prefetch_descriptor(SFB_tma.atom) if cutlass.const_expr(is_2cta): cute.arch.cluster_arrive_relaxed() @@ -261,12 +327,22 @@ def kernel( # [(BM, BK), (M/BM, K/BK)] gA_tiles = cute.zipped_divide(A_tma.tma_tensor, (BM, BK)) gB_tiles = cute.zipped_divide(B_tma.tma_tensor, (BN // cta_group, BK)) + if cutlass.const_expr(self.dtype == "mxfp8"): + BKsf = BK // 32 + gSFA_tiles = cute.zipped_divide( + SFA_tma.tma_tensor, (128 * BKsf // 8, BM // 128) + ) + gSFB_tiles = cute.zipped_divide( + SFB_tma.tma_tensor, (128 * BKsf // 8, BN // 128) + ) for bid in range(raw_bid, total_tiles, num_bids): bid_m = bid % grid_m bid_n = bid // grid_m if cutlass.const_expr(cta_group == 2): - bid_n = bid_n * cta_group + cta_rank + bid_n_b = bid_n * cta_group + cta_rank + else: + bid_n_b = bid_n for iter_k in cutlass.range(cute.ceil_div(K, BK), unroll=1): mbar = tma_full_mbar_ + tma_stage @@ -282,10 +358,23 @@ def kernel( ) simple_tma_copy( B_tma.atom, - gB_tiles[None, (bid_n, iter_k)], + gB_tiles[None, (bid_n_b, iter_k)], sB[None, None, tma_stage], mbar, ) + if cutlass.const_expr(self.dtype == "mxfp8"): + simple_tma_copy( + SFA_tma.atom, + gSFA_tiles[None, (iter_k, bid_m)], + sSFA[None, None, tma_stage], + mbar, + ) + simple_tma_copy( + SFB_tma.atom, + gSFB_tiles[None, (iter_k, bid_n)], + sSFB[None, None, tma_stage], + mbar, + ) tma_stage = (tma_stage + 1) % num_stages if tma_stage == 0: @@ -303,8 +392,15 @@ def kernel( MMA_M = BM * cta_group MMA_N = BN - idesc = _tcgen05.make_bf16_idesc(MMA_M, MMA_N) sdesc = _tcgen05.make_sdesc_128B_swizzle(0) + if cutlass.const_expr(self.dtype == "mxfp8"): + # SF smem descriptor: no swizzle, SBO = 8 * 16 bytes. + sdesc_sf = Uint64((8 << 32) | (1 << 46)) + sfa_tmem = 496 + sfb_tmem = sfa_tmem + 4 + idesc = _tcgen05.make_mxfp8_idesc(MMA_M, MMA_N) + else: + idesc = _tcgen05.make_bf16_idesc(MMA_M, MMA_N) multicast_mask = Uint16((1 << self.cta_group) - 1) for bid in range(raw_bid, total_tiles, num_bids): @@ -319,17 +415,54 @@ def kernel( b_addr = sB[None, None, tma_stage].iterator.toint() a_desc = sdesc | (a_addr >> 4) b_desc = sdesc | (b_addr >> 4) + if cutlass.const_expr(self.dtype == "mxfp8"): + sfa_desc = sdesc_sf | ( + sSFA[None, None, tma_stage].iterator.toint() >> 4 + ) + sfb_desc = sdesc_sf | ( + sSFB[None, None, tma_stage].iterator.toint() >> 4 + ) cute.arch.mbarrier_wait( tma_full_mbar + tma_stage, tma_full_parity ) _tcgen05.fence_after_thread_sync() - for mma_k in cutlass.range_constexpr(BK // 16): - enable_d = iter_k > 0 or mma_k > 0 - _tcgen05.mma_f16( - d_tmem, a_desc, b_desc, idesc, enable_d, cta_group + # load SF from smem->tmem + if cutlass.const_expr(self.dtype == "mxfp8"): + _tcgen05.cp( + sfa_tmem, sfa_desc, "32x128b", "warpx4", cta_group ) + for j in cutlass.range_constexpr(BN // 128): + _tcgen05.cp( + sfb_tmem + j * 4, + sfb_desc, + "32x128b", + "warpx4", + cta_group, + ) + sfb_desc += (128 * 4) >> 4 + + # 4 = 128B / 32B + for mma_k in cutlass.range_constexpr(4): + enable_d = iter_k > 0 or mma_k > 0 + if cutlass.const_expr(self.dtype == "mxfp8"): + # select SF IDs + idesc_k = idesc + ((mma_k << 4) | (mma_k << 29)) + _tcgen05.mma_mxfp8( + d_tmem, + a_desc, + b_desc, + idesc_k, + sfa_tmem, + sfb_tmem, + enable_d, + cta_group, + ) + else: + _tcgen05.mma_f16( + d_tmem, a_desc, b_desc, idesc, enable_d, cta_group + ) a_desc += 32 >> 4 b_desc += 32 >> 4 _tcgen05.commit( @@ -638,9 +771,9 @@ def release_gemm_tile( def compile( rank: int, num_ranks: int, - BN: int, cta_group: int, all_reduce: bool = False, + dtype: str = "bf16", ): M = cute.sym_int() padded_M = cute.sym_int() @@ -649,12 +782,37 @@ def compile( local_M = cute.sym_int() num_flags = cute.sym_int() + ab_dtype = Float8E4M3FN if dtype == "mxfp8" else BFloat16 + # Row stride must keep rows 16B-aligned: 8 BF16 or 16 FP8 elements. + row_div = 16 if dtype == "mxfp8" else 8 A = make_fake_tensor( - BFloat16, (M, K), (cute.sym_int64(divisibility=8), 1), assumed_align=16 + ab_dtype, + (M, K), + (cute.sym_int64(divisibility=row_div), 1), + assumed_align=16, ) B = make_fake_tensor( - BFloat16, (N, K), (cute.sym_int64(divisibility=8), 1), assumed_align=16 + ab_dtype, + (N, K), + (cute.sym_int64(divisibility=row_div), 1), + assumed_align=16, ) + if dtype == "mxfp8": + SFA = make_fake_tensor( + Float8E8M0FNU, + (cute.sym_int(divisibility=512),), + (1,), + assumed_align=16, + ) + SFB = make_fake_tensor( + Float8E8M0FNU, + (cute.sym_int(divisibility=512),), + (1,), + assumed_align=16, + ) + else: + SFA = None + SFB = None partial = make_fake_tensor( BFloat16, (padded_M, N), @@ -677,13 +835,15 @@ def compile( peer_flag_ptr = nullptr(Int64, cute.AddressSpace.gmem, assumed_align=8) stream = make_fake_stream(use_tvm_ffi_env_stream=True) - kernel = Sm100GemmRsArBF16( - rank, num_ranks, BN, cta_group, all_reduce=all_reduce + kernel = Sm100GemmRsAr( + rank, num_ranks, cta_group, all_reduce=all_reduce, dtype=dtype ) return cute.compile( kernel, A, B, + SFA, + SFB, partial, partial_mc_ptr, output, @@ -727,6 +887,13 @@ def __init__(self, *, max_M: int, N: int, all_reduce: bool = False) -> None: self.N = N self.device = device self.all_reduce = all_reduce + # dtypes with at least one compatible projection, as recorded by + # can_run; warmup compiles the dispatches for each of them. + self.dtypes: set[str] = {"bf16"} + # MXFP8 weight scales registered by can_run, keyed by the weight's + # data_ptr (transpose views share it), so callers only pass the + # weight to __call__. + self._weight_scales: dict[int, torch.Tensor] = {} self.partial = symm_mem.empty((max_M, N), dtype=torch.bfloat16, device=device) self.partial_handle = symm_mem.rendezvous(self.partial, group) @@ -767,20 +934,69 @@ def __init__(self, *, max_M: int, N: int, all_reduce: bool = False) -> None: def can_run(self, linear: LinearBase) -> bool: # Validate projection-invariant requirements once during model init. - # only supports BF16 for now - if not isinstance(linear.quant_method, UnquantizedLinearMethod): + if isinstance(linear.quant_method, UnquantizedLinearMethod): + w = linear.weight + if w.ndim != 2: + return False + K = w.shape[1] + return ( + w.shape == (self.N, K) + and K % 64 == 0 + and w.dtype == torch.bfloat16 + and w.device == self.device + and w.is_contiguous() + ) + return self._can_run_mxfp8(linear) + + def _can_run_mxfp8(self, linear: LinearBase) -> bool: + # Only the FlashInfer MXFP8 kernels keep the [N, K] row-major weight + # plus a flat swizzled scale, which is what the kernel consumes. + from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( + FlashInferCutedslMxfp8LinearKernel, + FlashInferCutlassMxfp8LinearKernel, + ) + + kernel = getattr(linear.quant_method, "kernel", None) + if not isinstance( + kernel, + (FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel), + ): return False - w = linear.weight - if w.ndim != 2: + operands = self._mxfp8_operands(linear) + if operands is None: return False - K = w.shape[1] - return ( - w.shape == (self.N, K) - and K % 64 == 0 - and w.dtype == torch.bfloat16 - and w.device == self.device - and w.is_contiguous() - ) + w, w_sf, K = operands + if K % 128 != 0: + return False + self._weight_scales[w.data_ptr()] = w_sf + self.dtypes.add("mxfp8") + return True + + def _mxfp8_operands( + self, linear: LinearBase + ) -> tuple[torch.Tensor, torch.Tensor, int] | None: + """Return (weight [N, K] FP8, flat swizzled scale, K), or None.""" + w = linear.weight + w_sf = getattr(linear, "weight_scale", None) + if ( + w.ndim != 2 + or w.dtype != torch.float8_e4m3fn + or w.device != self.device + or w_sf is None + or w_sf.ndim != 1 + or w_sf.dtype != torch.uint8 + or not w_sf.is_contiguous() + ): + return None + if w.shape[0] == self.N and w.is_contiguous(): + pass + elif w.shape[1] == self.N and w.t().is_contiguous(): + # FlashInfer CuTe-DSL stores the weight as a [K, N] column-major + # view of the row-major [N, K] buffer. + w = w.t() + else: + return None + return w, w_sf, w.shape[1] def warn_incompatible_projection(self) -> None: logger.warning_once( @@ -794,18 +1010,46 @@ def should_run(self, x: torch.Tensor) -> bool: # supported but faster on the existing LL path. return x.shape[0] >= 128 - def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: + def __call__( + self, + x: torch.Tensor, + w: torch.Tensor, + w_sf: torch.Tensor | None = None, + ) -> torch.Tensor: assert x.ndim == 2 M, K = x.shape assert 0 < M <= self.max_M - assert w.shape == (self.N, K) and K % 64 == 0 - assert w.dtype == torch.bfloat16 - assert w.device == self.device - assert w.is_contiguous() assert x.dtype == torch.bfloat16 assert x.device == self.device assert x.is_contiguous() - N = w.shape[0] + N = self.N + dtype = "mxfp8" if w.dtype == torch.float8_e4m3fn else "bf16" + x_sf = None + if dtype == "mxfp8": + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + ) + + if w_sf is None: + w_sf = self._weight_scales.get(w.data_ptr()) + assert w_sf is not None + assert w_sf.ndim == 1 and w_sf.is_contiguous() + if w.shape[0] != N: + assert w.shape[1] == N + w = w.t() + assert w.shape == (N, K) and K % 128 == 0 + assert w.is_contiguous() + x, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True) + # The kernel declares the scales as E8M0; the stored/quantized + # tensors carry the same bits as uint8. + x_sf = x_sf.view(torch.float8_e8m0fnu) + w_sf = w_sf.view(torch.float8_e8m0fnu) + else: + assert w_sf is None + assert w.shape == (N, K) and K % 64 == 0 + assert w.dtype == torch.bfloat16 + assert w.device == self.device + assert w.is_contiguous() padded_M = (M + self.world_size - 1) // self.world_size padded_M *= self.world_size local_M = padded_M // self.world_size @@ -814,10 +1058,8 @@ def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: # Avoid padding small odd grids; 2-CTA wins consistently for M >= 1024. cta_group = 2 if M >= 1024 or grid_m % 2 == 0 else 1 grid_m = (grid_m + cta_group - 1) // cta_group * cta_group - BN = 256 if M * K >= 24 * 1024 * 1024 else 128 - assert N % BN == 0 - num_tiles = grid_m * (N // BN) + num_tiles = grid_m * (N // 128) num_ctas = min(num_tiles, self.num_sms) num_ctas = num_ctas // cta_group * cta_group assert self.flags.numel() >= num_tiles + num_ctas @@ -825,16 +1067,18 @@ def __call__(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor: output = None if not self.all_reduce: output = torch.empty((local_M, N), dtype=torch.bfloat16, device=self.device) - compiled = Sm100GemmRsArBF16.compile( + compiled = Sm100GemmRsAr.compile( self.rank, self.world_size, - BN, cta_group, self.all_reduce, + dtype=dtype, ) compiled( x, w, + x_sf, + w_sf, self.partial[:padded_M], self.partial_mc_ptr, output, @@ -877,16 +1121,18 @@ def warmup_gemm_rs_ar() -> int: if _gemm_rs_ar is None: return 0 # Keep these profiles in sync with the dispatch in GemmRsAr.__call__. - profiles = ((128, 1), (128, 2), (256, 2)) - for BN, cta_group in profiles: - Sm100GemmRsArBF16.compile( - _gemm_rs_ar.rank, - _gemm_rs_ar.world_size, - BN, - cta_group, - _gemm_rs_ar.all_reduce, - ) - return len(profiles) + count = 0 + for dtype in sorted(_gemm_rs_ar.dtypes): + for cta_group in (1, 2): + Sm100GemmRsAr.compile( + _gemm_rs_ar.rank, + _gemm_rs_ar.world_size, + cta_group, + _gemm_rs_ar.all_reduce, + dtype=dtype, + ) + count += 1 + return count def get_gemm_rs_ar() -> GemmRsAr: From de1e31ca850105713cad30038011fcf09df9c5d1 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Thu, 17 Sep 2026 08:01:43 +0000 Subject: [PATCH 02/11] simplify Signed-off-by: Thien Tran --- tests/kernels/test_kimi_k3_gemm_rs_ar.py | 160 +++++------------- .../kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py | 111 +++--------- 2 files changed, 72 insertions(+), 199 deletions(-) diff --git a/tests/kernels/test_kimi_k3_gemm_rs_ar.py b/tests/kernels/test_kimi_k3_gemm_rs_ar.py index dcfd024b7029..1762afafa042 100644 --- a/tests/kernels/test_kimi_k3_gemm_rs_ar.py +++ b/tests/kernels/test_kimi_k3_gemm_rs_ar.py @@ -17,6 +17,7 @@ from vllm.utils.system_utils import update_environment_variables _SHAPES = ( + (129, 512), (129, 768), (257, 768), (1023, 4224), @@ -26,22 +27,34 @@ _N = 512 +def _round_to_mxfp8(x: torch.Tensor) -> torch.Tensor: + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + dequant_mxfp8_to_bf16, + mxfp8_e4m3_quantize, + ) + + values, scales = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=False) + return dequant_mxfp8_to_bf16(values, scales) + + def _reference( x: torch.Tensor, weight: torch.Tensor, world_size: int, group: dist.ProcessGroup, all_reduce: bool, - out_dtype: torch.dtype | None = None, + mxfp8: bool, ) -> torch.Tensor: + if mxfp8: + x = _round_to_mxfp8(x) M = x.shape[0] padded_M = (M + world_size - 1) // world_size * world_size partial = torch.empty( (padded_M, weight.shape[0]), - dtype=out_dtype or x.dtype, + dtype=x.dtype, device=x.device, ) - partial[:M] = torch.mm(x, weight.T) + torch.mm(x, weight.T, out=partial[:M]) if padded_M > M: partial[M:].zero_() @@ -86,7 +99,7 @@ def _run_mode( rank: int, world_size: int, weights: dict[int, torch.Tensor], - weights_mxfp8: dict[int, tuple[torch.Tensor, torch.Tensor, torch.Tensor]], + mxfp8: bool, ) -> None: # cute_dsl is unavailable off CUDA, so import it only inside the GPU worker. from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import GemmRsAr @@ -96,6 +109,22 @@ def _run_mode( N=_N, all_reduce=all_reduce, ) + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + ) + + operands = {} + reference_weights = {} + for K, w in weights.items(): + if mxfp8: + w_q, w_sf = mxfp8_e4m3_quantize(w, is_sf_swizzled_layout=True) + # Exercise column-major square weights as well as row-major weights. + operands[K] = (w_q.t() if K == _N else w_q, w_sf) + reference_weights[K] = _round_to_mxfp8(w) + else: + operands[K] = (w, None) + reference_weights[K] = w + input_generator = torch.Generator(device=device) # Alternating shapes exercise producer-flag reuse across different grids @@ -109,36 +138,10 @@ def _run_mode( device=device, generator=input_generator, ) - expected = _reference(x, weights[K], world_size, group, all_reduce) - actual = gemm_rs_ar(x, weights[K]) - torch.accelerator.synchronize(device) - _assert_valid_rows_close(actual, expected, M, rank, all_reduce) - - # MXFP8 dispatch over the same shapes. The wrapper quantizes activations - # internally; mirror that quantization for the reference. FP8 values times - # a power-of-2 scale are exact in BF16, so the reference only differs from - # the kernel in accumulation order and output rounding. - from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( - dequant_mxfp8_to_bf16, - mxfp8_e4m3_quantize, - ) - - for M, K in (*_SHAPES, *_SHAPES[::-1]): - input_generator.manual_seed(4000 + M + K) - x = torch.randn( - M, - K, - dtype=torch.bfloat16, - device=device, - generator=input_generator, - ) - x_q, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=False) - x_ref = dequant_mxfp8_to_bf16(x_q, x_sf) - w_q, w_sf, w_ref = weights_mxfp8[K] expected = _reference( - x_ref, w_ref, world_size, group, all_reduce, out_dtype=torch.bfloat16 + x, reference_weights[K], world_size, group, all_reduce, mxfp8 ) - actual = gemm_rs_ar(x, w_q, w_sf) + actual = gemm_rs_ar(x, *operands[K]) torch.accelerator.synchronize(device) _assert_valid_rows_close(actual, expected, M, rank, all_reduce) @@ -153,9 +156,9 @@ def _run_mode( device=device, generator=input_generator, ) - first_output = gemm_rs_ar(lifetime_x, weights[lifetime_K]) + first_output = gemm_rs_ar(lifetime_x, *operands[lifetime_K]) first_snapshot = first_output.clone() - second_output = gemm_rs_ar(-lifetime_x, weights[lifetime_K]) + second_output = gemm_rs_ar(-lifetime_x, *operands[lifetime_K]) torch.accelerator.synchronize(device) assert first_output.data_ptr() != second_output.data_ptr() torch.testing.assert_close(first_output, first_snapshot, rtol=0, atol=0) @@ -171,10 +174,11 @@ def _run_mode( ) graph_expected = _reference( graph_x, - weights[graph_K], + reference_weights[graph_K], world_size, group, all_reduce, + mxfp8, ) graph_output = torch.empty_like(graph_expected) @@ -182,13 +186,13 @@ def _run_mode( capture_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(capture_stream): for _ in range(3): - graph_output.copy_(gemm_rs_ar(graph_x, weights[graph_K])) + graph_output.copy_(gemm_rs_ar(graph_x, *operands[graph_K])) capture_stream.synchronize() dist.barrier(group=group) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=capture_stream): - graph_output.copy_(gemm_rs_ar(graph_x, weights[graph_K])) + graph_output.copy_(gemm_rs_ar(graph_x, *operands[graph_K])) torch.cuda.current_stream().wait_stream(capture_stream) dist.barrier(group=group) @@ -203,63 +207,11 @@ def _run_mode( all_reduce, ) - # MXFP8 under CUDA graph capture: the wrapper runs the activation-quant - # kernel inside the captured region. - w_q, w_sf, w_ref = weights_mxfp8[graph_K] - input_generator.manual_seed(4500) - graph_x_mxfp8 = torch.randn( - graph_M, - graph_K, - dtype=torch.bfloat16, - device=device, - generator=input_generator, - ) - x_q, x_sf = mxfp8_e4m3_quantize(graph_x_mxfp8, is_sf_swizzled_layout=False) - graph_expected_mxfp8 = _reference( - dequant_mxfp8_to_bf16(x_q, x_sf), - w_ref, - world_size, - group, - all_reduce, - out_dtype=torch.bfloat16, - ) - graph_output_mxfp8 = torch.empty_like(graph_expected_mxfp8) - - with torch.cuda.stream(capture_stream): - for _ in range(3): - graph_output_mxfp8.copy_(gemm_rs_ar(graph_x_mxfp8, w_q, w_sf)) - capture_stream.synchronize() - dist.barrier(group=group) - - graph_mxfp8 = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph_mxfp8, stream=capture_stream): - graph_output_mxfp8.copy_(gemm_rs_ar(graph_x_mxfp8, w_q, w_sf)) - torch.cuda.current_stream().wait_stream(capture_stream) - dist.barrier(group=group) - - for _ in range(3): - graph_mxfp8.replay() - torch.accelerator.synchronize(device) - _assert_valid_rows_close( - graph_output_mxfp8, - graph_expected_mxfp8, - graph_M, - rank, - all_reduce, - ) - dist.barrier(group=group) - del ( - gemm_rs_ar, - capture_stream, - graph, - graph_output, - graph_mxfp8, - graph_output_mxfp8, - ) + del gemm_rs_ar, capture_stream, graph, graph_output -def _worker(local_rank: int, world_size: int, master_port: int) -> None: +def _worker(local_rank: int, world_size: int, master_port: int, mxfp8: bool) -> None: # This module pulls in cute_dsl, which is unavailable off CUDA, so importing # it at module level would fail collection. Import it where it is used, as # `kda.py` does. @@ -285,13 +237,6 @@ def _worker(local_rank: int, world_size: int, master_port: int) -> None: weight_generator = torch.Generator(device=device) weights = {} - weights_mxfp8 = {} - - from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( - dequant_mxfp8_to_bf16, - mxfp8_e4m3_quantize, - ) - for K in {K for _, K in _SHAPES}: weight_generator.manual_seed(1000 + rank * 10 + K) weights[K] = torch.randn( @@ -302,18 +247,6 @@ def _worker(local_rank: int, world_size: int, master_port: int) -> None: generator=weight_generator, ) - weight_generator.manual_seed(5000 + rank * 10 + K) - w = torch.randn( - _N, - K, - dtype=torch.bfloat16, - device=device, - generator=weight_generator, - ) - w_q, w_sf = mxfp8_e4m3_quantize(w, is_sf_swizzled_layout=True) - _, w_sf_2d = mxfp8_e4m3_quantize(w, is_sf_swizzled_layout=False) - weights_mxfp8[K] = (w_q, w_sf, dequant_mxfp8_to_bf16(w_q, w_sf_2d)) - # Production binds one mode per worker; exercise both mode-bound instances # sequentially in this test without implying that both are initialized. for all_reduce in (False, True): @@ -324,17 +257,18 @@ def _worker(local_rank: int, world_size: int, master_port: int) -> None: rank=rank, world_size=world_size, weights=weights, - weights_mxfp8=weights_mxfp8, + mxfp8=mxfp8, ) cleanup_dist_env_and_memory() +@pytest.mark.parametrize("mxfp8", [False, True], ids=["bf16", "mxfp8"]) @pytest.mark.distributed(num_gpus=2) @pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="Kimi-K3 GEMM-RS/AR requires SM100", ) -def test_kimi_k3_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch) -> None: +def test_kimi_k3_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch, mxfp8: bool) -> None: world_size = 2 if torch.accelerator.device_count() < world_size: pytest.skip("GEMM-RS/AR requires two GPUs") @@ -344,7 +278,7 @@ def test_kimi_k3_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch) -> None: try: mp.spawn( _worker, - args=(world_size, get_open_port()), + args=(world_size, get_open_port(), mxfp8), nprocs=world_size, ) finally: diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py index 64e8c71579b9..b71bc5abed9f 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py @@ -887,13 +887,6 @@ def __init__(self, *, max_M: int, N: int, all_reduce: bool = False) -> None: self.N = N self.device = device self.all_reduce = all_reduce - # dtypes with at least one compatible projection, as recorded by - # can_run; warmup compiles the dispatches for each of them. - self.dtypes: set[str] = {"bf16"} - # MXFP8 weight scales registered by can_run, keyed by the weight's - # data_ptr (transpose views share it), so callers only pass the - # weight to __call__. - self._weight_scales: dict[int, torch.Tensor] = {} self.partial = symm_mem.empty((max_M, N), dtype=torch.bfloat16, device=device) self.partial_handle = symm_mem.rendezvous(self.partial, group) @@ -934,69 +927,19 @@ def __init__(self, *, max_M: int, N: int, all_reduce: bool = False) -> None: def can_run(self, linear: LinearBase) -> bool: # Validate projection-invariant requirements once during model init. - if isinstance(linear.quant_method, UnquantizedLinearMethod): - w = linear.weight - if w.ndim != 2: - return False - K = w.shape[1] - return ( - w.shape == (self.N, K) - and K % 64 == 0 - and w.dtype == torch.bfloat16 - and w.device == self.device - and w.is_contiguous() - ) - return self._can_run_mxfp8(linear) - - def _can_run_mxfp8(self, linear: LinearBase) -> bool: - # Only the FlashInfer MXFP8 kernels keep the [N, K] row-major weight - # plus a flat swizzled scale, which is what the kernel consumes. - from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( - FlashInferCutedslMxfp8LinearKernel, - FlashInferCutlassMxfp8LinearKernel, - ) - - kernel = getattr(linear.quant_method, "kernel", None) - if not isinstance( - kernel, - (FlashInferCutedslMxfp8LinearKernel, FlashInferCutlassMxfp8LinearKernel), - ): + if not isinstance(linear.quant_method, UnquantizedLinearMethod): return False - operands = self._mxfp8_operands(linear) - if operands is None: - return False - w, w_sf, K = operands - if K % 128 != 0: - return False - self._weight_scales[w.data_ptr()] = w_sf - self.dtypes.add("mxfp8") - return True - - def _mxfp8_operands( - self, linear: LinearBase - ) -> tuple[torch.Tensor, torch.Tensor, int] | None: - """Return (weight [N, K] FP8, flat swizzled scale, K), or None.""" w = linear.weight - w_sf = getattr(linear, "weight_scale", None) - if ( - w.ndim != 2 - or w.dtype != torch.float8_e4m3fn - or w.device != self.device - or w_sf is None - or w_sf.ndim != 1 - or w_sf.dtype != torch.uint8 - or not w_sf.is_contiguous() - ): - return None - if w.shape[0] == self.N and w.is_contiguous(): - pass - elif w.shape[1] == self.N and w.t().is_contiguous(): - # FlashInfer CuTe-DSL stores the weight as a [K, N] column-major - # view of the row-major [N, K] buffer. - w = w.t() - else: - return None - return w, w_sf, w.shape[1] + if w.ndim != 2: + return False + K = w.shape[1] + return ( + w.shape == (self.N, K) + and K % 64 == 0 + and w.dtype == torch.bfloat16 + and w.device == self.device + and w.is_contiguous() + ) def warn_incompatible_projection(self) -> None: logger.warning_once( @@ -1022,6 +965,7 @@ def __call__( assert x.dtype == torch.bfloat16 assert x.device == self.device assert x.is_contiguous() + assert w.device == self.device N = self.N dtype = "mxfp8" if w.dtype == torch.float8_e4m3fn else "bf16" x_sf = None @@ -1030,12 +974,12 @@ def __call__( mxfp8_e4m3_quantize, ) - if w_sf is None: - w_sf = self._weight_scales.get(w.data_ptr()) assert w_sf is not None + assert w_sf.dtype == torch.uint8 and w_sf.device == self.device assert w_sf.ndim == 1 and w_sf.is_contiguous() - if w.shape[0] != N: - assert w.shape[1] == N + assert w_sf.numel() == N * (K // 32) + # CuTe-DSL stores [K, N] column-major, including when K == N. + if not w.is_contiguous(): w = w.t() assert w.shape == (N, K) and K % 128 == 0 assert w.is_contiguous() @@ -1048,7 +992,6 @@ def __call__( assert w_sf is None assert w.shape == (N, K) and K % 64 == 0 assert w.dtype == torch.bfloat16 - assert w.device == self.device assert w.is_contiguous() padded_M = (M + self.world_size - 1) // self.world_size padded_M *= self.world_size @@ -1120,19 +1063,15 @@ def warmup_gemm_rs_ar() -> int: # Initialization can be disabled or fail when multicast is unavailable. if _gemm_rs_ar is None: return 0 - # Keep these profiles in sync with the dispatch in GemmRsAr.__call__. - count = 0 - for dtype in sorted(_gemm_rs_ar.dtypes): - for cta_group in (1, 2): - Sm100GemmRsAr.compile( - _gemm_rs_ar.rank, - _gemm_rs_ar.world_size, - cta_group, - _gemm_rs_ar.all_reduce, - dtype=dtype, - ) - count += 1 - return count + + for cta_group in (1, 2): + Sm100GemmRsAr.compile( + _gemm_rs_ar.rank, + _gemm_rs_ar.world_size, + cta_group, + _gemm_rs_ar.all_reduce, + ) + return 2 def get_gemm_rs_ar() -> GemmRsAr: From 3d318d1726c969568ec3438bd169573cb0c842b9 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Thu, 17 Sep 2026 08:06:54 +0000 Subject: [PATCH 03/11] Simplify MXFP8 GEMM-RS/AR reference tests Use torch._scaled_mm for the MXFP8 reference and require contiguous weights in the kernel wrapper. Co-authored-by: OpenAI Codex Signed-off-by: Thien Tran --- tests/kernels/test_kimi_k3_gemm_rs_ar.py | 52 ++++++++----------- .../kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py | 3 -- 2 files changed, 22 insertions(+), 33 deletions(-) diff --git a/tests/kernels/test_kimi_k3_gemm_rs_ar.py b/tests/kernels/test_kimi_k3_gemm_rs_ar.py index 1762afafa042..6bc55c6d3fad 100644 --- a/tests/kernels/test_kimi_k3_gemm_rs_ar.py +++ b/tests/kernels/test_kimi_k3_gemm_rs_ar.py @@ -27,26 +27,14 @@ _N = 512 -def _round_to_mxfp8(x: torch.Tensor) -> torch.Tensor: - from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( - dequant_mxfp8_to_bf16, - mxfp8_e4m3_quantize, - ) - - values, scales = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=False) - return dequant_mxfp8_to_bf16(values, scales) - - def _reference( x: torch.Tensor, weight: torch.Tensor, + weight_scale: torch.Tensor | None, world_size: int, group: dist.ProcessGroup, all_reduce: bool, - mxfp8: bool, ) -> torch.Tensor: - if mxfp8: - x = _round_to_mxfp8(x) M = x.shape[0] padded_M = (M + world_size - 1) // world_size * world_size partial = torch.empty( @@ -54,7 +42,21 @@ def _reference( dtype=x.dtype, device=x.device, ) - torch.mm(x, weight.T, out=partial[:M]) + if weight_scale is None: + torch.mm(x, weight.T, out=partial[:M]) + else: + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + mxfp8_e4m3_quantize, + ) + + x_q, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True) + partial[:M] = torch._scaled_mm( + x_q, + weight.T, + x_sf.view(torch.float8_e8m0fnu), + weight_scale.view(torch.float8_e8m0fnu), + out_dtype=x.dtype, + ) if padded_M > M: partial[M:].zero_() @@ -113,17 +115,10 @@ def _run_mode( mxfp8_e4m3_quantize, ) - operands = {} - reference_weights = {} - for K, w in weights.items(): - if mxfp8: - w_q, w_sf = mxfp8_e4m3_quantize(w, is_sf_swizzled_layout=True) - # Exercise column-major square weights as well as row-major weights. - operands[K] = (w_q.t() if K == _N else w_q, w_sf) - reference_weights[K] = _round_to_mxfp8(w) - else: - operands[K] = (w, None) - reference_weights[K] = w + operands: dict[int, tuple[torch.Tensor, torch.Tensor | None]] = { + K: mxfp8_e4m3_quantize(w, is_sf_swizzled_layout=True) if mxfp8 else (w, None) + for K, w in weights.items() + } input_generator = torch.Generator(device=device) @@ -138,9 +133,7 @@ def _run_mode( device=device, generator=input_generator, ) - expected = _reference( - x, reference_weights[K], world_size, group, all_reduce, mxfp8 - ) + expected = _reference(x, *operands[K], world_size, group, all_reduce) actual = gemm_rs_ar(x, *operands[K]) torch.accelerator.synchronize(device) _assert_valid_rows_close(actual, expected, M, rank, all_reduce) @@ -174,11 +167,10 @@ def _run_mode( ) graph_expected = _reference( graph_x, - reference_weights[graph_K], + *operands[graph_K], world_size, group, all_reduce, - mxfp8, ) graph_output = torch.empty_like(graph_expected) diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py index b71bc5abed9f..b3ae17e137ab 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py @@ -978,9 +978,6 @@ def __call__( assert w_sf.dtype == torch.uint8 and w_sf.device == self.device assert w_sf.ndim == 1 and w_sf.is_contiguous() assert w_sf.numel() == N * (K // 32) - # CuTe-DSL stores [K, N] column-major, including when K == N. - if not w.is_contiguous(): - w = w.t() assert w.shape == (N, K) and K % 128 == 0 assert w.is_contiguous() x, x_sf = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=True) From 673d58c61371053b24eb885cba548694b8eadfbf Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Thu, 17 Sep 2026 08:46:13 +0000 Subject: [PATCH 04/11] Integrate MXFP8 projections with Kimi-K3 GEMM-RS/AR Select compatible backends before weight loading, adapt loaded weights and scales at callsites, and warm up each selected dtype. Exercise BF16 and both FlashInfer MXFP8 layouts against torch._scaled_mm. Co-authored-by: OpenAI Codex Signed-off-by: Thien Tran --- tests/kernels/test_kimi_k3_gemm_rs_ar.py | 85 +++++++++++++------ vllm/models/kimi_k3/nvidia/kda.py | 2 +- vllm/models/kimi_k3/nvidia/mla.py | 2 +- vllm/models/kimi_k3/nvidia/model.py | 2 +- .../kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py | 75 +++++++++++----- 5 files changed, 116 insertions(+), 50 deletions(-) diff --git a/tests/kernels/test_kimi_k3_gemm_rs_ar.py b/tests/kernels/test_kimi_k3_gemm_rs_ar.py index 6bc55c6d3fad..83430fffda6e 100644 --- a/tests/kernels/test_kimi_k3_gemm_rs_ar.py +++ b/tests/kernels/test_kimi_k3_gemm_rs_ar.py @@ -6,7 +6,7 @@ import torch.distributed as dist import torch.multiprocessing as mp -from tests.utils import ensure_current_vllm_config +from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.distributed import cleanup_dist_env_and_memory from vllm.distributed.parallel_state import ( init_distributed_environment, @@ -101,7 +101,7 @@ def _run_mode( rank: int, world_size: int, weights: dict[int, torch.Tensor], - mxfp8: bool, + backend: str, ) -> None: # cute_dsl is unavailable off CUDA, so import it only inside the GPU worker. from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import GemmRsAr @@ -111,14 +111,37 @@ def _run_mode( N=_N, all_reduce=all_reduce, ) - from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( - mxfp8_e4m3_quantize, + from vllm.config.quantization import QuantizationConfigArgs + from vllm.model_executor.layers.linear import RowParallelLinear + from vllm.model_executor.layers.quantization.online.base import ( + OnlineQuantizationConfig, ) - operands: dict[int, tuple[torch.Tensor, torch.Tensor | None]] = { - K: mxfp8_e4m3_quantize(w, is_sf_swizzled_layout=True) if mxfp8 else (w, None) - for K, w in weights.items() - } + quant_config = ( + None + if backend == "bf16" + else OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8")) + ) + projections = {} + operands = {} + for K, w in weights.items(): + with torch.device(device): + linear = RowParallelLinear( + K * world_size, + _N, + bias=False, + params_dtype=torch.bfloat16, + quant_config=quant_config, + ) + # Select the fused path before online quantization replaces meta weights. + assert gemm_rs_ar.can_run(linear) + linear.weight = torch.nn.Parameter(w.clone(), requires_grad=False) + linear.quant_method.process_weights_after_loading(linear) + projections[K] = linear + weight = linear.weight + if backend == "flashinfer_cutedsl": + weight = weight.t() + operands[K] = (weight, getattr(linear, "weight_scale", None)) input_generator = torch.Generator(device=device) @@ -134,7 +157,7 @@ def _run_mode( generator=input_generator, ) expected = _reference(x, *operands[K], world_size, group, all_reduce) - actual = gemm_rs_ar(x, *operands[K]) + actual = gemm_rs_ar.apply(x, projections[K]) torch.accelerator.synchronize(device) _assert_valid_rows_close(actual, expected, M, rank, all_reduce) @@ -149,9 +172,9 @@ def _run_mode( device=device, generator=input_generator, ) - first_output = gemm_rs_ar(lifetime_x, *operands[lifetime_K]) + first_output = gemm_rs_ar.apply(lifetime_x, projections[lifetime_K]) first_snapshot = first_output.clone() - second_output = gemm_rs_ar(-lifetime_x, *operands[lifetime_K]) + second_output = gemm_rs_ar.apply(-lifetime_x, projections[lifetime_K]) torch.accelerator.synchronize(device) assert first_output.data_ptr() != second_output.data_ptr() torch.testing.assert_close(first_output, first_snapshot, rtol=0, atol=0) @@ -178,13 +201,13 @@ def _run_mode( capture_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(capture_stream): for _ in range(3): - graph_output.copy_(gemm_rs_ar(graph_x, *operands[graph_K])) + graph_output.copy_(gemm_rs_ar.apply(graph_x, projections[graph_K])) capture_stream.synchronize() dist.barrier(group=group) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=capture_stream): - graph_output.copy_(gemm_rs_ar(graph_x, *operands[graph_K])) + graph_output.copy_(gemm_rs_ar.apply(graph_x, projections[graph_K])) torch.cuda.current_stream().wait_stream(capture_stream) dist.barrier(group=group) @@ -203,7 +226,7 @@ def _run_mode( del gemm_rs_ar, capture_stream, graph, graph_output -def _worker(local_rank: int, world_size: int, master_port: int, mxfp8: bool) -> None: +def _worker(local_rank: int, world_size: int, master_port: int, backend: str) -> None: # This module pulls in cute_dsl, which is unavailable off CUDA, so importing # it at module level would fail collection. Import it where it is used, as # `kda.py` does. @@ -221,7 +244,10 @@ def _worker(local_rank: int, world_size: int, master_port: int, mxfp8: bool) -> ) init_distributed_environment() - with ensure_current_vllm_config(): + config = VllmConfig() + config.model_config = ModelConfig(dtype="bfloat16") + config.kernel_config.linear_backend = "auto" if backend == "bf16" else backend + with set_current_vllm_config(config): initialize_model_parallel(tensor_model_parallel_size=world_size) group = dist.group.WORLD @@ -241,26 +267,29 @@ def _worker(local_rank: int, world_size: int, master_port: int, mxfp8: bool) -> # Production binds one mode per worker; exercise both mode-bound instances # sequentially in this test without implying that both are initialized. - for all_reduce in (False, True): - _run_mode( - all_reduce=all_reduce, - device=device, - group=group, - rank=rank, - world_size=world_size, - weights=weights, - mxfp8=mxfp8, - ) + with set_current_vllm_config(config): + for all_reduce in (False, True): + _run_mode( + all_reduce=all_reduce, + device=device, + group=group, + rank=rank, + world_size=world_size, + weights=weights, + backend=backend, + ) cleanup_dist_env_and_memory() -@pytest.mark.parametrize("mxfp8", [False, True], ids=["bf16", "mxfp8"]) +@pytest.mark.parametrize( + "backend", ["bf16", "flashinfer_cutlass", "flashinfer_cutedsl"] +) @pytest.mark.distributed(num_gpus=2) @pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="Kimi-K3 GEMM-RS/AR requires SM100", ) -def test_kimi_k3_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch, mxfp8: bool) -> None: +def test_kimi_k3_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch, backend: str) -> None: world_size = 2 if torch.accelerator.device_count() < world_size: pytest.skip("GEMM-RS/AR requires two GPUs") @@ -270,7 +299,7 @@ def test_kimi_k3_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch, mxfp8: bool) -> Non try: mp.spawn( _worker, - args=(world_size, get_open_port(), mxfp8), + args=(world_size, get_open_port(), backend), nprocs=world_size, ) finally: diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index 92326944f592..172d2213bf6b 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -869,7 +869,7 @@ def forward( ) core_attn_out = rearrange(core_attn_out, "1 n h d -> n (h d)") if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(core_attn_out): - return self.gemm_rs_ar(core_attn_out, self.o_proj.weight) + return self.gemm_rs_ar.apply(core_attn_out, self.o_proj) return self.o_proj(core_attn_out)[0] @eager_break_during_capture diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 87bf59749c1d..568b6a51be34 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -631,7 +631,7 @@ def forward( attn_out = _gate_sigmoid_mul(attn_out, gate) if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(attn_out): - return self.gemm_rs_ar(attn_out, self.o_proj.weight) + return self.gemm_rs_ar.apply(attn_out, self.o_proj) return self.o_proj(attn_out)[0] diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 7fe612b65466..88197c3a0fd1 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -309,7 +309,7 @@ def forward(self, x): x = self.act_fn(gate_up) if self.gemm_rs_ar is not None and self.gemm_rs_ar.should_run(x): - return self.gemm_rs_ar(x, self.down_proj.weight) + return self.gemm_rs_ar.apply(x, self.down_proj) x, _ = self.down_proj(x) if self.shard_sequence_parallel: diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py index b3ae17e137ab..cb1d1122a376 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py @@ -887,6 +887,7 @@ def __init__(self, *, max_M: int, N: int, all_reduce: bool = False) -> None: self.N = N self.device = device self.all_reduce = all_reduce + self.dtypes: set[str] = set() self.partial = symm_mem.empty((max_M, N), dtype=torch.bfloat16, device=device) self.partial_handle = symm_mem.rendezvous(self.partial, group) @@ -926,21 +927,55 @@ def __init__(self, *, max_M: int, N: int, all_reduce: bool = False) -> None: tp_group.barrier() def can_run(self, linear: LinearBase) -> bool: - # Validate projection-invariant requirements once during model init. - if not isinstance(linear.quant_method, UnquantizedLinearMethod): - return False + from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( + FlashInferCutedslMxfp8LinearKernel, + FlashInferCutlassMxfp8LinearKernel, + ) + + # Called before weight loading; online quantization starts on meta. + if isinstance(linear.quant_method, UnquantizedLinearMethod): + dtype, k_alignment = "bf16", 64 + if linear.weight.dtype != torch.bfloat16: + return False + else: + method = getattr(linear, "scheme", linear.quant_method) + kernel = getattr(method, "kernel", None) + if not isinstance( + kernel, + ( + FlashInferCutedslMxfp8LinearKernel, + FlashInferCutlassMxfp8LinearKernel, + ), + ): + return False + dtype, k_alignment = "mxfp8", 128 + w = linear.weight - if w.ndim != 2: + if ( + w.ndim != 2 + or w.shape[0] != self.N + or w.shape[1] % k_alignment != 0 + or (w.device != self.device and not w.is_meta) + or not w.is_contiguous() + ): return False - K = w.shape[1] - return ( - w.shape == (self.N, K) - and K % 64 == 0 - and w.dtype == torch.bfloat16 - and w.device == self.device - and w.is_contiguous() + self.dtypes.add(dtype) + return True + + def apply(self, x: torch.Tensor, linear: LinearBase) -> torch.Tensor: + from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( + FlashInferCutedslMxfp8LinearKernel, ) + method = getattr(linear, "scheme", linear.quant_method) + w = linear.weight + if isinstance( + getattr(method, "kernel", None), FlashInferCutedslMxfp8LinearKernel + ): + # This backend stores a column-major [K, N] view after loading. + w = w.t() + return self(x, w, getattr(linear, "weight_scale", None)) + def warn_incompatible_projection(self) -> None: logger.warning_once( "Some projections are incompatible with GEMM-RS/AR; using the " @@ -1061,14 +1096,16 @@ def warmup_gemm_rs_ar() -> int: if _gemm_rs_ar is None: return 0 - for cta_group in (1, 2): - Sm100GemmRsAr.compile( - _gemm_rs_ar.rank, - _gemm_rs_ar.world_size, - cta_group, - _gemm_rs_ar.all_reduce, - ) - return 2 + for dtype in sorted(_gemm_rs_ar.dtypes): + for cta_group in (1, 2): + Sm100GemmRsAr.compile( + _gemm_rs_ar.rank, + _gemm_rs_ar.world_size, + cta_group, + _gemm_rs_ar.all_reduce, + dtype=dtype, + ) + return 2 * len(_gemm_rs_ar.dtypes) def get_gemm_rs_ar() -> GemmRsAr: From a277a8a61eff1177dd665bacd1db3a752b46365b Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Thu, 17 Sep 2026 08:50:45 +0000 Subject: [PATCH 05/11] Consolidate GEMM-RS/AR execution in apply Fold the tensor launch path into apply and migrate the tracked benchmark to the projection-based API. Co-authored-by: OpenAI Codex Signed-off-by: Thien Tran --- .../kernels/benchmark_kimi_k3_gemm_rs_ar.py | 10 +++++- .../kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py | 32 ++++++++----------- 2 files changed, 22 insertions(+), 20 deletions(-) diff --git a/benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py b/benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py index e4daea26bbce..796bafaf9e7d 100644 --- a/benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py +++ b/benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py @@ -26,6 +26,7 @@ init_distributed_environment, initialize_model_parallel, ) +from vllm.model_executor.layers.linear import RowParallelLinear from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import GemmRsAr # Shared-expert down-proj and attention O-proj. @@ -231,8 +232,15 @@ def run() -> torch.Tensor: def make_fused_gemm_collective( x: torch.Tensor, weight: torch.Tensor ) -> Callable[[], torch.Tensor]: + with torch.device("meta"): + linear = RowParallelLinear( + K * world_size, N, bias=False, params_dtype=weight.dtype + ) + linear.weight = torch.nn.Parameter(weight, requires_grad=False) + assert gemm_rs_ar.can_run(linear) + def run() -> torch.Tensor: - return gemm_rs_ar(x, weight) + return gemm_rs_ar.apply(x, linear) return run diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py index cb1d1122a376..c226b28db32c 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py @@ -962,6 +962,18 @@ def can_run(self, linear: LinearBase) -> bool: self.dtypes.add(dtype) return True + def warn_incompatible_projection(self) -> None: + logger.warning_once( + "Some projections are incompatible with GEMM-RS/AR; using the " + "unfused path instead.", + scope="global", + ) + + def should_run(self, x: torch.Tensor) -> bool: + # Use the same threshold for RS and AR for now. Small-M shapes are + # supported but faster on the existing LL path. + return x.shape[0] >= 128 + def apply(self, x: torch.Tensor, linear: LinearBase) -> torch.Tensor: from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( FlashInferCutedslMxfp8LinearKernel, @@ -974,26 +986,8 @@ def apply(self, x: torch.Tensor, linear: LinearBase) -> torch.Tensor: ): # This backend stores a column-major [K, N] view after loading. w = w.t() - return self(x, w, getattr(linear, "weight_scale", None)) - - def warn_incompatible_projection(self) -> None: - logger.warning_once( - "Some projections are incompatible with GEMM-RS/AR; using the " - "unfused path instead.", - scope="global", - ) + w_sf = getattr(linear, "weight_scale", None) - def should_run(self, x: torch.Tensor) -> bool: - # Use the same threshold for RS and AR for now. Small-M shapes are - # supported but faster on the existing LL path. - return x.shape[0] >= 128 - - def __call__( - self, - x: torch.Tensor, - w: torch.Tensor, - w_sf: torch.Tensor | None = None, - ) -> torch.Tensor: assert x.ndim == 2 M, K = x.shape assert 0 < M <= self.max_M From 129db1a2d49a4fd251e39761b68f298d5d5bf8e9 Mon Sep 17 00:00:00 2001 From: Canlin Date: Fri, 18 Sep 2026 02:34:31 +0800 Subject: [PATCH 06/11] [Perf][DSV4.1] Fuse MXFP8 wo_b GEMM with sequence-parallel reduce-scatter Bind the DeepSeek-V4.1 sequence-parallel wo_b projection to the SM100 BF16/MXFP8 GEMM-RS kernel (opt-in via VLLM_DSV41_GEMM_RS=1) and move that kernel out of the Kimi-K3 tree into vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py so both models share one workspace, init gate, warmup, test and benchmark. - DeepseekV4Attention gains bind_gemm_rs()/_wo_b_proj(); every _o_proj routes wo_b through it and the decoder layer skips its own sp_reduce_scatter when the fused path is bound. DSpark draft layers bind as well. - GemmRsAr.apply zeroes RS output rows beyond M so callers that pass an unpadded token count see the same padding as sp_reduce_scatter. - Benchmark gains --model deepseek_v41 and --backend flashinfer_cutlass|flashinfer_cutedsl presets. Builds on gau-nernst's MXFP8 support in the Kimi-K3 GEMM-RS/AR kernel. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Canlin --- .buildkite/test_areas/distributed.yaml | 11 +- ..._gemm_rs_ar.py => benchmark_gemm_rs_ar.py} | 193 +++++++++++++----- ...mi_k3_gemm_rs_ar.py => test_gemm_rs_ar.py} | 10 +- vllm/envs.py | 5 + .../kernels/linear}/cute_dsl/gemm_rs_ar.py | 74 ++++++- vllm/model_executor/warmup/kernel_warmup.py | 15 +- vllm/models/deepseek_v4/nvidia/ops/o_proj.py | 8 +- vllm/models/deepseek_v41/attention.py | 40 ++++ vllm/models/deepseek_v41/nvidia/dspark.py | 7 + .../nvidia/flash_mla_mega_attn.py | 2 +- .../deepseek_v41/nvidia/flashinfer_sparse.py | 4 +- vllm/models/deepseek_v41/nvidia/flashmla.py | 2 +- vllm/models/deepseek_v41/nvidia/model.py | 38 +++- vllm/models/kimi_k3/nvidia/kda.py | 2 +- vllm/models/kimi_k3/nvidia/mla.py | 2 +- vllm/models/kimi_k3/nvidia/model.py | 51 +---- 16 files changed, 343 insertions(+), 121 deletions(-) rename benchmarks/kernels/{benchmark_kimi_k3_gemm_rs_ar.py => benchmark_gemm_rs_ar.py} (72%) rename tests/kernels/{test_kimi_k3_gemm_rs_ar.py => test_gemm_rs_ar.py} (95%) rename vllm/{models/kimi_k3/nvidia/ops => model_executor/kernels/linear}/cute_dsl/gemm_rs_ar.py (93%) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 3ea26ce89c1d..9b2ad524d649 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -310,17 +310,18 @@ steps: - pytest -v -s tests/v1/distributed/test_dbo.py - pytest -v -s tests/distributed/test_mnnvl_alltoall.py -- label: ":nvidia: (B200) Kimi-K3 GEMM-RS/AR" - key: kimi-k3-gemm-rs-ar-2xb200 +- label: ":nvidia: (B200) GEMM-RS/AR" + key: gemm-rs-ar-2xb200 timeout_in_minutes: 20 device: b200-k8s working_dir: "/vllm-workspace/" num_devices: 2 source_file_dependencies: - - vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py - - tests/kernels/test_kimi_k3_gemm_rs_ar.py + - vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py + - vllm/cute_utils/ + - tests/kernels/test_gemm_rs_ar.py commands: - - pytest -v -s tests/kernels/test_kimi_k3_gemm_rs_ar.py + - pytest -v -s tests/kernels/test_gemm_rs_ar.py - label: ":nvidia: (L4) Distributed 2-Node" key: 2-node-test-4-gpus diff --git a/benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py b/benchmarks/kernels/benchmark_gemm_rs_ar.py similarity index 72% rename from benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py rename to benchmarks/kernels/benchmark_gemm_rs_ar.py index 796bafaf9e7d..7234cd33bc5a 100644 --- a/benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py +++ b/benchmarks/kernels/benchmark_gemm_rs_ar.py @@ -1,11 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Benchmark the SM100 Kimi-K3 GEMM-RS/AR kernel. +"""Benchmark the SM100 GEMM-RS/AR kernel. -All ranks must belong to one NVLink domain. For example, run a TP8 sweep with: +All ranks must belong to one NVLink domain. For example, run a TP8 sweep over +the Kimi-K3 projections with: - torchrun --nproc-per-node=8 \ - benchmarks/kernels/benchmark_kimi_k3_gemm_rs_ar.py + torchrun --nproc-per-node=8 benchmarks/kernels/benchmark_gemm_rs_ar.py + +or the DeepSeek-V4.1 MXFP8 ``wo_b`` projection at TP4 with: + + torchrun --nproc-per-node=4 benchmarks/kernels/benchmark_gemm_rs_ar.py \ + --model deepseek_v41 --backend flashinfer_cutedsl """ import argparse @@ -19,18 +24,23 @@ import torch.distributed as dist import torch.distributed._symmetric_memory as symm_mem -from vllm.config import VllmConfig, set_current_vllm_config +from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.distributed import cleanup_dist_env_and_memory from vllm.distributed.parallel_state import ( get_tp_group, init_distributed_environment, initialize_model_parallel, ) +from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr from vllm.model_executor.layers.linear import RowParallelLinear -from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import GemmRsAr -# Shared-expert down-proj and attention O-proj. -_KIMI_K3_PROJECTION_K = (6144, 12288) +# Per model: (hidden size N, global input widths K). Kimi-K3 lists the +# shared-expert down-proj and attention O-proj; DeepSeek-V4.1 lists the +# ``wo_b`` output projection (o_groups * o_lora_rank). +_MODEL_PROJECTIONS = { + "kimi_k3": (7168, (6144, 12288)), + "deepseek_v41": (5120, (8192,)), +} @dataclass @@ -55,17 +65,32 @@ def parse_args() -> argparse.Namespace: default=[128, 512, 2048, 8192, 32768], help="Global token counts to benchmark.", ) + parser.add_argument( + "--model", + choices=tuple(_MODEL_PROJECTIONS), + default="kimi_k3", + help="Model whose projection shapes set the default --n and --k.", + ) + parser.add_argument( + "--backend", + choices=("bf16", "flashinfer_cutlass", "flashinfer_cutedsl"), + default="bf16", + help="Weight dtype/layout: BF16 or an online-MXFP8 FlashInfer kernel.", + ) parser.add_argument( "--k", type=int, nargs="+", help=( - "Per-rank input dimensions. By default, derive the Kimi-K3 " - "shared-expert down-proj and O-proj dimensions from the TP " - "world size." + "Per-rank input dimensions. By default, derive the selected " + "model's projection dimensions from the TP world size." ), ) - parser.add_argument("--n", type=int, default=7168) + parser.add_argument( + "--n", + type=int, + help="Output width. Defaults to the selected model's hidden size.", + ) parser.add_argument( "--num-workspaces", type=int, @@ -144,9 +169,47 @@ def valid_rows(M: int, local_M: int, rank: int) -> int: return min(max(M - rank * local_M, 0), local_M) +def make_projection( + gemm_rs_ar: GemmRsAr, + weight: torch.Tensor, + world_size: int, + backend: str, +) -> RowParallelLinear: + """Wrap a BF16 [N, K] weight the way the model would hand it to GEMM-RS. + + For the MXFP8 backends this quantizes online through the FlashInfer + kernel's post-load step, so the fused kernel and the unfused baseline see + the production weight and scale layouts. + """ + N, K = weight.shape + quant_config = None + if backend != "bf16": + from vllm.config.quantization import QuantizationConfigArgs + from vllm.model_executor.layers.quantization.online.base import ( + OnlineQuantizationConfig, + ) + + quant_config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8")) + with torch.device(weight.device): + linear = RowParallelLinear( + K * world_size, + N, + bias=False, + params_dtype=weight.dtype, + quant_config=quant_config, + return_bias=False, + ) + # Select the fused path before online quantization replaces the weights. + assert gemm_rs_ar.can_run(linear) + linear.weight = torch.nn.Parameter(weight, requires_grad=False) + linear.quant_method.process_weights_after_loading(linear) + return linear + + def benchmark_shape( gemm_rs_ar: GemmRsAr, mode: str, + backend: str, M: int, N: int, K: int, @@ -170,8 +233,13 @@ def benchmark_shape( torch.randn(M, K, dtype=torch.bfloat16, device=device, generator=rng) for _ in range(num_workspaces) ] - weights = [ - torch.randn(N, K, dtype=torch.bfloat16, device=device, generator=rng) + projections = [ + make_projection( + gemm_rs_ar, + torch.randn(N, K, dtype=torch.bfloat16, device=device, generator=rng), + world_size, + backend, + ) for _ in range(num_workspaces) ] @@ -199,11 +267,20 @@ def benchmark_shape( partial[M:].zero_() symm_partial[M:].zero_() + def gemm_into(x: torch.Tensor, linear: RowParallelLinear, out: torch.Tensor): + # BF16 writes straight into the collective buffer; the MXFP8 kernels + # allocate their own output, so the baseline pays one extra copy, as + # the model's unfused path does before its collective. + if backend == "bf16": + torch.mm(x, linear.weight.T, out=out) + else: + out.copy_(linear(x)) + def make_torch_ring_ll_gemm_collective( - x: torch.Tensor, weight: torch.Tensor + x: torch.Tensor, linear: RowParallelLinear ) -> Callable[[], torch.Tensor]: def run() -> torch.Tensor: - torch.mm(x, weight.T, out=partial[:M]) + gemm_into(x, linear, partial[:M]) if all_reduce: dist.all_reduce(partial, group=device_group) return partial[:M] @@ -213,10 +290,10 @@ def run() -> torch.Tensor: return run def make_torch_ldmc_gemm_collective( - x: torch.Tensor, weight: torch.Tensor + x: torch.Tensor, linear: RowParallelLinear ) -> Callable[[], torch.Tensor]: def run() -> torch.Tensor: - torch.mm(x, weight.T, out=symm_partial[:M]) + gemm_into(x, linear, symm_partial[:M]) if all_reduce: dist.all_reduce(symm_partial, group=device_group) return symm_partial[:M] @@ -230,15 +307,8 @@ def run() -> torch.Tensor: return run def make_fused_gemm_collective( - x: torch.Tensor, weight: torch.Tensor + x: torch.Tensor, linear: RowParallelLinear ) -> Callable[[], torch.Tensor]: - with torch.device("meta"): - linear = RowParallelLinear( - K * world_size, N, bias=False, params_dtype=weight.dtype - ) - linear.weight = torch.nn.Parameter(weight, requires_grad=False) - assert gemm_rs_ar.can_run(linear) - def run() -> torch.Tensor: return gemm_rs_ar.apply(x, linear) @@ -246,10 +316,11 @@ def run() -> torch.Tensor: def make_torch_gemm( x: torch.Tensor, - weight: torch.Tensor, + linear: RowParallelLinear, ) -> Callable[[], torch.Tensor]: def run() -> torch.Tensor: - return torch.mm(x, weight.T, out=gemm_output) + gemm_into(x, linear, gemm_output) + return gemm_output return run @@ -281,22 +352,23 @@ def run() -> torch.Tensor: return run + pairs = list(zip(inputs, projections)) candidates = ( Candidate( "ring_ll_us", - [make_torch_ring_ll_gemm_collective(x, w) for x, w in zip(inputs, weights)], + [make_torch_ring_ll_gemm_collective(x, w) for x, w in pairs], ), Candidate( "ldmc_us", - [make_torch_ldmc_gemm_collective(x, w) for x, w in zip(inputs, weights)], + [make_torch_ldmc_gemm_collective(x, w) for x, w in pairs], ), Candidate( "gemm_rs_ar_us", - [make_fused_gemm_collective(x, w) for x, w in zip(inputs, weights)], + [make_fused_gemm_collective(x, w) for x, w in pairs], ), Candidate( "torch_gemm_us", - [make_torch_gemm(x, w) for x, w in zip(inputs, weights)], + [make_torch_gemm(x, w) for x, w in pairs], check_correctness=False, ), Candidate( @@ -345,6 +417,7 @@ def run() -> torch.Tensor: best_nccl_us = min(times["ring_ll_collective_us"], times["ldmc_collective_us"]) return { "mode": mode.upper(), + "backend": backend, "M": M, "N": N, "K": K, @@ -406,8 +479,11 @@ def print_results(results: list[dict[str, float | int | str]]) -> None: def main() -> None: args = parse_args() + default_n, default_k = _MODEL_PROJECTIONS[args.model] + if args.n is None: + args.n = default_n assert args.m and min(args.m) > 0 - assert args.n % 256 == 0 + assert args.n % 128 == 0 assert args.num_workspaces > 0 assert args.warmup_replays >= 0 assert args.samples > 0 @@ -417,16 +493,23 @@ def main() -> None: init_distributed_environment() world_size = dist.get_world_size() if args.k is None: - assert all(K % world_size == 0 for K in _KIMI_K3_PROJECTION_K) - K_values = [K // world_size for K in _KIMI_K3_PROJECTION_K] + assert all(K % world_size == 0 for K in default_k) + K_values = [K // world_size for K in default_k] else: K_values = args.k - assert all(K % 64 == 0 for K in K_values) + k_alignment = 64 if args.backend == "bf16" else 128 + assert all(K % k_alignment == 0 for K in K_values) # Reserve symmetric memory for the NCCL-managed benchmark allocations. os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0" # NCCL-managed symmetric allocations select the NVLS/LDMC collective path. symm_mem.set_backend("NCCL") - with set_current_vllm_config(VllmConfig()): + config = VllmConfig() + config.model_config = ModelConfig(dtype="bfloat16") + # Pin the online-MXFP8 kernel the projections quantize through. + config.kernel_config.linear_backend = ( + "auto" if args.backend == "bf16" else args.backend + ) + with set_current_vllm_config(config): initialize_model_parallel(tensor_model_parallel_size=world_size) tp_group = get_tp_group() @@ -447,23 +530,25 @@ def device_barrier() -> None: N=args.n, all_reduce=args.mode == "ar", ) - results = [ - benchmark_shape( - gemm_rs_ar, - args.mode, - M, - args.n, - K, - args.num_workspaces, - args.warmup_replays, - args.samples, - tp_group.device_group, - tp_group.cpu_group, - device_barrier, - ) - for K in K_values - for M in args.m - ] + with set_current_vllm_config(config): + results = [ + benchmark_shape( + gemm_rs_ar, + args.mode, + args.backend, + M, + args.n, + K, + args.num_workspaces, + args.warmup_replays, + args.samples, + tp_group.device_group, + tp_group.cpu_group, + device_barrier, + ) + for K in K_values + for M in args.m + ] del gemm_rs_ar if tp_group.rank_in_group == 0: diff --git a/tests/kernels/test_kimi_k3_gemm_rs_ar.py b/tests/kernels/test_gemm_rs_ar.py similarity index 95% rename from tests/kernels/test_kimi_k3_gemm_rs_ar.py rename to tests/kernels/test_gemm_rs_ar.py index 83430fffda6e..b1c79af2b925 100644 --- a/tests/kernels/test_kimi_k3_gemm_rs_ar.py +++ b/tests/kernels/test_gemm_rs_ar.py @@ -91,6 +91,10 @@ def _assert_valid_rows_close( rtol=5e-2, atol=4.0, ) + # Sequence-parallel callers may pass an unpadded M and read the padding + # rows back (they flow through per-token kernels before being dropped), + # so RS must leave them zero like ``sp_reduce_scatter`` does. + assert not actual[valid_rows:].any() def _run_mode( @@ -104,7 +108,7 @@ def _run_mode( backend: str, ) -> None: # cute_dsl is unavailable off CUDA, so import it only inside the GPU worker. - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import GemmRsAr + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr gemm_rs_ar = GemmRsAr( max_M=max(M for M, _ in _SHAPES), @@ -287,9 +291,9 @@ def _worker(local_rank: int, world_size: int, master_port: int, backend: str) -> @pytest.mark.distributed(num_gpus=2) @pytest.mark.skipif( not current_platform.is_device_capability_family(100), - reason="Kimi-K3 GEMM-RS/AR requires SM100", + reason="GEMM-RS/AR requires SM100", ) -def test_kimi_k3_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch, backend: str) -> None: +def test_gemm_rs_ar(monkeypatch: pytest.MonkeyPatch, backend: str) -> None: world_size = 2 if torch.accelerator.device_count() < world_size: pytest.skip("GEMM-RS/AR requires two GPUs") diff --git a/vllm/envs.py b/vllm/envs.py index c512df457adc..038917b1b63b 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -213,6 +213,7 @@ VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False VLLM_KIMI_K3_GEMM_AR: bool = True VLLM_KIMI_K3_GEMM_RS: bool = False + VLLM_DSV41_GEMM_RS: bool = False VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True VLLM_USE_FLASHINFER_MOE_INT4: bool = False VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None @@ -1610,6 +1611,10 @@ def _resolve_rust_cli_path() -> str | None: # Use the SM100 BF16 GEMM-RS kernel for eligible Kimi-K3 sequence-parallel # row-parallel projections. All TP ranks must belong to one NVLink domain. "VLLM_KIMI_K3_GEMM_RS": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_RS", "0"))), + # Fuse the DeepSeek-V4.1 sequence-parallel ``wo_b`` projection with its + # TP reduce-scatter using the SM100 BF16/MXFP8 GEMM-RS kernel. All TP + # ranks must belong to one NVLink domain. Off until model evals land. + "VLLM_DSV41_GEMM_RS": lambda: bool(int(os.getenv("VLLM_DSV41_GEMM_RS", "0"))), # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py b/vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py similarity index 93% rename from vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py rename to vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py index c226b28db32c..ec40fb056a13 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/gemm_rs_ar.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/gemm_rs_ar.py @@ -1,7 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""SM100 BF16/MXFP8 GEMM with fused tensor-parallel reduce-scatter/all-reduce.""" +"""SM100 BF16/MXFP8 GEMM with fused tensor-parallel reduce-scatter/all-reduce. + +Model-neutral: any row-parallel projection whose output is reduced across TP +(Kimi-K3 ``o_proj``/``down_proj``, DeepSeek-V4.1 ``wo_b``) can bind to the +process-wide workspace through ``maybe_init_gemm_rs_ar`` + ``GemmRsAr.apply``. +""" # Based on CUTLASS's Blackwell distributed GEMM-RS example at dcf215a. # See https://github.com/NVIDIA/cutlass/issues/3117 for memory semantics. @@ -9,6 +14,7 @@ # https://github.com/gau-nernst/gn-kernels (sm100_mm_mxfp8.py). from functools import cache +from typing import TYPE_CHECKING import cutlass import torch @@ -36,6 +42,10 @@ from vllm.distributed import get_tp_group from vllm.logger import init_logger from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod +from vllm.platforms import current_platform + +if TYPE_CHECKING: + from vllm.config import VllmConfig logger = init_logger(__name__) @@ -857,7 +867,7 @@ def compile( class GemmRsAr: - """Own the symmetric workspace for Kimi-K3 GEMM-RS/AR launches. + """Own the symmetric workspace for GEMM-RS/AR launches. All TP ranks must belong to one NVLink domain for multimem instructions. @@ -1036,6 +1046,14 @@ def apply(self, x: torch.Tensor, linear: LinearBase) -> torch.Tensor: output = None if not self.all_reduce: output = torch.empty((local_M, N), dtype=torch.bfloat16, device=self.device) + # The kernel never touches rows at or beyond M. Callers that hand + # in an unpadded M (DeepSeek-V4.1 slices the SP all-gather back to + # the real token count) expect the same zero padding rows that + # ``sp_reduce_scatter`` produces, so clear them here. This is a + # no-op whenever M is already a multiple of the TP size. + valid_rows = min(max(M - self.rank * local_M, 0), local_M) + if valid_rows < local_M: + output[valid_rows:].zero_() compiled = Sm100GemmRsAr.compile( self.rank, self.world_size, @@ -1084,6 +1102,58 @@ def init_gemm_rs_ar(max_M: int, N: int, *, all_reduce: bool = False) -> None: _gemm_rs_ar = GemmRsAr(max_M=max_M, N=N, all_reduce=all_reduce) +def maybe_init_gemm_rs_ar( + vllm_config: "VllmConfig", *, N: int, all_reduce: bool +) -> bool: + """Collectively initialize the mode-bound GEMM-RS/AR state if supported. + + Callers gate on their own feature flag first; this checks the worker + topology (SM100-family CUDA, BF16 model dtype, no ubatching, a TP size in + 2-16 that divides 128) and the NVLink multicast rendezvous. Returns whether + projections may bind to the kernel through ``get_gemm_rs_ar().can_run``. + """ + mode = "GEMM-AR" if all_reduce else "GEMM-RS" + parallel_config = vllm_config.parallel_config + tp_size = parallel_config.tensor_parallel_size + if parallel_config.use_ubatching: + reason = "ubatching is enabled" + elif vllm_config.model_config.dtype != torch.bfloat16: + reason = "the model dtype is not BF16" + elif not current_platform.is_cuda(): + reason = "the device is not CUDA" + elif not current_platform.is_device_capability_family(100): + reason = "the device is not SM100-family" + elif not 1 < tp_size <= 16: + reason = "TP size is not in the supported range 2-16" + elif 128 % tp_size != 0: + reason = "TP size does not divide 128" + elif N % 128 != 0: + reason = f"the output width {N} is not a multiple of 128" + else: + reason = None + + if reason is not None: + logger.warning_once("%s is disabled because %s.", mode, reason) + return False + + try: + init_gemm_rs_ar( + max_M=vllm_config.scheduler_config.max_num_batched_tokens, + N=N, + all_reduce=all_reduce, + ) + except RuntimeError as e: + logger.warning_once( + "%s is disabled because initialization failed: %s. This may mean " + "the TP ranks do not share one NVLink domain.", + mode, + e, + ) + return False + logger.info_once("%s is enabled.", mode) + return True + + def warmup_gemm_rs_ar() -> int: """Compile every reachable dispatch for the initialized GEMM-RS/AR mode.""" # Initialization can be disabled or fail when multicast is unavailable. diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index 81eca7ac50e5..33fe7002aa2f 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -136,16 +136,17 @@ def _warmup_bf16x3_router_gemm( logger.info_once("Warmed up BF16x3 router GEMM configs: %s.", configs) -def _warmup_kimi_k3_gemm_rs_ar() -> None: - # Kimi-K3 model construction imports this module only when GEMM-RS/AR is - # enabled and initializes its singleton before kernel_warmup runs. Avoid - # importing it here so other models do not compile the RS/AR variants. - module = sys.modules.get("vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar") +def _warmup_gemm_rs_ar() -> None: + # Model construction (Kimi-K3, DeepSeek-V4.1) imports this module only + # when GEMM-RS/AR is enabled and initializes its singleton before + # kernel_warmup runs. Avoid importing it here so other models do not + # compile the RS/AR variants. + module = sys.modules.get("vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar") if module is None: return compiled = module.warmup_gemm_rs_ar() if compiled: - logger.info_once("Warmed up %d Kimi-K3 GEMM-RS/AR variants.", compiled) + logger.info_once("Warmed up %d GEMM-RS/AR variants.", compiled) def _autotune_kimi_k3_kda_qkvg(model: torch.nn.Module) -> None: @@ -208,7 +209,7 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False): if current_platform.has_device_capability(90): _warmup_ll_bf16_router_gemm(worker.get_model()) - _warmup_kimi_k3_gemm_rs_ar() + _warmup_gemm_rs_ar() if worker.vllm_config.kernel_config.enable_cutedsl_warmup: # TODO(roberto): Remove after registered CuTeDSL warmups are migrated diff --git a/vllm/models/deepseek_v4/nvidia/ops/o_proj.py b/vllm/models/deepseek_v4/nvidia/ops/o_proj.py index 72ce43b46cbf..ec36bf763ebb 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/o_proj.py +++ b/vllm/models/deepseek_v4/nvidia/ops/o_proj.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable + import torch import torch.nn as nn @@ -31,7 +33,7 @@ def deep_gemm_fp8_o_proj( positions: torch.Tensor, cos_sin_cache: torch.Tensor, wo_a: nn.Module, - wo_b: nn.Module, + wo_b: Callable[[torch.Tensor], torch.Tensor], *, n_groups: int, heads_per_group: int, @@ -44,7 +46,9 @@ def deep_gemm_fp8_o_proj( """O projection: inverse RoPE + grouped wo_a + wo_b. Shared by the FlashMLA and FlashInfer CUDA backends. The attention - layer selects the recipe at initialization. + layer selects the recipe at initialization. ``wo_b`` is any callable over + the flattened ``z``: the projection module itself, or a wrapper that also + reduce-scatters its output (DeepSeek-V4.1 GEMM-RS). """ use_fp8 = wo_a.weight.dtype == torch.float8_e4m3fn o_proj_input, o_scale = fused_inv_rope_fp8_quant( diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 5b26a98c1c53..35f9fa7c2248 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -26,6 +26,7 @@ from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer from vllm.model_executor.layers.sparse_mqa_indexer import SparseMQAIndexer from vllm.models.common.ops import fused_q_kv_rmsnorm +from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter from vllm.models.deepseek_v41.common.ops import ( MXFP4_BLOCK_SIZE, fused_indexer_q_rope_quant, @@ -33,6 +34,7 @@ ) if TYPE_CHECKING: + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr from vllm.v1.attention.backends.mla.sparse_swa import ( DeepseekSparseSWAMetadata, ) @@ -389,6 +391,9 @@ def __init__( return_bias=False, prefix=f"{prefix}.wo_b", ) + # Set by ``bind_gemm_rs`` when the decoder layer runs sequence + # parallel and the fused GEMM + reduce-scatter kernel accepts wo_b. + self.gemm_rs: GemmRsAr | None = None # Initialize rotary embedding before the indexer/compressor consume it. self.rotary_emb = build_deepseek_v4_rope( @@ -643,6 +648,41 @@ def forward( ) return self._o_proj(attn_out, positions) + def bind_gemm_rs(self) -> None: + """Fuse ``wo_b`` with the sequence-parallel TP reduce-scatter. + + The decoder layer calls this after the model initialized the + process-wide GEMM-RS workspace and before weights load: the + eligibility check inspects the projection's linear kernel and weight + shape, which online quantization may later re-layout. + """ + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import ( + get_gemm_rs_ar, + ) + + # The layer owns the reduction under sequence parallel. + assert not self.wo_b.reduce_results + gemm_rs = get_gemm_rs_ar() + if gemm_rs.can_run(self.wo_b): + self.gemm_rs = gemm_rs + else: + gemm_rs.warn_incompatible_projection() + + def _wo_b_proj(self, z: torch.Tensor) -> torch.Tensor: + """Apply ``wo_b``; with GEMM-RS bound, also reduce-scatter the result. + + Every ``_o_proj`` implementation projects through this so the decoder + layer sees one contract: when ``gemm_rs`` is bound the output is + already the local sequence-parallel shard, otherwise it is the + unreduced TP partial. + """ + if self.gemm_rs is None: + return self.wo_b(z) + if self.gemm_rs.should_run(z): + return self.gemm_rs.apply(z, self.wo_b) + # Small batches stay on the unfused path, which is faster there. + return sp_reduce_scatter(self.wo_b(z)) + def _alloc_attn_out( self, num_tokens: int, hidden_states: torch.Tensor ) -> "torch.Tensor | QuantizedActivation": diff --git a/vllm/models/deepseek_v41/nvidia/dspark.py b/vllm/models/deepseek_v41/nvidia/dspark.py index 9a07189bc3bb..c2a97515d9fc 100644 --- a/vllm/models/deepseek_v41/nvidia/dspark.py +++ b/vllm/models/deepseek_v41/nvidia/dspark.py @@ -56,6 +56,7 @@ _linear_scale_param_name, _use_sequence_parallel, make_deepseek_v4_expert_params_mapping, + maybe_init_gemm_rs, ) logger = init_logger(__name__) @@ -109,12 +110,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) current_vllm_config = get_current_vllm_config() + # The target model already holds the GEMM-RS workspace (same hidden + # size, same TP group); this only re-checks and binds the draft layers. + run_gemm_rs = maybe_init_gemm_rs( + current_vllm_config, self.use_sequence_parallel + ) self.layers = nn.ModuleList( [ DeepseekV4DecoderLayer( current_vllm_config, prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"), topk_indices_buffer=self.topk_indices_buffer, + run_gemm_rs=run_gemm_rs, ) for i in range(self.num_dspark_layers) ] diff --git a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py index 9722673dd1ee..c11f330e3336 100644 --- a/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py +++ b/vllm/models/deepseek_v41/nvidia/flash_mla_mega_attn.py @@ -225,7 +225,7 @@ def _o_proj( z, recipe=self._einsum_recipe, ) - return self.wo_b(z.flatten(1)) + return self._wo_b_proj(z.flatten(1)) # ---- weights ----------------------------------------------------------- diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 409725434a17..5c8a1a6627d7 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py @@ -208,7 +208,7 @@ def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tens positions, self.rotary_emb.cos_sin_cache, self.wo_a, - self.wo_b, + self._wo_b_proj, n_groups=self.n_local_groups, heads_per_group=self.n_local_heads // self.n_local_groups, nope_dim=self.nope_head_dim, @@ -566,7 +566,7 @@ def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tens positions, self.rotary_emb.cos_sin_cache, self.wo_a, - self.wo_b, + self._wo_b_proj, n_groups=self.n_local_groups, heads_per_group=self.n_local_heads // self.n_local_groups, nope_dim=self.nope_head_dim, diff --git a/vllm/models/deepseek_v41/nvidia/flashmla.py b/vllm/models/deepseek_v41/nvidia/flashmla.py index 870098975b71..96dde36c61d1 100644 --- a/vllm/models/deepseek_v41/nvidia/flashmla.py +++ b/vllm/models/deepseek_v41/nvidia/flashmla.py @@ -65,7 +65,7 @@ def _o_proj(self, attn_out: torch.Tensor, positions: torch.Tensor) -> torch.Tens positions, self.rotary_emb.cos_sin_cache, self.wo_a, - self.wo_b, + self._wo_b_proj, n_groups=self.n_local_groups, heads_per_group=self.n_local_heads // self.n_local_groups, nope_dim=self.nope_head_dim, diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 21fa00e5c963..747792073087 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -177,6 +177,30 @@ def _use_sequence_parallel(vllm_config: VllmConfig) -> bool: ) +def maybe_init_gemm_rs(vllm_config: VllmConfig, use_sequence_parallel: bool) -> bool: + """Set up the fused ``wo_b`` GEMM + reduce-scatter when opted in. + + Sequence parallel is the only topology where ``wo_b`` ends in a + reduce-scatter, so the kernel is bound to RS mode. The workspace is a + process-wide singleton shared with any other model code in this worker + (the DSpark drafter reuses it), and its NVLink multicast rendezvous is + collective, so every TP rank must take the same decision here. + """ + if not (use_sequence_parallel and envs.VLLM_DSV41_GEMM_RS): + return False + + # The kernel module pulls in cute_dsl, so import it only once opted in. + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import ( + maybe_init_gemm_rs_ar, + ) + + hidden_size = vllm_config.model_config.hf_config.hidden_size + if not maybe_init_gemm_rs_ar(vllm_config, N=hidden_size, all_reduce=False): + return False + logger.info_once("To disable DeepSeek-V4.1 GEMM-RS, set VLLM_DSV41_GEMM_RS=0.") + return True + + class DeepseekV4DecoderLayer(nn.Module): def __init__( self, @@ -186,6 +210,7 @@ def __init__( aux_stream_list: list[torch.cuda.Stream] | None = None, candidate_block_buffer: torch.Tensor | None = None, engram_layout: EngramLayout | None = None, + run_gemm_rs: bool = False, ): super().__init__() @@ -216,6 +241,10 @@ def __init__( ) if self.use_sequence_parallel: self.attn.wo_b.reduce_results = False + if run_gemm_rs: + # Binds only when wo_b's kernel and shape qualify; otherwise + # forward keeps the separate reduce-scatter below. + self.attn.bind_gemm_rs() self.ffn = DeepseekV4MoE( vllm_config, prefix=f"{prefix}.ffn", @@ -442,7 +471,9 @@ def forward( x = sp_all_gather(x)[: positions.shape[0]] x = self.attn(positions, x, None) - if self.use_sequence_parallel: + # With GEMM-RS bound, the attention output is already this rank's + # sequence-parallel shard (see DeepseekV4Attention._wo_b_proj). + if self.use_sequence_parallel and self.attn.gemm_rs is None: x = sp_reduce_scatter(x) residual, post_mix, res_mix, x, ffn_pre, _ = mhc_shifted_post_pre( @@ -529,6 +560,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.engram_layout = EngramLayout.from_config(config) + # GEMM-RS uses NCCL symmetric-memory multicast, which requires all TP + # ranks to belong to one NVLink domain. Collective: run before layers. + self.run_gemm_rs = maybe_init_gemm_rs(vllm_config, self.use_sequence_parallel) + self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, lambda prefix: DeepseekV4DecoderLayer( @@ -538,6 +573,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): aux_stream_list=aux_stream_list, candidate_block_buffer=self.candidate_block_buffer, engram_layout=self.engram_layout, + run_gemm_rs=self.run_gemm_rs, ), prefix=f"{prefix}.layers", ) diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index 172d2213bf6b..bac7ebac26d0 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -783,7 +783,7 @@ def __init__( ) self.gemm_rs_ar = None if run_gemm_rs_ar: - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import ( + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import ( get_gemm_rs_ar, ) diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 568b6a51be34..1d7625ab10e8 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -301,7 +301,7 @@ def __init__( ) self.gemm_rs_ar = None if run_gemm_rs_ar: - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import ( + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import ( get_gemm_rs_ar, ) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 88197c3a0fd1..ec382ce4f1ee 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -164,50 +164,19 @@ def maybe_init_gemm_rs_ar(vllm_config: VllmConfig, use_sequence_parallel: bool) if not enabled: return False - parallel_config = vllm_config.parallel_config - tp_size = parallel_config.tensor_parallel_size - if parallel_config.use_ubatching: - reason = "ubatching is enabled" - elif vllm_config.model_config.dtype != torch.bfloat16: - reason = "the model dtype is not BF16" - elif not current_platform.is_cuda(): - reason = "the device is not CUDA" - elif not current_platform.is_device_capability_family(100): - reason = "the device is not SM100-family" - elif not 1 < tp_size <= 16: - reason = "TP size is not in the supported range 2-16" - elif 128 % tp_size != 0: - reason = "TP size does not divide 128" - else: - reason = None - - if reason is not None: - logger.warning_once("%s is disabled because %s.", mode, reason) - return False - - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import init_gemm_rs_ar + # The kernel module pulls in cute_dsl, so import it only once the flag + # asks for it. + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import ( + maybe_init_gemm_rs_ar as maybe_init_shared_gemm_rs_ar, + ) config = vllm_config.model_config.hf_text_config - try: - init_gemm_rs_ar( - max_M=vllm_config.scheduler_config.max_num_batched_tokens, - N=config.hidden_size, - all_reduce=all_reduce, - ) - except RuntimeError as e: - logger.warning_once( - "%s is disabled because initialization failed: %s. This may mean " - "the TP ranks do not share one NVLink domain.", - mode, - e, - ) + if not maybe_init_shared_gemm_rs_ar( + vllm_config, N=config.hidden_size, all_reduce=all_reduce + ): return False if all_reduce: - logger.info_once( - "GEMM-AR is enabled. To disable it, set VLLM_KIMI_K3_GEMM_AR=0." - ) - else: - logger.info_once("GEMM-RS is enabled.") + logger.info_once("To disable %s, set VLLM_KIMI_K3_GEMM_AR=0.", mode) return True @@ -276,7 +245,7 @@ def __init__( not use_sequence_parallel and reduce_results ) if use_gemm_rs_ar and run_gemm_rs_ar: - from vllm.models.kimi_k3.nvidia.ops.cute_dsl.gemm_rs_ar import ( + from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import ( get_gemm_rs_ar, ) From 21e53f3b269dc3a14661951a2c46f461541fa229 Mon Sep 17 00:00:00 2001 From: Canlin Date: Fri, 18 Sep 2026 12:57:49 +0800 Subject: [PATCH 07/11] [Bench] Fix the MXFP8 GEMM-RS baseline and add the vLLM reduce-scatter path The online-MXFP8 baseline projections all-reduced inside RowParallelLinear, so the reference partial was already summed and every collective candidate failed the correctness check by the TP factor. Build them with reduce_results=False like sequence-parallel layers do, label assertion failures with the candidate and M, and add a candidate that runs the model's unfused path (projection kernel + sp_reduce_scatter) so the table compares against what the fused kernel actually replaces. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Canlin --- benchmarks/kernels/benchmark_gemm_rs_ar.py | 33 +++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/benchmarks/kernels/benchmark_gemm_rs_ar.py b/benchmarks/kernels/benchmark_gemm_rs_ar.py index 7234cd33bc5a..f817c2687593 100644 --- a/benchmarks/kernels/benchmark_gemm_rs_ar.py +++ b/benchmarks/kernels/benchmark_gemm_rs_ar.py @@ -25,7 +25,10 @@ import torch.distributed._symmetric_memory as symm_mem from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config -from vllm.distributed import cleanup_dist_env_and_memory +from vllm.distributed import ( + cleanup_dist_env_and_memory, + tensor_model_parallel_all_reduce, +) from vllm.distributed.parallel_state import ( get_tp_group, init_distributed_environment, @@ -197,6 +200,9 @@ def make_projection( bias=False, params_dtype=weight.dtype, quant_config=quant_config, + # The unfused baseline wants the local partial, as under sequence + # parallel; the collective candidates reduce it themselves. + reduce_results=False, return_bias=False, ) # Select the fused path before online quantization replaces the weights. @@ -306,6 +312,20 @@ def run() -> torch.Tensor: return run + def make_vllm_gemm_collective( + x: torch.Tensor, linear: RowParallelLinear + ) -> Callable[[], torch.Tensor]: + # The model's unfused sequence-parallel path: the projection's own + # kernel followed by vLLM's reduce-scatter dispatch. + from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter + + def run() -> torch.Tensor: + if all_reduce: + return tensor_model_parallel_all_reduce(linear(x)) + return sp_reduce_scatter(linear(x)) + + return run + def make_fused_gemm_collective( x: torch.Tensor, linear: RowParallelLinear ) -> Callable[[], torch.Tensor]: @@ -362,6 +382,10 @@ def run() -> torch.Tensor: "ldmc_us", [make_torch_ldmc_gemm_collective(x, w) for x, w in pairs], ), + Candidate( + "vllm_us", + [make_vllm_gemm_collective(x, w) for x, w in pairs], + ), Candidate( "gemm_rs_ar_us", [make_fused_gemm_collective(x, w) for x, w in pairs], @@ -390,11 +414,13 @@ def run() -> torch.Tensor: continue actual = candidate.runs[0]() torch.accelerator.synchronize(device) + label = f"{candidate.name} vs ring_ll at M={M}" torch.testing.assert_close( actual[:rows], expected[:rows], rtol=5e-2, atol=4.0, + msg=lambda m, label=label: f"{label}: {m}", ) candidate_graphs = {} @@ -425,6 +451,7 @@ def run() -> torch.Tensor: "best_nccl_us": best_nccl_us, "speedup_vs_ring_ll": times["ring_ll_us"] / times["gemm_rs_ar_us"], "speedup_vs_ldmc": times["ldmc_us"] / times["gemm_rs_ar_us"], + "speedup_vs_vllm": times["vllm_us"] / times["gemm_rs_ar_us"], } @@ -439,18 +466,22 @@ def print_results(results: list[dict[str, float | int | str]]) -> None: "K", "ring_ll_us", "ldmc_us", + "vllm_us", "gemm_rs_ar_us", "speedup_vs_ring_ll", "speedup_vs_ldmc", + "speedup_vs_vllm", ] ] .rename( columns={ "ring_ll_us": f"Torch GEMM + NCCL {collective} (RING_LL) (us)", "ldmc_us": f"Torch GEMM + NCCL {collective} (LDMC) (us)", + "vllm_us": f"vLLM linear + {collective} (us)", "gemm_rs_ar_us": f"GEMM-{collective} (us)", "speedup_vs_ring_ll": "Speedup vs RING_LL", "speedup_vs_ldmc": "Speedup vs LDMC", + "speedup_vs_vllm": "Speedup vs vLLM", } ) .round(3) From 33dd68ee438faf0630b59c928cfc817e02612e7b Mon Sep 17 00:00:00 2001 From: Canlin Date: Sat, 19 Sep 2026 20:06:16 +0800 Subject: [PATCH 08/11] [Bench] Compare GEMM-RS/AR only against vLLM's unfused path Drop the hand-built NCCL RING_LL/LDMC candidates, the bare GEMM and bare collective candidates, and the NCCL symmetric-memory scratch they needed. The baseline is now what a model runs without the fusion: the projection's linear kernel followed by sp_reduce_scatter or tensor_model_parallel_all_reduce, with the collective backend chosen by vLLM's dispatch as in serving (FlashInfer all-reduce first). One table with the latency change, optional --json, no pandas dependency. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Canlin --- benchmarks/kernels/benchmark_gemm_rs_ar.py | 422 +++++---------------- 1 file changed, 105 insertions(+), 317 deletions(-) diff --git a/benchmarks/kernels/benchmark_gemm_rs_ar.py b/benchmarks/kernels/benchmark_gemm_rs_ar.py index f817c2687593..84cd6f939cf7 100644 --- a/benchmarks/kernels/benchmark_gemm_rs_ar.py +++ b/benchmarks/kernels/benchmark_gemm_rs_ar.py @@ -1,28 +1,31 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Benchmark the SM100 GEMM-RS/AR kernel. +"""Benchmark the SM100 GEMM-RS/AR kernel against the model's unfused path. -All ranks must belong to one NVLink domain. For example, run a TP8 sweep over -the Kimi-K3 projections with: +The baseline is what a model runs without the fusion: the projection's own +linear kernel followed by vLLM's collective (``sp_reduce_scatter`` for RS, +``tensor_model_parallel_all_reduce`` for AR), with the collective backend +chosen by vLLM's dispatch exactly as in serving. - torchrun --nproc-per-node=8 benchmarks/kernels/benchmark_gemm_rs_ar.py +All ranks must belong to one NVLink domain. Run a TP8 sweep over the Kimi-K3 +projections with: + + torchrun --nproc-per-node=8 -- benchmarks/kernels/benchmark_gemm_rs_ar.py or the DeepSeek-V4.1 MXFP8 ``wo_b`` projection at TP4 with: - torchrun --nproc-per-node=4 benchmarks/kernels/benchmark_gemm_rs_ar.py \ + torchrun --nproc-per-node=4 -- benchmarks/kernels/benchmark_gemm_rs_ar.py \ --model deepseek_v41 --backend flashinfer_cutedsl """ import argparse +import json import os import statistics from collections.abc import Callable -from dataclasses import dataclass -import pandas as pd import torch import torch.distributed as dist -import torch.distributed._symmetric_memory as symm_mem from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.distributed import ( @@ -36,6 +39,7 @@ ) from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr from vllm.model_executor.layers.linear import RowParallelLinear +from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter # Per model: (hidden size N, global input widths K). Kimi-K3 lists the # shared-expert down-proj and attention O-proj; DeepSeek-V4.1 lists the @@ -46,13 +50,6 @@ } -@dataclass -class Candidate: - name: str - runs: list[Callable[[], torch.Tensor]] - check_correctness: bool = True - - def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument( @@ -102,29 +99,68 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("--warmup-replays", type=int, default=5) parser.add_argument("--samples", type=int, default=20) + parser.add_argument("--json", type=str, help="Also write the results here.") return parser.parse_args() +def make_projection( + gemm_rs_ar: GemmRsAr, + weight: torch.Tensor, + world_size: int, + backend: str, +) -> RowParallelLinear: + """Wrap a BF16 [N, K] weight the way the model would hand it to GEMM-RS. + + For the MXFP8 backends this quantizes online through the FlashInfer + kernel's post-load step, so the fused kernel and the unfused baseline see + the production weight and scale layouts. + """ + N, K = weight.shape + quant_config = None + if backend != "bf16": + from vllm.config.quantization import QuantizationConfigArgs + from vllm.model_executor.layers.quantization.online.base import ( + OnlineQuantizationConfig, + ) + + quant_config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8")) + with torch.device(weight.device): + linear = RowParallelLinear( + K * world_size, + N, + bias=False, + params_dtype=weight.dtype, + quant_config=quant_config, + # The baseline reduces the local partial itself, like the models. + reduce_results=False, + return_bias=False, + ) + # Select the fused path before online quantization replaces the weights. + assert gemm_rs_ar.can_run(linear) + linear.weight = torch.nn.Parameter(weight, requires_grad=False) + linear.quant_method.process_weights_after_loading(linear) + return linear + + def capture_graph( op: Callable[[], torch.Tensor], stream: torch.cuda.Stream, cpu_group: dist.ProcessGroup, -) -> tuple[torch.cuda.CUDAGraph, list[torch.Tensor | None]]: - result: list[torch.Tensor | None] = [None] +) -> torch.cuda.CUDAGraph: stream.wait_stream(torch.cuda.current_stream()) dist.barrier(group=cpu_group) with torch.cuda.stream(stream): for _ in range(3): - result[0] = op() + op() stream.synchronize() dist.barrier(group=cpu_group) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=stream): - result[0] = op() + op() torch.cuda.current_stream().wait_stream(stream) dist.barrier(group=cpu_group) - return graph, result + return graph def benchmark_graphs( @@ -134,23 +170,22 @@ def benchmark_graphs( device_group: dist.ProcessGroup, device_barrier: Callable[[], None], ) -> dict[str, float]: - candidate_names = list(candidate_graphs) + """Median over samples of the slowest rank's replay time, in us.""" + names = list(candidate_graphs) for round_index in range(warmup_replays): - for candidate_index in range(len(candidate_names)): - candidate_id = (round_index + candidate_index) % len(candidate_names) - name = candidate_names[candidate_id] + for offset in range(len(names)): + name = names[(round_index + offset) % len(names)] graphs = candidate_graphs[name] device_barrier() graphs[round_index % len(graphs)].replay() torch.accelerator.synchronize() - timings: dict[str, list[float]] = {name: [] for name in candidate_names} + timings: dict[str, list[float]] = {name: [] for name in names} start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) for sample_index in range(samples): - for candidate_index in range(len(candidate_names)): - candidate_id = (sample_index + candidate_index) % len(candidate_names) - name = candidate_names[candidate_id] + for offset in range(len(names)): + name = names[(sample_index + offset) % len(names)] graphs = candidate_graphs[name] device_barrier() start.record() @@ -168,50 +203,6 @@ def benchmark_graphs( return {name: statistics.median(values) for name, values in timings.items()} -def valid_rows(M: int, local_M: int, rank: int) -> int: - return min(max(M - rank * local_M, 0), local_M) - - -def make_projection( - gemm_rs_ar: GemmRsAr, - weight: torch.Tensor, - world_size: int, - backend: str, -) -> RowParallelLinear: - """Wrap a BF16 [N, K] weight the way the model would hand it to GEMM-RS. - - For the MXFP8 backends this quantizes online through the FlashInfer - kernel's post-load step, so the fused kernel and the unfused baseline see - the production weight and scale layouts. - """ - N, K = weight.shape - quant_config = None - if backend != "bf16": - from vllm.config.quantization import QuantizationConfigArgs - from vllm.model_executor.layers.quantization.online.base import ( - OnlineQuantizationConfig, - ) - - quant_config = OnlineQuantizationConfig(QuantizationConfigArgs(linear="mxfp8")) - with torch.device(weight.device): - linear = RowParallelLinear( - K * world_size, - N, - bias=False, - params_dtype=weight.dtype, - quant_config=quant_config, - # The unfused baseline wants the local partial, as under sequence - # parallel; the collective candidates reduce it themselves. - reduce_results=False, - return_bias=False, - ) - # Select the fused path before online quantization replaces the weights. - assert gemm_rs_ar.can_run(linear) - linear.weight = torch.nn.Parameter(weight, requires_grad=False) - linear.quant_method.process_weights_after_loading(linear) - return linear - - def benchmark_shape( gemm_rs_ar: GemmRsAr, mode: str, @@ -230,8 +221,7 @@ def benchmark_shape( world_size = dist.get_world_size(device_group) rank = dist.get_rank(device_group) device = torch.device("cuda", torch.accelerator.current_device_index()) - padded_M = (M + world_size - 1) // world_size * world_size - local_M = padded_M // world_size + local_M = (M + world_size - 1) // world_size rng = torch.Generator(device=device) rng.manual_seed(1000 + rank * 10 + M + K) @@ -249,263 +239,60 @@ def benchmark_shape( for _ in range(num_workspaces) ] - partial = torch.empty((padded_M, N), dtype=torch.bfloat16, device=device) - symm_partial = symm_mem.empty((padded_M, N), dtype=torch.bfloat16, device=device) - symm_partial_handle = symm_mem.rendezvous(symm_partial, device_group) - collective_inputs = [torch.empty_like(partial) for _ in range(num_workspaces)] - symm_collective_inputs = [] - symm_collective_handles = [] - for _ in range(num_workspaces): - collective_input = symm_mem.empty( - (padded_M, N), - dtype=torch.bfloat16, - device=device, - ) - symm_collective_inputs.append(collective_input) - symm_collective_handles.append( - symm_mem.rendezvous(collective_input, device_group) - ) - torch_output = torch.empty((local_M, N), dtype=torch.bfloat16, device=device) - symm_output = torch.empty_like(torch_output) - gemm_output = torch.empty((M, N), dtype=torch.bfloat16, device=device) - - if padded_M > M: - partial[M:].zero_() - symm_partial[M:].zero_() - - def gemm_into(x: torch.Tensor, linear: RowParallelLinear, out: torch.Tensor): - # BF16 writes straight into the collective buffer; the MXFP8 kernels - # allocate their own output, so the baseline pays one extra copy, as - # the model's unfused path does before its collective. - if backend == "bf16": - torch.mm(x, linear.weight.T, out=out) - else: - out.copy_(linear(x)) - - def make_torch_ring_ll_gemm_collective( - x: torch.Tensor, linear: RowParallelLinear - ) -> Callable[[], torch.Tensor]: - def run() -> torch.Tensor: - gemm_into(x, linear, partial[:M]) - if all_reduce: - dist.all_reduce(partial, group=device_group) - return partial[:M] - dist.reduce_scatter_single(torch_output, partial, group=device_group) - return torch_output - - return run - - def make_torch_ldmc_gemm_collective( - x: torch.Tensor, linear: RowParallelLinear - ) -> Callable[[], torch.Tensor]: - def run() -> torch.Tensor: - gemm_into(x, linear, symm_partial[:M]) - if all_reduce: - dist.all_reduce(symm_partial, group=device_group) - return symm_partial[:M] - dist.reduce_scatter_single( - symm_output, - symm_partial, - group=device_group, - ) - return symm_output - - return run - - def make_vllm_gemm_collective( - x: torch.Tensor, linear: RowParallelLinear - ) -> Callable[[], torch.Tensor]: - # The model's unfused sequence-parallel path: the projection's own - # kernel followed by vLLM's reduce-scatter dispatch. - from vllm.models.common.ops.sequence_parallel import sp_reduce_scatter - - def run() -> torch.Tensor: - if all_reduce: - return tensor_model_parallel_all_reduce(linear(x)) - return sp_reduce_scatter(linear(x)) - - return run - - def make_fused_gemm_collective( - x: torch.Tensor, linear: RowParallelLinear - ) -> Callable[[], torch.Tensor]: - def run() -> torch.Tensor: - return gemm_rs_ar.apply(x, linear) - - return run - - def make_torch_gemm( - x: torch.Tensor, - linear: RowParallelLinear, - ) -> Callable[[], torch.Tensor]: - def run() -> torch.Tensor: - gemm_into(x, linear, gemm_output) - return gemm_output - - return run - - def make_ring_ll_collective( - collective_input: torch.Tensor, - ) -> Callable[[], torch.Tensor]: - def run() -> torch.Tensor: - if all_reduce: - dist.all_reduce(collective_input, group=device_group) - return collective_input - dist.reduce_scatter_single( - torch_output, collective_input, group=device_group - ) - return torch_output - - return run - - def make_ldmc_collective( - collective_input: torch.Tensor, - ) -> Callable[[], torch.Tensor]: - def run() -> torch.Tensor: - if all_reduce: - dist.all_reduce(collective_input, group=device_group) - return collective_input - dist.reduce_scatter_single( - symm_output, collective_input, group=device_group - ) - return symm_output + def unfused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor: + if all_reduce: + return tensor_model_parallel_all_reduce(linear(x)) + return sp_reduce_scatter(linear(x)) - return run + def fused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor: + return gemm_rs_ar.apply(x, linear) - pairs = list(zip(inputs, projections)) - candidates = ( - Candidate( - "ring_ll_us", - [make_torch_ring_ll_gemm_collective(x, w) for x, w in pairs], - ), - Candidate( - "ldmc_us", - [make_torch_ldmc_gemm_collective(x, w) for x, w in pairs], - ), - Candidate( - "vllm_us", - [make_vllm_gemm_collective(x, w) for x, w in pairs], - ), - Candidate( - "gemm_rs_ar_us", - [make_fused_gemm_collective(x, w) for x, w in pairs], - ), - Candidate( - "torch_gemm_us", - [make_torch_gemm(x, w) for x, w in pairs], - check_correctness=False, - ), - Candidate( - "ring_ll_collective_us", - [make_ring_ll_collective(x) for x in collective_inputs], - check_correctness=False, - ), - Candidate( - "ldmc_collective_us", - [make_ldmc_collective(x) for x in symm_collective_inputs], - check_correctness=False, - ), - ) + candidates = {"vllm": unfused, "gemm_rs_ar": fused} - expected = candidates[0].runs[0]() - rows = M if all_reduce else valid_rows(M, local_M, rank) - for candidate in candidates[1:]: - if not candidate.check_correctness: - continue - actual = candidate.runs[0]() - torch.accelerator.synchronize(device) - label = f"{candidate.name} vs ring_ll at M={M}" - torch.testing.assert_close( - actual[:rows], - expected[:rows], - rtol=5e-2, - atol=4.0, - msg=lambda m, label=label: f"{label}: {m}", - ) + expected = unfused(inputs[0], projections[0]) + actual = fused(inputs[0], projections[0]) + torch.accelerator.synchronize(device) + rows = M if all_reduce else min(max(M - rank * local_M, 0), local_M) + torch.testing.assert_close( + actual[:rows], expected[:rows], rtol=5e-2, atol=4.0, msg=f"M={M}, K={K}" + ) candidate_graphs = {} - graph_keepalive: list[object] = [symm_partial_handle, *symm_collective_handles] - for candidate in candidates: + keepalive: list[object] = [] + for name, fn in candidates.items(): stream = torch.cuda.Stream() - bundles = [capture_graph(run, stream, cpu_group) for run in candidate.runs] - candidate_graphs[candidate.name] = [graph for graph, _ in bundles] - graph_keepalive.extend(bundles) - graph_keepalive.append(stream) + candidate_graphs[name] = [ + capture_graph(lambda x=x, w=w, fn=fn: fn(x, w), stream, cpu_group) + for x, w in zip(inputs, projections) + ] + keepalive.append(stream) times = benchmark_graphs( - candidate_graphs, - warmup_replays, - samples, - device_group, - device_barrier, + candidate_graphs, warmup_replays, samples, device_group, device_barrier ) - - best_nccl_us = min(times["ring_ll_collective_us"], times["ldmc_collective_us"]) return { "mode": mode.upper(), "backend": backend, "M": M, "N": N, "K": K, - **times, - "best_nccl_us": best_nccl_us, - "speedup_vs_ring_ll": times["ring_ll_us"] / times["gemm_rs_ar_us"], - "speedup_vs_ldmc": times["ldmc_us"] / times["gemm_rs_ar_us"], - "speedup_vs_vllm": times["vllm_us"] / times["gemm_rs_ar_us"], + "vllm_us": times["vllm"], + "gemm_rs_ar_us": times["gemm_rs_ar"], + "latency_change": (times["gemm_rs_ar"] - times["vllm"]) / times["vllm"], } def print_results(results: list[dict[str, float | int | str]]) -> None: - results_df = pd.DataFrame(results) - collective = str(results_df["mode"].iloc[0]) - end_to_end = ( - results_df[ - [ - "M", - "N", - "K", - "ring_ll_us", - "ldmc_us", - "vllm_us", - "gemm_rs_ar_us", - "speedup_vs_ring_ll", - "speedup_vs_ldmc", - "speedup_vs_vllm", - ] - ] - .rename( - columns={ - "ring_ll_us": f"Torch GEMM + NCCL {collective} (RING_LL) (us)", - "ldmc_us": f"Torch GEMM + NCCL {collective} (LDMC) (us)", - "vllm_us": f"vLLM linear + {collective} (us)", - "gemm_rs_ar_us": f"GEMM-{collective} (us)", - "speedup_vs_ring_ll": "Speedup vs RING_LL", - "speedup_vs_ldmc": "Speedup vs LDMC", - "speedup_vs_vllm": "Speedup vs vLLM", - } + collective = results[0]["mode"] + backend = results[0]["backend"] + print(f"### GEMM-{collective} vs vLLM unfused path ({backend} weights)") + print(f"| M | N | K | vLLM (us) | GEMM-{collective} (us) | latency |") + print("|---:|---:|---:|---:|---:|---:|") + for r in results: + print( + f"| {r['M']} | {r['N']} | {r['K']} | {r['vllm_us']:.1f} | " + f"{r['gemm_rs_ar_us']:.1f} | {r['latency_change']:+.1%} |" ) - .round(3) - ) - components = ( - results_df[["M", "N", "K", "torch_gemm_us", "best_nccl_us", "gemm_rs_ar_us"]] - .rename( - columns={ - "torch_gemm_us": "Torch GEMM (us)", - "best_nccl_us": f"NCCL {collective} (best) (us)", - "gemm_rs_ar_us": f"GEMM-{collective} (us)", - } - ) - .round(2) - ) - - print(f"### {collective} end-to-end latency") - print(end_to_end.to_markdown(index=False)) - print(f"\n### {collective} component latency") - print(components.to_markdown(index=False)) - print( - f"\nNCCL {collective} (best) is the faster of RING_LL and LDMC " - "for each shape.\n" - ) def main() -> None: @@ -530,10 +317,7 @@ def main() -> None: K_values = args.k k_alignment = 64 if args.backend == "bf16" else 128 assert all(K % k_alignment == 0 for K in K_values) - # Reserve symmetric memory for the NCCL-managed benchmark allocations. - os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0" - # NCCL-managed symmetric allocations select the NVLS/LDMC collective path. - symm_mem.set_backend("NCCL") + config = VllmConfig() config.model_config = ModelConfig(dtype="bfloat16") # Pin the online-MXFP8 kernel the projections quantize through. @@ -541,11 +325,12 @@ def main() -> None: "auto" if args.backend == "bf16" else args.backend ) with set_current_vllm_config(config): + # Builds the TP communicator as in serving, so the baseline's + # collective is whatever vLLM's dispatch selects (FlashInfer all-reduce + # first, then NCCL symmetric memory, custom all-reduce, pynccl). initialize_model_parallel(tensor_model_parallel_size=world_size) tp_group = get_tp_group() - group_warmup = torch.zeros(1, device=torch.accelerator.current_device_index()) - dist.all_reduce(group_warmup, group=tp_group.device_group) pynccl_comm = tp_group.device_communicator.pynccl_comm assert pynccl_comm is not None sync_input = torch.zeros(1, device=torch.accelerator.current_device_index()) @@ -584,6 +369,9 @@ def device_barrier() -> None: if tp_group.rank_in_group == 0: print_results(results) + if args.json: + with open(args.json, "w") as f: + json.dump(results, f, indent=2) dist.barrier(group=tp_group.cpu_group) cleanup_dist_env_and_memory() From 50cd05b430ca9454118493f420697cf6351d5ff3 Mon Sep 17 00:00:00 2001 From: Canlin Date: Mon, 21 Sep 2026 12:21:06 +0800 Subject: [PATCH 09/11] [Bench] Time the bare GEMM and collective as the GEMM-RS/AR floor Add the GEMM alone and the collective alone (both through vLLM's dispatch: the projection's MXFP8 linear kernel, sp_reduce_scatter / tensor_model_parallel_all_reduce on the GEMM's own partial) as candidates. Their max is the floor a perfectly overlapped fused kernel could reach; the table now prints both and the fused kernel's gap to the floor next to its change against the unfused path. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Canlin --- benchmarks/kernels/benchmark_gemm_rs_ar.py | 52 ++++++++++++++++------ 1 file changed, 38 insertions(+), 14 deletions(-) diff --git a/benchmarks/kernels/benchmark_gemm_rs_ar.py b/benchmarks/kernels/benchmark_gemm_rs_ar.py index 84cd6f939cf7..3b20d6dad3e1 100644 --- a/benchmarks/kernels/benchmark_gemm_rs_ar.py +++ b/benchmarks/kernels/benchmark_gemm_rs_ar.py @@ -5,7 +5,9 @@ The baseline is what a model runs without the fusion: the projection's own linear kernel followed by vLLM's collective (``sp_reduce_scatter`` for RS, ``tensor_model_parallel_all_reduce`` for AR), with the collective backend -chosen by vLLM's dispatch exactly as in serving. +chosen by vLLM's dispatch exactly as in serving. The same two kernels are also +timed bare (the GEMM alone, the collective alone on a precomputed partial); +their max is the floor a perfectly overlapped fused kernel could reach. All ranks must belong to one NVLink domain. Run a TP8 sweep over the Kimi-K3 projections with: @@ -239,15 +241,23 @@ def benchmark_shape( for _ in range(num_workspaces) ] - def unfused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor: + def collective(partial: torch.Tensor) -> torch.Tensor: if all_reduce: - return tensor_model_parallel_all_reduce(linear(x)) - return sp_reduce_scatter(linear(x)) + return tensor_model_parallel_all_reduce(partial) + return sp_reduce_scatter(partial) + + def gemm(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor: + return linear(x) + + def unfused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor: + return collective(linear(x)) def fused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor: return gemm_rs_ar.apply(x, linear) - candidates = {"vllm": unfused, "gemm_rs_ar": fused} + # The bare collective runs on the GEMM's own output so it sees the same + # message size and layout the unfused path hands to the collective. + partials = [linear(x) for x, linear in zip(inputs, projections)] expected = unfused(inputs[0], projections[0]) actual = fused(inputs[0], projections[0]) @@ -257,28 +267,37 @@ def fused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor: actual[:rows], expected[:rows], rtol=5e-2, atol=4.0, msg=f"M={M}, K={K}" ) + candidates: dict[str, list[Callable[[], torch.Tensor]]] = { + "gemm": [lambda x=x, w=w: gemm(x, w) for x, w in zip(inputs, projections)], + "collective": [lambda p=p: collective(p) for p in partials], + "vllm": [lambda x=x, w=w: unfused(x, w) for x, w in zip(inputs, projections)], + "gemm_rs_ar": [lambda x=x, w=w: fused(x, w) for x, w in zip(inputs, projections)], + } + candidate_graphs = {} keepalive: list[object] = [] - for name, fn in candidates.items(): + for name, ops in candidates.items(): stream = torch.cuda.Stream() - candidate_graphs[name] = [ - capture_graph(lambda x=x, w=w, fn=fn: fn(x, w), stream, cpu_group) - for x, w in zip(inputs, projections) - ] + candidate_graphs[name] = [capture_graph(op, stream, cpu_group) for op in ops] keepalive.append(stream) times = benchmark_graphs( candidate_graphs, warmup_replays, samples, device_group, device_barrier ) + floor = max(times["gemm"], times["collective"]) return { "mode": mode.upper(), "backend": backend, "M": M, "N": N, "K": K, + "gemm_us": times["gemm"], + "collective_us": times["collective"], + "floor_us": floor, "vllm_us": times["vllm"], "gemm_rs_ar_us": times["gemm_rs_ar"], "latency_change": (times["gemm_rs_ar"] - times["vllm"]) / times["vllm"], + "floor_gap": (times["gemm_rs_ar"] - floor) / floor, } @@ -286,12 +305,17 @@ def print_results(results: list[dict[str, float | int | str]]) -> None: collective = results[0]["mode"] backend = results[0]["backend"] print(f"### GEMM-{collective} vs vLLM unfused path ({backend} weights)") - print(f"| M | N | K | vLLM (us) | GEMM-{collective} (us) | latency |") - print("|---:|---:|---:|---:|---:|---:|") + print( + f"| M | N | K | GEMM (us) | {collective} (us) | floor (us) | vLLM (us) | " + f"GEMM-{collective} (us) | vs vLLM | vs floor |" + ) + print("|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|") for r in results: print( - f"| {r['M']} | {r['N']} | {r['K']} | {r['vllm_us']:.1f} | " - f"{r['gemm_rs_ar_us']:.1f} | {r['latency_change']:+.1%} |" + f"| {r['M']} | {r['N']} | {r['K']} | {r['gemm_us']:.1f} | " + f"{r['collective_us']:.1f} | {r['floor_us']:.1f} | {r['vllm_us']:.1f} | " + f"{r['gemm_rs_ar_us']:.1f} | {r['latency_change']:+.1%} | " + f"{r['floor_gap']:+.1%} |" ) From 8f272b25bf20da5f9360ffba8acda100bd50a784 Mon Sep 17 00:00:00 2001 From: Canlin Date: Mon, 21 Sep 2026 12:22:12 +0800 Subject: [PATCH 10/11] [Bench] Format the candidate table Co-Authored-By: Claude Fable 5.1 Signed-off-by: Canlin --- benchmarks/kernels/benchmark_gemm_rs_ar.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/benchmarks/kernels/benchmark_gemm_rs_ar.py b/benchmarks/kernels/benchmark_gemm_rs_ar.py index 3b20d6dad3e1..3e47f61c74f0 100644 --- a/benchmarks/kernels/benchmark_gemm_rs_ar.py +++ b/benchmarks/kernels/benchmark_gemm_rs_ar.py @@ -271,7 +271,9 @@ def fused(x: torch.Tensor, linear: RowParallelLinear) -> torch.Tensor: "gemm": [lambda x=x, w=w: gemm(x, w) for x, w in zip(inputs, projections)], "collective": [lambda p=p: collective(p) for p in partials], "vllm": [lambda x=x, w=w: unfused(x, w) for x, w in zip(inputs, projections)], - "gemm_rs_ar": [lambda x=x, w=w: fused(x, w) for x, w in zip(inputs, projections)], + "gemm_rs_ar": [ + lambda x=x, w=w: fused(x, w) for x, w in zip(inputs, projections) + ], } candidate_graphs = {} From 6cd2dba683d3432ec8e04f50ef4f1decb46a707e Mon Sep 17 00:00:00 2001 From: Canlin Date: Mon, 21 Sep 2026 16:49:32 +0800 Subject: [PATCH 11/11] Gate Kimi-K3 and DeepSeek-V4.1 GEMM-RS behind one VLLM_ENABLE_GEMM_RS flag Both models bind the same SM100 GEMM-RS kernel, so replace the per-model VLLM_KIMI_K3_GEMM_RS and VLLM_DSV41_GEMM_RS switches (both default off) with a single VLLM_ENABLE_GEMM_RS. VLLM_KIMI_K3_GEMM_AR is unchanged. Co-Authored-By: Claude Fable 5.1 Signed-off-by: Canlin --- vllm/envs.py | 15 ++++++--------- vllm/models/deepseek_v41/nvidia/model.py | 4 ++-- vllm/models/kimi_k3/nvidia/model.py | 6 +++--- 3 files changed, 11 insertions(+), 14 deletions(-) diff --git a/vllm/envs.py b/vllm/envs.py index e443b9680b02..4c8b1ef92491 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -213,8 +213,7 @@ VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False VLLM_KIMI_K3_GEMM_AR: bool = True - VLLM_KIMI_K3_GEMM_RS: bool = False - VLLM_DSV41_GEMM_RS: bool = False + VLLM_ENABLE_GEMM_RS: bool = False VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True VLLM_USE_FLASHINFER_MOE_INT4: bool = False VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None @@ -1623,13 +1622,11 @@ def _resolve_rust_cli_path() -> str | None: # Use the SM100 BF16 GEMM-AR kernel for eligible Kimi-K3 row-parallel # attention projections. All TP ranks must belong to one NVLink domain. "VLLM_KIMI_K3_GEMM_AR": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_AR", "1"))), - # Use the SM100 BF16 GEMM-RS kernel for eligible Kimi-K3 sequence-parallel - # row-parallel projections. All TP ranks must belong to one NVLink domain. - "VLLM_KIMI_K3_GEMM_RS": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_RS", "0"))), - # Fuse the DeepSeek-V4.1 sequence-parallel ``wo_b`` projection with its - # TP reduce-scatter using the SM100 BF16/MXFP8 GEMM-RS kernel. All TP - # ranks must belong to one NVLink domain. Off until model evals land. - "VLLM_DSV41_GEMM_RS": lambda: bool(int(os.getenv("VLLM_DSV41_GEMM_RS", "0"))), + # Fuse eligible sequence-parallel row-parallel projections with their TP + # reduce-scatter using the SM100 BF16/MXFP8 GEMM-RS kernel (Kimi-K3 + # attention/shared-expert projections, DeepSeek-V4.1 ``wo_b``). All TP + # ranks must belong to one NVLink domain. + "VLLM_ENABLE_GEMM_RS": lambda: bool(int(os.getenv("VLLM_ENABLE_GEMM_RS", "0"))), # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 3ce69b6282ce..12a39d3b4bbc 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -202,7 +202,7 @@ def maybe_init_gemm_rs(vllm_config: VllmConfig, use_sequence_parallel: bool) -> (the DSpark drafter reuses it), and its NVLink multicast rendezvous is collective, so every TP rank must take the same decision here. """ - if not (use_sequence_parallel and envs.VLLM_DSV41_GEMM_RS): + if not (use_sequence_parallel and envs.VLLM_ENABLE_GEMM_RS): return False # The kernel module pulls in cute_dsl, so import it only once opted in. @@ -213,7 +213,7 @@ def maybe_init_gemm_rs(vllm_config: VllmConfig, use_sequence_parallel: bool) -> hidden_size = vllm_config.model_config.hf_config.hidden_size if not maybe_init_gemm_rs_ar(vllm_config, N=hidden_size, all_reduce=False): return False - logger.info_once("To disable DeepSeek-V4.1 GEMM-RS, set VLLM_DSV41_GEMM_RS=0.") + logger.info_once("To disable DeepSeek-V4.1 GEMM-RS, set VLLM_ENABLE_GEMM_RS=0.") return True diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 595f9e9ba884..9ce3f1990efa 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -160,7 +160,7 @@ def maybe_init_gemm_rs_ar(vllm_config: VllmConfig, use_sequence_parallel: bool) # its singleton to exactly one mode. all_reduce = not use_sequence_parallel mode = "GEMM-AR" if all_reduce else "GEMM-RS" - enabled = envs.VLLM_KIMI_K3_GEMM_AR if all_reduce else envs.VLLM_KIMI_K3_GEMM_RS + enabled = envs.VLLM_KIMI_K3_GEMM_AR if all_reduce else envs.VLLM_ENABLE_GEMM_RS if not enabled: return False @@ -175,8 +175,8 @@ def maybe_init_gemm_rs_ar(vllm_config: VllmConfig, use_sequence_parallel: bool) vllm_config, N=config.hidden_size, all_reduce=all_reduce ): return False - if all_reduce: - logger.info_once("To disable %s, set VLLM_KIMI_K3_GEMM_AR=0.", mode) + flag = "VLLM_KIMI_K3_GEMM_AR" if all_reduce else "VLLM_ENABLE_GEMM_RS" + logger.info_once("To disable %s, set %s=0.", mode, flag) return True