From 9676182482d210a43790ed1bea3edc3bbf194b7b Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Mon, 25 May 2026 21:11:42 -0500 Subject: [PATCH 01/27] use tilelang mhc on rocm Signed-off-by: tjtanaa --- tests/kernels/test_mhc_kernels.py | 123 +++++++++++++++++++- vllm/_tilelang_ops.py | 91 ++++++++------- vllm/model_executor/kernels/mhc/tilelang.py | 95 ++++++++++----- vllm/model_executor/layers/mhc.py | 74 ++++++++++-- vllm/models/deepseek_v4/amd/model.py | 12 +- vllm/platforms/cuda.py | 8 ++ vllm/platforms/interface.py | 6 + 7 files changed, 326 insertions(+), 83 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 81fceeceac10..a2e76e9f8acf 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -92,8 +92,88 @@ def hc_head_ref( @pytest.mark.skipif( - not current_platform.is_cuda(), - reason="CUDA required", + not (current_platform.is_cuda() or current_platform.is_rocm()), + reason="CUDA or ROCm required", +) +@pytest.mark.parametrize("num_tokens", [4, 256]) +@pytest.mark.parametrize("hidden_size", [1280, 7168]) +@pytest.mark.parametrize("hc_mult", [4]) +def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult): + torch.set_default_device(DEVICE) + set_random_seed(0) + + residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) + hc_mult2 = hc_mult * hc_mult + hc_mult3 = 2 * hc_mult + hc_mult2 + fn = ( + torch.randn((hc_mult3, hc_mult, hidden_size), dtype=torch.float) + * 1e-4 + * (1 + torch.arange(hc_mult).mul(0.01).view(1, -1, 1)) + ).flatten(1, 2) + hc_scale = torch.randn((3,), dtype=torch.float) * 0.1 + hc_base = torch.randn((hc_mult3,), dtype=torch.float) * 0.1 + + hc_sinkhorn_eps = hc_pre_eps = rms_eps = 1e-6 + sinkhorn_repeat = 20 + hc_post_alpha = 1.0 + + ref = mhc_pre_ref( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_alpha, + sinkhorn_repeat, + ) + out = torch.ops.vllm.mhc_pre_tilelang( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_alpha, + sinkhorn_repeat, + ) + + for actual, expected in zip(out, ref, strict=True): + torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2) + + +@pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_rocm()), + reason="CUDA or ROCm required", +) +@pytest.mark.parametrize("num_tokens", [4, 256]) +@pytest.mark.parametrize("hidden_size", [1280, 7168]) +@pytest.mark.parametrize("hc_mult", [4]) +def test_mhc_post_tilelang(num_tokens, hidden_size, hc_mult): + torch.set_default_device(DEVICE) + set_random_seed(0) + + x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16) + residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) + post_layer_mix = torch.randn((num_tokens, hc_mult, 1), dtype=torch.float32) + comb_res_mix = torch.randn((num_tokens, hc_mult, hc_mult), dtype=torch.float32) + + ref = mhc_post_ref(x, residual, post_layer_mix, comb_res_mix) + out = torch.ops.vllm.mhc_post_tilelang( + x, + residual, + post_layer_mix, + comb_res_mix, + ) + + torch.testing.assert_close(out, ref, atol=5e-2, rtol=1e-2) + + +@pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_rocm()), + reason="CUDA or ROCm required", ) @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128]) @pytest.mark.parametrize("hidden_size", [4096, 7168]) @@ -196,3 +276,42 @@ def test_hc_head_triton(num_tokens, hidden_size, hc_mult): out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps) torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2) + + +@pytest.mark.skipif( + not (current_platform.is_cuda() or current_platform.is_rocm()), + reason="CUDA or ROCm required", +) +@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128]) +@pytest.mark.parametrize("hidden_size", [4096, 7168]) +@pytest.mark.parametrize("hc_mult", [4]) +def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult): + torch.set_default_device(DEVICE) + set_random_seed(0) + + residual = torch.randn((num_tokens, hc_mult, hidden_size), dtype=torch.bfloat16) + fn = torch.randn((hc_mult, hc_mult * hidden_size), dtype=torch.float32) * 1e-4 + hc_scale = torch.randn((1,), dtype=torch.float32) * 0.1 + hc_base = torch.randn((hc_mult,), dtype=torch.float32) * 0.1 + rms_eps = hc_eps = 1e-6 + + out = torch.empty((num_tokens, hidden_size), dtype=torch.bfloat16) + out.fill_(float("nan")) + + result = torch.ops.vllm.hc_head_fused_kernel_tilelang( + residual, + fn, + hc_scale, + hc_base, + out, + hidden_size, + rms_eps, + hc_eps, + hc_mult, + ) + + assert result is None + assert not torch.isnan(out).any() + + out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps) + torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2) diff --git a/vllm/_tilelang_ops.py b/vllm/_tilelang_ops.py index aa742fe50320..9f2827cceee3 100644 --- a/vllm/_tilelang_ops.py +++ b/vllm/_tilelang_ops.py @@ -10,8 +10,20 @@ from vllm.utils.import_utils import has_tilelang from vllm.utils.math_utils import cdiv -# tilelang is only available on CUDA platforms -if TYPE_CHECKING or current_platform.is_cuda(): +def _is_tilelang_platform() -> bool: + return current_platform.is_cuda() or current_platform.is_rocm() + + +def _is_pdl_supported() -> bool: + is_arch_support_pdl = getattr(current_platform, "is_arch_support_pdl", None) + if not callable(is_arch_support_pdl): + return False + return is_arch_support_pdl() + + +# TileLang is used for MHC on CUDA and ROCm. Keep non-GPU imports cheap so +# registering the Python wrapper modules does not require TileLang everywhere. +if TYPE_CHECKING or _is_tilelang_platform(): if not has_tilelang(): raise ImportError( "tilelang is required for mhc but is not installed. Install it with " @@ -23,6 +35,7 @@ tilelang = None # type: ignore[assignment] T = None # type: ignore[assignment] +ENABLE_PDL = _is_pdl_supported() @cache def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: @@ -37,12 +50,17 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: return split_k +pass_configs = { + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, +} + +if current_platform.is_cuda(): + pass_configs[tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL] = 10 + + @tilelang.jit( - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10, - }, + pass_configs=pass_configs, ) def mhc_pre_big_fuse_tilelang( gemm_out_mul, @@ -78,7 +96,8 @@ def mhc_pre_big_fuse_tilelang( layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type] with T.Kernel(num_tokens, threads=96) as i: - T.pdl_sync() + if ENABLE_PDL: + T.pdl_sync() ################################################################## # _pre_norm_fn_fwd_norm rms = T.alloc_fragment(1, T.float32) @@ -174,18 +193,16 @@ def mhc_pre_big_fuse_tilelang( ol[i1_h] += pre * xl[i_hc, i1_h] T.copy(ol, layer_input[i, i0_h * hidden_block]) - T.pdl_trigger() + + if ENABLE_PDL: + T.pdl_trigger() # Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478 @tilelang.jit( - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10, - }, + pass_configs=pass_configs, ) def mhc_pre_big_fuse_with_norm_tilelang( gemm_out_mul, @@ -230,7 +247,8 @@ def mhc_pre_big_fuse_with_norm_tilelang( T.clear(mixes) rms[0] = 0 - T.pdl_sync() + if ENABLE_PDL: + T.pdl_sync() for i_split in T.serial(n_splits): rms[0] += gemm_out_sqrsum[i_split, i] @@ -341,15 +359,12 @@ def mhc_pre_big_fuse_with_norm_tilelang( T.copy(ol, layer_input[i, i0_h * hidden_block]) - T.pdl_trigger() + if ENABLE_PDL: + T.pdl_trigger() @tilelang.jit( - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10, - }, + pass_configs=pass_configs, ) def mhc_fused_tilelang( comb_mix, @@ -390,8 +405,8 @@ def mhc_fused_tilelang( with T.Kernel(m, n_tiles, split_k, threads=n_thr) as (i_n, i_nt, i_ks): tid = T.get_thread_binding() - warp_id = T.get_warp_idx() - lane = T.get_lane_idx() + warp_id = tid // 32 + lane = tid % 32 s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32) s_post = T.alloc_shared((hc,), T.float32) @@ -407,7 +422,8 @@ def mhc_fused_tilelang( T.clear(sqr) h_split_start = i_ks * h_per_split - T.pdl_sync() + if ENABLE_PDL: + T.pdl_sync() T.copy(post_mix[i_n, 0], s_post) T.copy(comb_mix[i_n, 0, 0], s_comb) @@ -466,15 +482,12 @@ def mhc_fused_tilelang( v2 += s_warp[w, tile_n] rp_out[i_ks, i_n] = v2 - T.pdl_trigger() + if ENABLE_PDL: + T.pdl_trigger() @tilelang.jit( - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10, - }, + pass_configs=pass_configs, ) def mhc_post_tilelang( a, @@ -507,7 +520,8 @@ def mhc_post_tilelang( a_local = T.alloc_fragment((hc, hc), T.float32) c_local = T.alloc_fragment(hc, T.float32) - T.pdl_sync() + if ENABLE_PDL: + T.pdl_sync() T.copy(a[i_n, 0, 0], a_local) T.copy(c[i_n, 0], c_local) @@ -523,15 +537,12 @@ def mhc_post_tilelang( x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h] T.copy(x_local, x[i_n, 0, i0_h * h_blk]) - T.pdl_trigger() + if ENABLE_PDL: + T.pdl_trigger() @tilelang.jit( - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10, - }, + pass_configs=pass_configs, ) def hc_head_fuse_tilelang( residual, @@ -566,7 +577,8 @@ def hc_head_fuse_tilelang( out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef,valid-type] with T.Kernel(num_tokens, threads=n_thr) as i: - T.pdl_sync() + if ENABLE_PDL: + T.pdl_sync() # ------------------------------------------------------------------ # Pass 1 – for each residual channel m_c and h_block: @@ -624,4 +636,5 @@ def hc_head_fuse_tilelang( T.copy(ol, out[i, i0_h * h_block], disable_tma=True) - T.pdl_trigger() + if ENABLE_PDL: + T.pdl_trigger() diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index c242fef2d026..3000cd402dd1 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -5,6 +5,25 @@ from vllm.utils.torch_utils import direct_register_custom_op +def _can_use_deep_gemm_hc_prenorm() -> bool: + from vllm.utils.deep_gemm import is_deep_gemm_supported + + return is_deep_gemm_supported() + + +def _torch_hc_prenorm_gemm( + x: torch.Tensor, + fn: torch.Tensor, + out: torch.Tensor, + sqrsum: torch.Tensor, +) -> None: + assert out.shape[0] == 1 + assert sqrsum.shape[0] == 1 + x_float = x.float() + out[0].copy_(x_float @ fn.t()) + sqrsum[0].copy_(x_float.square().sum(dim=-1)) + + def mhc_pre_tilelang( residual: torch.Tensor, fn: torch.Tensor, @@ -80,10 +99,16 @@ def mhc_pre_tilelang( residual_flat = residual.view(-1, hc_mult, hidden_size) num_tokens = residual_flat.shape[0] - # these numbers are from deepgemm kernel impl - block_k = 64 - block_m = 64 - n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m)) + use_deep_gemm = _can_use_deep_gemm_hc_prenorm() + if use_deep_gemm: + # these numbers are from deepgemm kernel impl + block_k = 64 + block_m = 64 + n_splits = compute_num_split( + block_k, hc_hidden_size, cdiv(num_tokens, block_m) + ) + else: + n_splits = 1 post_mix = torch.empty( num_tokens, hc_mult, dtype=torch.float32, device=residual.device @@ -102,13 +127,17 @@ def mhc_pre_tilelang( n_splits, num_tokens, dtype=torch.float32, device=residual.device ) - tf32_hc_prenorm_gemm( - residual_flat.view(num_tokens, hc_mult * hidden_size), - fn, - gemm_out_mul, - gemm_out_sqrsum, - n_splits, - ) + residual_2d = residual_flat.view(num_tokens, hc_mult * hidden_size) + if use_deep_gemm: + tf32_hc_prenorm_gemm( + residual_2d, + fn, + gemm_out_mul, + gemm_out_sqrsum, + n_splits, + ) + else: + _torch_hc_prenorm_gemm(residual_2d, fn, gemm_out_mul, gemm_out_sqrsum) if norm_weight is None: mhc_pre_big_fuse_tilelang( @@ -304,16 +333,22 @@ def mhc_fused_post_pre_tilelang( post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult) comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult) - fma_token_threshold = 16 - if num_tokens <= fma_token_threshold: + use_deep_gemm = _can_use_deep_gemm_hc_prenorm() + use_small_fma = num_tokens <= 16 + if use_small_fma: # TODO(gnovack): investigate autotuning these heuristics tile_n = 2 if num_tokens < 8 else 3 n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4 else: - # these number are from deepgemm kernel impl - block_k = 64 - block_m = 64 - n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m)) + if use_deep_gemm: + # these number are from deepgemm kernel impl + block_k = 64 + block_m = 64 + n_splits = compute_num_split( + block_k, hc_hidden_size, cdiv(num_tokens, block_m) + ) + else: + n_splits = 1 gemm_out_mul = torch.empty( n_splits, @@ -348,7 +383,7 @@ def mhc_fused_post_pre_tilelang( device=residual.device, ) - if num_tokens <= fma_token_threshold: + if use_small_fma: mhc_fused_tilelang( comb_res_mix_flat, residual_flat, @@ -375,15 +410,21 @@ def mhc_fused_post_pre_tilelang( residual.shape[-1], ) - from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm - - tf32_hc_prenorm_gemm( - residual_cur.view(num_tokens, hc_mult * hidden_size), - fn, - gemm_out_mul, - gemm_out_sqrsum, - n_splits, - ) + residual_cur_2d = residual_cur.view(num_tokens, hc_mult * hidden_size) + if use_deep_gemm: + from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm + + tf32_hc_prenorm_gemm( + residual_cur_2d, + fn, + gemm_out_mul, + gemm_out_sqrsum, + n_splits, + ) + else: + _torch_hc_prenorm_gemm( + residual_cur_2d, fn, gemm_out_mul, gemm_out_sqrsum + ) if norm_weight is None: mhc_pre_big_fuse_tilelang( diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index cb3965f764ba..ecf3b7a3c024 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -3,6 +3,7 @@ import torch # this import will also register the custom ops +# import vllm.model_executor.kernels.mhc # noqa: F401 import vllm.model_executor.kernels.mhc as mhc_kernels from vllm.model_executor.custom_op import CustomOp @@ -85,7 +86,18 @@ def forward_hip( # sinkhorn_repeat, # ) # else: - return mhc_kernels.mhc_pre_torch( + # return mhc_kernels.mhc_pre_torch( + # residual, + # fn, + # hc_scale, + # hc_base, + # rms_eps, + # hc_pre_eps, + # hc_sinkhorn_eps, + # hc_post_mult_value, + # sinkhorn_repeat, + # ) + return torch.ops.vllm.mhc_pre_tilelang( residual, fn, hc_scale, @@ -95,6 +107,9 @@ def forward_hip( hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, + n_splits, + norm_weight, + norm_eps, ) def forward_native(self, *args, **kwargs): @@ -147,11 +162,14 @@ def forward_hip( # comb_res_mix, # ) # else: - return mhc_kernels.mhc_post_torch( - x, - residual, - post_layer_mix, - comb_res_mix, + # return mhc_kernels.mhc_post_torch( + # x, + # residual, + # post_layer_mix, + # comb_res_mix, + # ) + return torch.ops.vllm.mhc_post_tilelang( + x, residual, post_layer_mix, comb_res_mix ) def forward_native(self, *args, **kwargs): @@ -220,7 +238,7 @@ def forward_hip( out = torch.empty( num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device ) - torch.ops.vllm.hc_head_triton( + torch.ops.vllm.hc_head_fused_kernel_tilelang( hs_flat, hc_fn, hc_scale, @@ -290,10 +308,46 @@ def forward_cuda( norm_eps, ) - def forward_hip(self, *args, **kwargs): - raise NotImplementedError( - "Hip implementation of mhc_fused_post_pre is not available" + def forward_hip( + self, + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + tile_n: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.ops.vllm.mhc_fused_post_pre_tilelang( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + tile_n, + norm_weight, + norm_eps, ) + # raise NotImplementedError( + # "Hip implementation of mhc_fused_post_pre is not available" + # ) def forward_native(self, *args, **kwargs): raise NotImplementedError( diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 1540667d1a4b..1680d10a192e 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -1195,10 +1195,10 @@ def forward( ) -> tuple[ torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None ]: - if current_platform.is_rocm(): - return self._forward_rocm( - x, positions, input_ids, post_mix, res_mix, residual - ) + # if current_platform.is_rocm(): + # return self._forward_rocm( + # x, positions, input_ids, post_mix, res_mix, residual + # ) return self._forward_cuda(x, positions, input_ids, post_mix, res_mix, residual) @@ -1361,7 +1361,9 @@ def forward( res_mix, residual, ) - if layer is not None and current_platform.is_cuda(): + if layer is not None and ( + current_platform.is_cuda() or current_platform.is_rocm() + ): hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix) if not get_pp_group().is_last_rank: diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 4a5be741d06b..b5cfce40b1b9 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -592,6 +592,14 @@ def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConf default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm ) + @classmethod + def is_arch_support_pdl(cls) -> bool: + try: + device = torch.cuda.current_device() + major, _ = torch.cuda.get_device_capability(device) + except Exception: + return False + return major >= 9 # NVML utils # Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`, diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 9a93ef9f82a7..8c1046881113 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -1016,6 +1016,12 @@ def get_default_ir_op_priority( # Native always used by default. Platforms can override this behavior. return IrOpPriorityConfig.with_default(["native"]) + @classmethod + def is_arch_support_pdl(cls) -> bool: + """ + Does the current platform support PDL (Programmatic Dependent Launch)? + """ + return False class UnspecifiedPlatform(Platform): _enum = PlatformEnum.UNSPECIFIED From 5be8268972dcad982aab020898e3eb9b1e80bc0c Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Mon, 25 May 2026 21:21:36 -0500 Subject: [PATCH 02/27] clean up mhc test Signed-off-by: tjtanaa --- tests/kernels/test_mhc_kernels.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index a2e76e9f8acf..597f8e1ff7cd 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -6,6 +6,7 @@ import vllm.model_executor.kernels.mhc # noqa: F401 from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +from vllm.utils.import_utils import has_tilelang DEVICE = current_platform.device_type @@ -92,11 +93,11 @@ def hc_head_ref( @pytest.mark.skipif( - not (current_platform.is_cuda() or current_platform.is_rocm()), - reason="CUDA or ROCm required", + not (current_platform.is_cuda_alike() and has_tilelang()), + reason="CUDA or ROCm and tilelang required", ) -@pytest.mark.parametrize("num_tokens", [4, 256]) -@pytest.mark.parametrize("hidden_size", [1280, 7168]) +@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128]) +@pytest.mark.parametrize("hidden_size", [4096, 7168]) @pytest.mark.parametrize("hc_mult", [4]) def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult): torch.set_default_device(DEVICE) @@ -145,11 +146,11 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult): @pytest.mark.skipif( - not (current_platform.is_cuda() or current_platform.is_rocm()), - reason="CUDA or ROCm required", + not (current_platform.is_cuda_alike() and has_tilelang()), + reason="CUDA or ROCm and tilelang required", ) -@pytest.mark.parametrize("num_tokens", [4, 256]) -@pytest.mark.parametrize("hidden_size", [1280, 7168]) +@pytest.mark.parametrize("num_tokens", [1, 4, 8, 128]) +@pytest.mark.parametrize("hidden_size", [4096, 7168]) @pytest.mark.parametrize("hc_mult", [4]) def test_mhc_post_tilelang(num_tokens, hidden_size, hc_mult): torch.set_default_device(DEVICE) @@ -172,8 +173,8 @@ def test_mhc_post_tilelang(num_tokens, hidden_size, hc_mult): @pytest.mark.skipif( - not (current_platform.is_cuda() or current_platform.is_rocm()), - reason="CUDA or ROCm required", + not (current_platform.is_cuda_alike() and has_tilelang()), + reason="CUDA or ROCm and tilelang required", ) @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128]) @pytest.mark.parametrize("hidden_size", [4096, 7168]) @@ -279,8 +280,8 @@ def test_hc_head_triton(num_tokens, hidden_size, hc_mult): @pytest.mark.skipif( - not (current_platform.is_cuda() or current_platform.is_rocm()), - reason="CUDA or ROCm required", + not (current_platform.is_cuda_alike() and has_tilelang()), + reason="CUDA or ROCm and tilelang required", ) @pytest.mark.parametrize("num_tokens", [1, 4, 8, 128]) @pytest.mark.parametrize("hidden_size", [4096, 7168]) From 122ad4490afea61a884ffd6c88714bc8387b2e0e Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Mon, 25 May 2026 21:21:59 -0500 Subject: [PATCH 03/27] add tilelang to requirements.txt Signed-off-by: tjtanaa --- requirements/build/rocm.txt | 1 + requirements/rocm.txt | 1 + 2 files changed, 2 insertions(+) diff --git a/requirements/build/rocm.txt b/requirements/build/rocm.txt index e5c2176a2c8c..752fb1db786a 100644 --- a/requirements/build/rocm.txt +++ b/requirements/build/rocm.txt @@ -16,3 +16,4 @@ wheel jinja2>=3.1.6 amdsmi==7.0.2 timm>=1.0.17 +tilelang>=0.1.10 diff --git a/requirements/rocm.txt b/requirements/rocm.txt index 61fcbc07010c..37bdf1afd15a 100644 --- a/requirements/rocm.txt +++ b/requirements/rocm.txt @@ -22,3 +22,4 @@ timm>=1.0.17 # amd-quark: required for Quark quantization on ROCm # To be consistent with test_quark.py amd-quark>=0.8.99 +tilelang>=0.1.10 From 63d5159e30c5716bb1c908f7677adc169147eefc Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Mon, 25 May 2026 23:02:52 -0500 Subject: [PATCH 04/27] add tilelang_hc_pre_norm_gemm Signed-off-by: tjtanaa --- tests/kernels/test_mhc_kernels.py | 48 ++++++ vllm/_tilelang_ops.py | 181 ++++++++++++++++++++ vllm/model_executor/kernels/mhc/tilelang.py | 87 +++++++++- 3 files changed, 313 insertions(+), 3 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 597f8e1ff7cd..985cb6830105 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -4,6 +4,10 @@ import torch import vllm.model_executor.kernels.mhc # noqa: F401 +from vllm.model_executor.kernels.mhc.tilelang import ( + _tilelang_hc_prenorm_gemm, + _torch_hc_prenorm_gemm, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed from vllm.utils.import_utils import has_tilelang @@ -145,6 +149,50 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult): torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2) +@pytest.mark.skipif( + not (current_platform.is_cuda_alike() and has_tilelang()), + reason="CUDA or ROCm and tilelang required", +) +@pytest.mark.parametrize( + ("num_tokens", "hidden_size"), + [ + (1, 1280), + (512, 1280), + (2048, 1280), + (1, 4096), + (64, 4096), + (512, 4096), + (2048, 4096), + (1, 7168), + (64, 7168), + (512, 7168), + (2048, 7168), + ], +) +def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size): + torch.set_default_device(DEVICE) + set_random_seed(0) + + hc_mult = 4 + hc_mult3 = 2 * hc_mult + hc_mult * hc_mult + x = torch.randn( + (num_tokens, hc_mult * hidden_size), dtype=torch.bfloat16 + ) + fn = torch.randn( + (hc_mult3, hc_mult * hidden_size), dtype=torch.float32 + ) * 1e-4 + out_ref = torch.empty((1, num_tokens, hc_mult3), dtype=torch.float32) + sqrsum_ref = torch.empty((1, num_tokens), dtype=torch.float32) + out = torch.empty_like(out_ref) + sqrsum = torch.empty_like(sqrsum_ref) + + _torch_hc_prenorm_gemm(x, fn, out_ref, sqrsum_ref) + _tilelang_hc_prenorm_gemm(x, fn, out, sqrsum, hidden_size, hc_mult) + + torch.testing.assert_close(out, out_ref, atol=1e-5, rtol=1e-4) + torch.testing.assert_close(sqrsum, sqrsum_ref, atol=8.0, rtol=5e-4) + + @pytest.mark.skipif( not (current_platform.is_cuda_alike() and has_tilelang()), reason="CUDA or ROCm and tilelang required", diff --git a/vllm/_tilelang_ops.py b/vllm/_tilelang_ops.py index 9f2827cceee3..d30193f08a67 100644 --- a/vllm/_tilelang_ops.py +++ b/vllm/_tilelang_ops.py @@ -541,6 +541,187 @@ def mhc_post_tilelang( T.pdl_trigger() +@tilelang.jit( + pass_configs=pass_configs, +) +def hc_prenorm_gemm_tilelang( + x, + fn, + out, + sqrsum, + hidden_size: int, + hc_mult: int = 4, + n_out: int = 24, + n_thr: int = 512, + tile_n: int = 12, + n_splits: int = 1, +) -> tilelang.JITKernel: + num_tokens = T.dynamic("num_tokens") + hc_hidden_size = hc_mult * hidden_size + k_per_split = hc_hidden_size // n_splits + k_iters = k_per_split // n_thr + n_tiles = T.ceildiv(n_out, tile_n) + + x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16) # type: ignore[no-redef, valid-type] + fn: T.Tensor((n_out, hc_hidden_size), T.float32) # type: ignore[no-redef, valid-type] + out: T.Tensor((n_splits, num_tokens, n_out), T.float32) # type: ignore[no-redef, valid-type] + sqrsum: T.Tensor((n_splits, num_tokens), T.float32) # type: ignore[no-redef, valid-type] + + with T.Kernel(num_tokens, n_tiles, n_splits, threads=n_thr) as ( + i_n, + i_t, + i_s, + ): + tid = T.get_thread_binding() + acc = T.alloc_local((tile_n,), T.float32) + sqr = T.alloc_local((1,), T.float32) + T.clear(acc) + T.clear(sqr) + + if ENABLE_PDL: + T.pdl_sync() + + for it in T.serial(k_iters): + i_k = i_s * k_per_split + it * n_thr + tid + x_val = x[i_n, i_k] + for i_o in T.unroll(tile_n): + out_idx = i_t * tile_n + i_o + if out_idx < n_out: + acc[i_o] += x_val * fn[out_idx, i_k] + if i_t == 0: + sqr[0] += x_val * x_val + + for i_o in T.unroll(tile_n): + acc[i_o] = T.warp_reduce_sum(acc[i_o]) + if i_t == 0: + sqr[0] = T.warp_reduce_sum(sqr[0]) + + lane = tid % 32 + warp_id = tid // 32 + num_warps = n_thr // 32 + warp_acc = T.alloc_shared((num_warps, tile_n), T.float32) + warp_sqr = T.alloc_shared(num_warps, T.float32) + + if lane == 0: + for i_o in T.unroll(tile_n): + warp_acc[warp_id, i_o] = acc[i_o] + if i_t == 0: + warp_sqr[warp_id] = sqr[0] + T.sync_threads() + + if warp_id == 0: + if lane < tile_n: + reduced_acc = T.alloc_var(T.float32, init=0.0) + for i_w in T.unroll(num_warps): + reduced_acc += warp_acc[i_w, lane] + out_idx = i_t * tile_n + lane + if out_idx < n_out: + out[i_s, i_n, out_idx] = reduced_acc + if lane == 0 and i_t == 0: + reduced_sqr = T.alloc_var(T.float32, init=0.0) + for i_w in T.unroll(num_warps): + reduced_sqr += warp_sqr[i_w] + sqrsum[i_s, i_n] = reduced_sqr + + if ENABLE_PDL: + T.pdl_trigger() + + +@tilelang.jit( + pass_configs=pass_configs, +) +def hc_prenorm_gemm_block_m_tilelang( + x, + fn, + out, + sqrsum, + hidden_size: int, + hc_mult: int = 4, + n_out: int = 24, + n_thr: int = 512, + tile_n: int = 12, + block_m: int = 2, +) -> tilelang.JITKernel: + num_tokens = T.dynamic("num_tokens") + hc_hidden_size = hc_mult * hidden_size + k_iters = hc_hidden_size // n_thr + n_tiles = T.ceildiv(n_out, tile_n) + m_tiles = T.ceildiv(num_tokens, block_m) + + x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16) # type: ignore[no-redef, valid-type] + fn: T.Tensor((n_out, hc_hidden_size), T.float32) # type: ignore[no-redef, valid-type] + out: T.Tensor((1, num_tokens, n_out), T.float32) # type: ignore[no-redef, valid-type] + sqrsum: T.Tensor((1, num_tokens), T.float32) # type: ignore[no-redef, valid-type] + + with T.Kernel(m_tiles, n_tiles, threads=n_thr) as (i_mt, i_t): + tid = T.get_thread_binding() + acc = T.alloc_local((block_m, tile_n), T.float32) + sqr = T.alloc_local((block_m,), T.float32) + T.clear(acc) + T.clear(sqr) + + if ENABLE_PDL: + T.pdl_sync() + + for it in T.serial(k_iters): + i_k = it * n_thr + tid + fn_val = T.alloc_local((tile_n,), T.float32) + for i_o in T.unroll(tile_n): + out_idx = i_t * tile_n + i_o + if out_idx < n_out: + fn_val[i_o] = fn[out_idx, i_k] + else: + fn_val[i_o] = 0.0 + for i_m in T.unroll(block_m): + token_idx = i_mt * block_m + i_m + if token_idx < num_tokens: + x_val = x[token_idx, i_k] + for i_o in T.unroll(tile_n): + acc[i_m, i_o] += x_val * fn_val[i_o] + if i_t == 0: + sqr[i_m] += x_val * x_val + + for i_m in T.unroll(block_m): + for i_o in T.unroll(tile_n): + acc[i_m, i_o] = T.warp_reduce_sum(acc[i_m, i_o]) + if i_t == 0: + sqr[i_m] = T.warp_reduce_sum(sqr[i_m]) + + lane = tid % 32 + warp_id = tid // 32 + num_warps = n_thr // 32 + warp_acc = T.alloc_shared((num_warps, block_m, tile_n), T.float32) + warp_sqr = T.alloc_shared((num_warps, block_m), T.float32) + + if lane == 0: + for i_m in T.unroll(block_m): + for i_o in T.unroll(tile_n): + warp_acc[warp_id, i_m, i_o] = acc[i_m, i_o] + if i_t == 0: + warp_sqr[warp_id, i_m] = sqr[i_m] + T.sync_threads() + + if warp_id == 0: + for i_m in T.unroll(block_m): + token_idx = i_mt * block_m + i_m + if token_idx < num_tokens: + if lane < tile_n: + reduced_acc = T.alloc_var(T.float32, init=0.0) + for i_w in T.unroll(num_warps): + reduced_acc += warp_acc[i_w, i_m, lane] + out_idx = i_t * tile_n + lane + if out_idx < n_out: + out[0, token_idx, out_idx] = reduced_acc + if lane == 0 and i_t == 0: + reduced_sqr = T.alloc_var(T.float32, init=0.0) + for i_w in T.unroll(num_warps): + reduced_sqr += warp_sqr[i_w, i_m] + sqrsum[0, token_idx] = reduced_sqr + + if ENABLE_PDL: + T.pdl_trigger() + + @tilelang.jit( pass_configs=pass_configs, ) diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index 3000cd402dd1..16ef72fcf19c 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -24,6 +24,75 @@ def _torch_hc_prenorm_gemm( sqrsum[0].copy_(x_float.square().sum(dim=-1)) +def _tilelang_hc_prenorm_gemm( + x: torch.Tensor, + fn: torch.Tensor, + out: torch.Tensor, + sqrsum: torch.Tensor, + hidden_size: int, + hc_mult: int, + tile_n: int = 12, + n_thr: int = 512, + n_splits: int = 1, +) -> None: + from vllm._tilelang_ops import ( + hc_prenorm_gemm_block_m_tilelang, + hc_prenorm_gemm_tilelang, + ) + + assert out.shape[0] == n_splits + assert sqrsum.shape[0] == n_splits + assert x.shape[1] == hc_mult * hidden_size + assert x.shape[1] % n_splits == 0 + assert (x.shape[1] // n_splits) % n_thr == 0 + use_default_config = tile_n == 12 and n_thr == 512 + if n_splits == 1 and use_default_config and x.shape[0] >= 1024: + hc_prenorm_gemm_block_m_tilelang( + x, + fn, + out, + sqrsum, + hidden_size, + hc_mult, + fn.shape[0], + n_thr, + tile_n, + 2, + ) + return + if ( + n_splits == 1 + and use_default_config + and x.shape[0] < 128 + and x.shape[1] % 1024 == 0 + ): + hc_prenorm_gemm_tilelang( + x, + fn, + out, + sqrsum, + hidden_size, + hc_mult, + fn.shape[0], + 1024, + 4, + n_splits, + ) + return + hc_prenorm_gemm_tilelang( + x, + fn, + out, + sqrsum, + hidden_size, + hc_mult, + fn.shape[0], + n_thr, + tile_n, + n_splits, + ) + + def mhc_pre_tilelang( residual: torch.Tensor, fn: torch.Tensor, @@ -137,7 +206,14 @@ def mhc_pre_tilelang( n_splits, ) else: - _torch_hc_prenorm_gemm(residual_2d, fn, gemm_out_mul, gemm_out_sqrsum) + _tilelang_hc_prenorm_gemm( + residual_2d, + fn, + gemm_out_mul, + gemm_out_sqrsum, + hidden_size, + hc_mult, + ) if norm_weight is None: mhc_pre_big_fuse_tilelang( @@ -422,8 +498,13 @@ def mhc_fused_post_pre_tilelang( n_splits, ) else: - _torch_hc_prenorm_gemm( - residual_cur_2d, fn, gemm_out_mul, gemm_out_sqrsum + _tilelang_hc_prenorm_gemm( + residual_cur_2d, + fn, + gemm_out_mul, + gemm_out_sqrsum, + hidden_size, + hc_mult, ) if norm_weight is None: From c966a7dab19b33afd4c933d3f948b1b80fa814ab Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Tue, 26 May 2026 01:01:06 -0500 Subject: [PATCH 05/27] fix prefix commit; support fuse and unfused codepath Signed-off-by: tjtanaa --- tests/kernels/test_mhc_kernels.py | 10 +-- vllm/_tilelang_ops.py | 6 +- vllm/model_executor/kernels/mhc/tilelang.py | 4 +- vllm/model_executor/layers/mhc.py | 93 +++++++++++++++------ vllm/models/deepseek_v4/amd/model.py | 23 ++--- vllm/platforms/cuda.py | 1 + vllm/platforms/interface.py | 1 + 7 files changed, 91 insertions(+), 47 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 985cb6830105..e7d4cde43f1d 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -9,8 +9,8 @@ _torch_hc_prenorm_gemm, ) from vllm.platforms import current_platform -from vllm.utils.torch_utils import set_random_seed from vllm.utils.import_utils import has_tilelang +from vllm.utils.torch_utils import set_random_seed DEVICE = current_platform.device_type @@ -175,12 +175,8 @@ def test_hc_prenorm_gemm_tilelang(num_tokens, hidden_size): hc_mult = 4 hc_mult3 = 2 * hc_mult + hc_mult * hc_mult - x = torch.randn( - (num_tokens, hc_mult * hidden_size), dtype=torch.bfloat16 - ) - fn = torch.randn( - (hc_mult3, hc_mult * hidden_size), dtype=torch.float32 - ) * 1e-4 + x = torch.randn((num_tokens, hc_mult * hidden_size), dtype=torch.bfloat16) + fn = torch.randn((hc_mult3, hc_mult * hidden_size), dtype=torch.float32) * 1e-4 out_ref = torch.empty((1, num_tokens, hc_mult3), dtype=torch.float32) sqrsum_ref = torch.empty((1, num_tokens), dtype=torch.float32) out = torch.empty_like(out_ref) diff --git a/vllm/_tilelang_ops.py b/vllm/_tilelang_ops.py index d30193f08a67..3b34f7f098a7 100644 --- a/vllm/_tilelang_ops.py +++ b/vllm/_tilelang_ops.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math from functools import cache -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch @@ -10,6 +10,7 @@ from vllm.utils.import_utils import has_tilelang from vllm.utils.math_utils import cdiv + def _is_tilelang_platform() -> bool: return current_platform.is_cuda() or current_platform.is_rocm() @@ -37,6 +38,7 @@ def _is_pdl_supported() -> bool: ENABLE_PDL = _is_pdl_supported() + @cache def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: device_props = torch.cuda.get_device_properties(0) @@ -50,7 +52,7 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: return split_k -pass_configs = { +pass_configs: dict[tilelang.PassConfigKey, Any] = { tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, } diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index 16ef72fcf19c..61f5c6aa2fa7 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -173,9 +173,7 @@ def mhc_pre_tilelang( # these numbers are from deepgemm kernel impl block_k = 64 block_m = 64 - n_splits = compute_num_split( - block_k, hc_hidden_size, cdiv(num_tokens, block_m) - ) + n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m)) else: n_splits = 1 diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index ecf3b7a3c024..aadefac380f5 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -86,17 +86,6 @@ def forward_hip( # sinkhorn_repeat, # ) # else: - # return mhc_kernels.mhc_pre_torch( - # residual, - # fn, - # hc_scale, - # hc_base, - # rms_eps, - # hc_pre_eps, - # hc_sinkhorn_eps, - # hc_post_mult_value, - # sinkhorn_repeat, - # ) return torch.ops.vllm.mhc_pre_tilelang( residual, fn, @@ -112,8 +101,32 @@ def forward_hip( norm_eps, ) - def forward_native(self, *args, **kwargs): - raise NotImplementedError("Native implementation of mhc_pre is not available") + def forward_native( + self, + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + n_splits: int = 1, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 0.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return mhc_kernels.mhc_pre_torch( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + ) # --8<-- [start:mhc_post] @@ -162,18 +175,23 @@ def forward_hip( # comb_res_mix, # ) # else: - # return mhc_kernels.mhc_post_torch( - # x, - # residual, - # post_layer_mix, - # comb_res_mix, - # ) return torch.ops.vllm.mhc_post_tilelang( x, residual, post_layer_mix, comb_res_mix ) - def forward_native(self, *args, **kwargs): - raise NotImplementedError("Native implementation of mhc_post is not available") + def forward_native( + self, + x: torch.Tensor, + residual: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + ) -> torch.Tensor: + return mhc_kernels.mhc_post_torch( + x, + residual, + post_layer_mix, + comb_res_mix, + ) # --8<-- [start:hc_head] @@ -251,6 +269,36 @@ def forward_hip( ) return out.view(*outer_shape, hidden_size) + def _forward_triton( + self, + hidden_states: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_norm_eps: float, + hc_eps: float, + ) -> torch.Tensor: + hc_mult, hidden_size = hidden_states.shape[-2:] + outer_shape = hidden_states.shape[:-2] + hs_flat = hidden_states.view(-1, hc_mult, hidden_size) + num_tokens = hs_flat.shape[0] + + out = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device + ) + torch.ops.vllm.hc_head_triton( + hs_flat, + hc_fn, + hc_scale, + hc_base, + out, + hidden_size, + rms_norm_eps, + hc_eps, + hc_mult, + ) + return out.view(*outer_shape, hidden_size) + def forward_native(self, *args, **kwargs): raise NotImplementedError("Native implementation of hc_head is not available") @@ -345,9 +393,6 @@ def forward_hip( norm_weight, norm_eps, ) - # raise NotImplementedError( - # "Hip implementation of mhc_fused_post_pre is not available" - # ) def forward_native(self, *args, **kwargs): raise NotImplementedError( diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 1680d10a192e..035aee4bd3ab 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -61,6 +61,7 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.triton_utils import tl, triton +from vllm.utils.import_utils import has_tilelang from vllm.utils.torch_utils import direct_register_custom_op @@ -1074,6 +1075,7 @@ def __init__( self.mhc_pre = MHCPreOp() self.mhc_post = MHCPostOp() self.mhc_fused_post_pre = MHCFusedPostPreOp() + self.has_tilelang = has_tilelang() def hc_pre( self, @@ -1104,7 +1106,7 @@ def hc_post( ): return self.mhc_post(x, residual, post, comb) - def _forward_cuda( + def _forward_fused_post_pre( self, x: torch.Tensor, positions: torch.Tensor, @@ -1156,7 +1158,7 @@ def _forward_cuda( x = self.ffn(x, input_ids) return x, residual, post_mix, res_mix - def _forward_rocm( + def _forward_unfused_post_pre( self, x: torch.Tensor, positions: torch.Tensor, @@ -1195,12 +1197,13 @@ def forward( ) -> tuple[ torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None ]: - # if current_platform.is_rocm(): - # return self._forward_rocm( - # x, positions, input_ids, post_mix, res_mix, residual - # ) - - return self._forward_cuda(x, positions, input_ids, post_mix, res_mix, residual) + if not self.has_tilelang: + return self._forward_unfused_post_pre( + x, positions, input_ids, post_mix, res_mix, residual + ) + return self._forward_fused_post_pre( + x, positions, input_ids, post_mix, res_mix, residual + ) @support_torch_compile @@ -1361,9 +1364,7 @@ def forward( res_mix, residual, ) - if layer is not None and ( - current_platform.is_cuda() or current_platform.is_rocm() - ): + if layer is not None and self.has_tilelang: hidden_states = layer.hc_post(hidden_states, residual, post_mix, res_mix) if not get_pp_group().is_last_rank: diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index b5cfce40b1b9..4b6f52b12fea 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -601,6 +601,7 @@ def is_arch_support_pdl(cls) -> bool: return False return major >= 9 + # NVML utils # Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`, # all the related functions work on real physical device ids. diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 8c1046881113..cf774b7bda91 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -1023,6 +1023,7 @@ def is_arch_support_pdl(cls) -> bool: """ return False + class UnspecifiedPlatform(Platform): _enum = PlatformEnum.UNSPECIFIED device_type = "" From d2f3c6dadde6cbeec63ceb9242742e6b5f3e4165 Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Tue, 26 May 2026 02:32:57 -0500 Subject: [PATCH 06/27] use tilelang mhc in mtp as well Signed-off-by: tjtanaa --- vllm/model_executor/layers/mhc.py | 119 ++++++++++++++------------- vllm/models/deepseek_v4/amd/model.py | 1 + vllm/models/deepseek_v4/amd/mtp.py | 4 +- vllm/utils/import_utils.py | 1 + 4 files changed, 68 insertions(+), 57 deletions(-) diff --git a/vllm/model_executor/layers/mhc.py b/vllm/model_executor/layers/mhc.py index aadefac380f5..b720fa1f6fe2 100644 --- a/vllm/model_executor/layers/mhc.py +++ b/vllm/model_executor/layers/mhc.py @@ -6,6 +6,9 @@ # import vllm.model_executor.kernels.mhc # noqa: F401 import vllm.model_executor.kernels.mhc as mhc_kernels from vllm.model_executor.custom_op import CustomOp +from vllm.utils.import_utils import has_tilelang + +HAS_TILELANG = has_tilelang() # --8<-- [start:mhc_pre] @@ -86,20 +89,36 @@ def forward_hip( # sinkhorn_repeat, # ) # else: - return torch.ops.vllm.mhc_pre_tilelang( - residual, - fn, - hc_scale, - hc_base, - rms_eps, - hc_pre_eps, - hc_sinkhorn_eps, - hc_post_mult_value, - sinkhorn_repeat, - n_splits, - norm_weight, - norm_eps, - ) + if HAS_TILELANG: + return torch.ops.vllm.mhc_pre_tilelang( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + norm_weight, + norm_eps, + ) + else: + return self.forward_native( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + n_splits, + norm_weight, + norm_eps, + ) def forward_native( self, @@ -175,9 +194,12 @@ def forward_hip( # comb_res_mix, # ) # else: - return torch.ops.vllm.mhc_post_tilelang( - x, residual, post_layer_mix, comb_res_mix - ) + if HAS_TILELANG: + return torch.ops.vllm.mhc_post_tilelang( + x, residual, post_layer_mix, comb_res_mix + ) + else: + return self.forward_native(x, residual, post_layer_mix, comb_res_mix) def forward_native( self, @@ -256,47 +278,32 @@ def forward_hip( out = torch.empty( num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device ) - torch.ops.vllm.hc_head_fused_kernel_tilelang( - hs_flat, - hc_fn, - hc_scale, - hc_base, - out, - hidden_size, - rms_norm_eps, - hc_eps, - hc_mult, - ) - return out.view(*outer_shape, hidden_size) - def _forward_triton( - self, - hidden_states: torch.Tensor, - hc_fn: torch.Tensor, - hc_scale: torch.Tensor, - hc_base: torch.Tensor, - rms_norm_eps: float, - hc_eps: float, - ) -> torch.Tensor: - hc_mult, hidden_size = hidden_states.shape[-2:] - outer_shape = hidden_states.shape[:-2] - hs_flat = hidden_states.view(-1, hc_mult, hidden_size) - num_tokens = hs_flat.shape[0] + if HAS_TILELANG: + torch.ops.vllm.hc_head_fused_kernel_tilelang( + hs_flat, + hc_fn, + hc_scale, + hc_base, + out, + hidden_size, + rms_norm_eps, + hc_eps, + hc_mult, + ) + else: + torch.ops.vllm.hc_head_triton( + hs_flat, + hc_fn, + hc_scale, + hc_base, + out, + hidden_size, + rms_norm_eps, + hc_eps, + hc_mult, + ) - out = torch.empty( - num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device - ) - torch.ops.vllm.hc_head_triton( - hs_flat, - hc_fn, - hc_scale, - hc_base, - out, - hidden_size, - rms_norm_eps, - hc_eps, - hc_mult, - ) return out.view(*outer_shape, hidden_size) def forward_native(self, *args, **kwargs): diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 035aee4bd3ab..b1d8618f98bc 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -1295,6 +1295,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): requires_grad=False, ) self.hc_head_op = HCHeadOp() + self.has_tilelang = has_tilelang() # Pre-hc_head residual stream buffer for the MTP draft. Stable # address (outside the cudagraph pool) so the copy_ in forward() # refreshes it correctly across captured shapes. diff --git a/vllm/models/deepseek_v4/amd/mtp.py b/vllm/models/deepseek_v4/amd/mtp.py index bcdd76de4c29..7f7501d7a4f0 100644 --- a/vllm/models/deepseek_v4/amd/mtp.py +++ b/vllm/models/deepseek_v4/amd/mtp.py @@ -39,6 +39,7 @@ from vllm.model_executor.models.utils import maybe_prefix from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors +from vllm.utils.import_utils import has_tilelang from .model import ( DeepseekV4DecoderLayer, @@ -121,6 +122,7 @@ def __init__( ) self.hc_head_op = HCHeadOp() + self.has_tilelang = has_tilelang() def forward( self, @@ -147,7 +149,7 @@ def forward( hidden_states, residual, post_mix, res_mix = self.mtp_block( positions=positions, x=hidden_states, input_ids=None ) - if current_platform.is_cuda(): + if self.has_tilelang: hidden_states = self.mtp_block.hc_post( hidden_states, residual, post_mix, res_mix ) diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index 5822e5840afc..e97228bfa609 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -430,6 +430,7 @@ def has_triton_kernels() -> bool: return is_available +@cache def has_tilelang() -> bool: """Whether the optional `tilelang` package is available.""" return _has_module("tilelang") From 9111eeee93ed21619f18b8ff9cb29920599dc420 Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Tue, 26 May 2026 08:28:42 -0500 Subject: [PATCH 07/27] clean up code Signed-off-by: tjtanaa --- vllm/model_executor/kernels/mhc/tilelang.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index 61f5c6aa2fa7..d76123bb7625 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -5,12 +5,6 @@ from vllm.utils.torch_utils import direct_register_custom_op -def _can_use_deep_gemm_hc_prenorm() -> bool: - from vllm.utils.deep_gemm import is_deep_gemm_supported - - return is_deep_gemm_supported() - - def _torch_hc_prenorm_gemm( x: torch.Tensor, fn: torch.Tensor, @@ -168,7 +162,9 @@ def mhc_pre_tilelang( residual_flat = residual.view(-1, hc_mult, hidden_size) num_tokens = residual_flat.shape[0] - use_deep_gemm = _can_use_deep_gemm_hc_prenorm() + from vllm.utils.deep_gemm import is_deep_gemm_supported + + use_deep_gemm = is_deep_gemm_supported() if use_deep_gemm: # these numbers are from deepgemm kernel impl block_k = 64 @@ -407,7 +403,9 @@ def mhc_fused_post_pre_tilelang( post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult) comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult) - use_deep_gemm = _can_use_deep_gemm_hc_prenorm() + from vllm.utils.deep_gemm import is_deep_gemm_supported + + use_deep_gemm = is_deep_gemm_supported() use_small_fma = num_tokens <= 16 if use_small_fma: # TODO(gnovack): investigate autotuning these heuristics From fbb43d77fe2e05e1cae8a84f616252662b0a02ca Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Tue, 26 May 2026 09:14:26 -0500 Subject: [PATCH 08/27] clean up code Signed-off-by: tjtanaa --- vllm/_tilelang_ops.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/vllm/_tilelang_ops.py b/vllm/_tilelang_ops.py index 3b34f7f098a7..8808a1ef0abf 100644 --- a/vllm/_tilelang_ops.py +++ b/vllm/_tilelang_ops.py @@ -11,10 +11,6 @@ from vllm.utils.math_utils import cdiv -def _is_tilelang_platform() -> bool: - return current_platform.is_cuda() or current_platform.is_rocm() - - def _is_pdl_supported() -> bool: is_arch_support_pdl = getattr(current_platform, "is_arch_support_pdl", None) if not callable(is_arch_support_pdl): @@ -24,7 +20,7 @@ def _is_pdl_supported() -> bool: # TileLang is used for MHC on CUDA and ROCm. Keep non-GPU imports cheap so # registering the Python wrapper modules does not require TileLang everywhere. -if TYPE_CHECKING or _is_tilelang_platform(): +if TYPE_CHECKING or current_platform.is_cuda_alike(): if not has_tilelang(): raise ImportError( "tilelang is required for mhc but is not installed. Install it with " From ed3db3288417fa051abf113893590ed2be1d8178 Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Tue, 26 May 2026 09:49:22 -0500 Subject: [PATCH 09/27] clean up code Signed-off-by: tjtanaa --- vllm/_tilelang_ops.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/vllm/_tilelang_ops.py b/vllm/_tilelang_ops.py index 8808a1ef0abf..5cc91a470a31 100644 --- a/vllm/_tilelang_ops.py +++ b/vllm/_tilelang_ops.py @@ -10,14 +10,6 @@ from vllm.utils.import_utils import has_tilelang from vllm.utils.math_utils import cdiv - -def _is_pdl_supported() -> bool: - is_arch_support_pdl = getattr(current_platform, "is_arch_support_pdl", None) - if not callable(is_arch_support_pdl): - return False - return is_arch_support_pdl() - - # TileLang is used for MHC on CUDA and ROCm. Keep non-GPU imports cheap so # registering the Python wrapper modules does not require TileLang everywhere. if TYPE_CHECKING or current_platform.is_cuda_alike(): @@ -32,7 +24,7 @@ def _is_pdl_supported() -> bool: tilelang = None # type: ignore[assignment] T = None # type: ignore[assignment] -ENABLE_PDL = _is_pdl_supported() +ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda() @cache From 52a31ccecca203effd490a5b496dc5f8d9496654 Mon Sep 17 00:00:00 2001 From: Ashwin Giridharan Date: Wed, 27 May 2026 05:39:49 -0700 Subject: [PATCH 10/27] [Bugfix] Map reasoning_effort to enable_thinking in chat template kwargs (#43401) Signed-off-by: Ashwin Giridharan Signed-off-by: Chauncey Co-authored-by: Chauncey --- docs/features/reasoning_outputs.md | 39 ++++++- .../openai/test_reasoning_enable_thinking.py | 107 ++++++++++++++++++ .../openai/chat_completion/protocol.py | 22 +++- vllm/entrypoints/openai/responses/protocol.py | 23 +++- 4 files changed, 178 insertions(+), 13 deletions(-) create mode 100644 tests/entrypoints/openai/test_reasoning_enable_thinking.py diff --git a/docs/features/reasoning_outputs.md b/docs/features/reasoning_outputs.md index f1cc18a25cbb..92563a8b4bbd 100644 --- a/docs/features/reasoning_outputs.md +++ b/docs/features/reasoning_outputs.md @@ -15,6 +15,7 @@ vLLM currently supports the following reasoning models: | ------------ | ----------- | ---------------- | ----------- | | [Cohere Command A Reasoning](https://huggingface.co/CohereLabs/command-a-reasoning-08-2025) | `cohere_command3` | `json`, `regex` | ✅ | | [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d) | `deepseek_r1` | `json`, `regex` | ❌ | +| [Gemma 4 series](https://huggingface.co/google/gemma-4-26B-A4B-it) | `gemma4` | `json`, `regex` | ✅ | | [DeepSeek-V3.1](https://huggingface.co/collections/deepseek-ai/deepseek-v31-68a491bed32bd77e7fca048f) | `deepseek_v3` | `json`, `regex` | ❌ | | [ERNIE-4.5-VL series](https://huggingface.co/baidu/ERNIE-4.5-VL-28B-A3B-PT) | `ernie45` | `json`, `regex` | ❌ | | [ERNIE-4.5-21B-A3B-Thinking](https://huggingface.co/baidu/ERNIE-4.5-21B-A3B-Thinking) | `ernie45` | `json`, `regex` | ✅ | @@ -29,6 +30,7 @@ vLLM currently supports the following reasoning models: !!! note IBM Granite 3.2 and DeepSeek-V3.1 reasoning is disabled by default; to enable it, you must also pass `thinking=True` in your `chat_template_kwargs`. The reasoning feature for the Qwen3 series is enabled by default. To disable it, you must pass `enable_thinking=False` in your `chat_template_kwargs`. + Gemma 4 reasoning is disabled by default; to enable it, pass `enable_thinking=True` in your `chat_template_kwargs` or set `reasoning_effort` (which enables it automatically). DeepSeek-V3.1 tool calling is supported in non-thinking mode. Holo2 reasoning is enabled by default. To disable it, you must also pass `thinking=False` in your `chat_template_kwargs`. @@ -314,9 +316,44 @@ for output in outputs: print("text:", output.outputs[0].text) ``` +## Automatic `enable_thinking` Activation + +Some models (such as Gemma 4, DeepSeek-V4-Pro and IBM Granite 3.2) require `enable_thinking: true` in their chat template kwargs to activate thinking mode — without it, reasoning tokens are never generated regardless of other settings. + +When you set `reasoning_effort` in a Chat Completions request (or `reasoning.effort` in a Responses API request), vLLM automatically injects `enable_thinking` into the chat template kwargs: + +- `reasoning_effort` = `"low"`, `"medium"`, or `"high"` → `enable_thinking = true` +- `reasoning_effort` = `"none"` → `enable_thinking = false` +- `reasoning_effort` not set → `enable_thinking` is not injected (preserves existing behavior) + +This means you no longer need to manually pass `chat_template_kwargs: {"enable_thinking": true}` when using `reasoning_effort` — it is handled automatically. + +!!! note + If you explicitly set `enable_thinking` in `chat_template_kwargs`, your value takes priority over the automatic injection. This allows you to override the behavior if needed. + + For models whose templates don't declare `enable_thinking` (e.g., DeepSeek R1), the injected kwarg is harmlessly filtered out by `resolve_chat_template_kwargs`. + +### Example + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy") + +# reasoning_effort automatically enables thinking for models that need it +response = client.chat.completions.create( + model="google/gemma-4-26B-A4B-it", + messages=[{"role": "user", "content": "What is 15 * 37?"}], + reasoning_effort="high", # Automatically sets enable_thinking=true +) + +print(response.choices[0].message.reasoning) +print(response.choices[0].message.content) +``` + ## Limitations -- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`). +- The reasoning content is only available for online serving's chat completion endpoint (`/v1/chat/completions`), Anthropic Messages API (`/v1/messages`) and the Responses API (`/v1/responses`). ## How to support a new reasoning model diff --git a/tests/entrypoints/openai/test_reasoning_enable_thinking.py b/tests/entrypoints/openai/test_reasoning_enable_thinking.py new file mode 100644 index 000000000000..95606f626a0b --- /dev/null +++ b/tests/entrypoints/openai/test_reasoning_enable_thinking.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for reasoning_effort -> enable_thinking mapping. + +Models like Gemma4 require enable_thinking=True in chat_template_kwargs to +activate thinking mode. This mapping ensures that when a user requests +reasoning (via reasoning_effort or reasoning.effort), the template kwarg +is injected automatically. +""" + +import pytest +from openai.types.shared import Reasoning + +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) +from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + + +def _build_chat_request(**kwargs) -> ChatCompletionRequest: + defaults = dict( + model="test-model", + messages=[{"role": "user", "content": "Hello"}], + ) + defaults.update(kwargs) + return ChatCompletionRequest(**defaults) + + +def _build_responses_request(**kwargs) -> ResponsesRequest: + defaults = dict( + model="test-model", + input=[{"role": "user", "content": "Hello"}], + ) + defaults.update(kwargs) + return ResponsesRequest(**defaults) + + +class TestChatCompletionReasoningEffort: + """Chat Completions: reasoning_effort -> enable_thinking.""" + + @pytest.mark.parametrize("effort", ["low", "medium", "high"]) + def test_non_none_effort_injects_enable_thinking_true(self, effort): + request = _build_chat_request(reasoning_effort=effort) + params = request.build_chat_params(None, "auto") + assert params.chat_template_kwargs["enable_thinking"] is True + + def test_none_effort_injects_enable_thinking_false(self): + request = _build_chat_request(reasoning_effort="none") + params = request.build_chat_params(None, "auto") + assert params.chat_template_kwargs["enable_thinking"] is False + + def test_no_effort_does_not_inject(self): + request = _build_chat_request() + params = request.build_chat_params(None, "auto") + assert "enable_thinking" not in params.chat_template_kwargs + + def test_explicit_user_kwarg_not_overridden(self): + request = _build_chat_request( + reasoning_effort="high", + chat_template_kwargs={"enable_thinking": False}, + ) + params = request.build_chat_params(None, "auto") + assert params.chat_template_kwargs["enable_thinking"] is False + + def test_reasoning_effort_still_in_kwargs(self): + request = _build_chat_request(reasoning_effort="high") + params = request.build_chat_params(None, "auto") + assert params.chat_template_kwargs["reasoning_effort"] == "high" + + +class TestResponsesReasoningEffort: + """Responses API: reasoning.effort -> enable_thinking.""" + + @pytest.mark.parametrize("effort", ["low", "medium", "high"]) + def test_non_none_effort_injects_enable_thinking_true(self, effort): + request = _build_responses_request( + reasoning=Reasoning(effort=effort), + ) + params = request.build_chat_params(None, "auto") + assert params.chat_template_kwargs["enable_thinking"] is True + + def test_none_effort_injects_enable_thinking_false(self): + request = _build_responses_request( + reasoning=Reasoning(effort="none"), + ) + params = request.build_chat_params(None, "auto") + assert params.chat_template_kwargs["enable_thinking"] is False + + def test_no_reasoning_does_not_inject(self): + request = _build_responses_request() + params = request.build_chat_params(None, "auto") + assert "enable_thinking" not in params.chat_template_kwargs + + def test_explicit_user_kwarg_not_overridden(self): + request = _build_responses_request( + reasoning=Reasoning(effort="high"), + chat_template_kwargs={"enable_thinking": False}, + ) + params = request.build_chat_params(None, "auto") + assert params.chat_template_kwargs["enable_thinking"] is False + + def test_reasoning_effort_still_in_kwargs(self): + request = _build_responses_request( + reasoning=Reasoning(effort="high"), + ) + params = request.build_chat_params(None, "auto") + assert params.chat_template_kwargs["reasoning_effort"] == "high" diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index cada289e109f..73ecb3f35a1c 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -473,17 +473,27 @@ def build_chat_params( default_template: str | None, default_template_content_format: ChatTemplateContentFormatOption, ) -> ChatParams: + extra_kwargs: dict[str, Any] = dict( + add_generation_prompt=self.add_generation_prompt, + continue_final_message=self.continue_final_message, + documents=self.documents, + reasoning_effort=self.reasoning_effort, + ) + + # When reasoning is requested, activate thinking for models whose + # chat templates require explicit opt-in (e.g., Gemma4 defaults + # enable_thinking to false). For templates that don't declare the + # variable, resolve_chat_template_kwargs filters it out harmlessly. + user_kwargs = self.chat_template_kwargs or {} + if self.reasoning_effort is not None and "enable_thinking" not in user_kwargs: + extra_kwargs["enable_thinking"] = self.reasoning_effort != "none" + return ChatParams( chat_template=self.chat_template or default_template, chat_template_content_format=default_template_content_format, chat_template_kwargs=merge_kwargs( self.chat_template_kwargs, - dict( - add_generation_prompt=self.add_generation_prompt, - continue_final_message=self.continue_final_message, - documents=self.documents, - reasoning_effort=self.reasoning_effort, - ), + extra_kwargs, ), media_io_kwargs=self.media_io_kwargs, ) diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 10aa5bde392b..370ed9a825d4 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -298,17 +298,28 @@ def build_chat_params( continue_final = should_continue_final_message(self.input) reasoning = self.reasoning + reasoning_effort = None if reasoning is None else reasoning.effort + + extra_kwargs: dict[str, Any] = dict( + add_generation_prompt=not continue_final, + continue_final_message=continue_final, + reasoning_effort=reasoning_effort, + ) + + # When reasoning is requested, activate thinking for models whose + # chat templates require explicit opt-in (e.g., Gemma4 defaults + # enable_thinking to false). For templates that don't declare the + # variable, resolve_chat_template_kwargs filters it out harmlessly. + user_kwargs = self.chat_template_kwargs or {} + if reasoning_effort is not None and "enable_thinking" not in user_kwargs: + extra_kwargs["enable_thinking"] = reasoning_effort != "none" return ChatParams( chat_template=default_template, chat_template_content_format=default_template_content_format, - chat_template_kwargs=merge_kwargs( # To remove unset values + chat_template_kwargs=merge_kwargs( self.chat_template_kwargs, - dict( - add_generation_prompt=not continue_final, - continue_final_message=continue_final, - reasoning_effort=None if reasoning is None else reasoning.effort, - ), + extra_kwargs, ), media_io_kwargs=self.media_io_kwargs, ) From 03d9cc2fe296c58cdb7f9f77218d165a1997abdf Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Wed, 27 May 2026 11:25:36 -0400 Subject: [PATCH 11/27] [misc] Bump cutedsl version to 4.5.2 (#43745) Signed-off-by: Yongye Zhu --- requirements/cuda.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 1e0a945e5a77..99a45c9d3ca3 100644 --- a/requirements/cuda.txt +++ b/requirements/cuda.txt @@ -21,7 +21,7 @@ nvidia-cudnn-frontend>=1.13.0,<1.19.0 fastsafetensors >= 0.2.2 # QuACK and Cutlass DSL for FA4 (cute-DSL implementation) -nvidia-cutlass-dsl[cu13]==4.5.0 +nvidia-cutlass-dsl[cu13]==4.5.2 quack-kernels>=0.3.3 # Tokenspeed_MLA for faster mla with spec decode From 165460941f5e82b8ddb344515163b43e537f9c3f Mon Sep 17 00:00:00 2001 From: Injae Ryou Date: Thu, 28 May 2026 00:53:32 +0900 Subject: [PATCH 12/27] [BugFix] HFValidationError with cloud storage URIs when HF_HUB_OFFLINE=1 (#39155) Signed-off-by: Injae Ryou --- tests/engine/test_arg_utils.py | 28 ++++++++++++++++++++++++++++ tests/test_config.py | 20 ++++++++++++++++++++ vllm/config/model.py | 2 +- vllm/engine/arg_utils.py | 23 ++++++++++++++--------- 4 files changed, 63 insertions(+), 10 deletions(-) diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index f595ca6ecbd3..9b21f3eebc10 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -613,3 +613,31 @@ def test_ir_op_priority(): ) def test_expand_json_human_readable_numbers(input_json, expected_json): assert _expand_json_human_readable_numbers(input_json) == expected_json + + +@pytest.mark.parametrize( + "uri", + ["s3://bucket/model", "gs://bucket/model", "az://container/model"], +) +def test_cloud_storage_uri_skips_get_model_path(uri, monkeypatch): + """Cloud storage URIs should not be passed to get_model_path() + when HF_HUB_OFFLINE=1, as they are not valid HF repo IDs.""" + import huggingface_hub + + monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) + + args = EngineArgs(model=uri) + # model should remain the original cloud URI, not raise + assert args.model == uri + + +def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch): + """Cloud storage tokenizer URI should not be passed to + get_model_path() when HF_HUB_OFFLINE=1.""" + import huggingface_hub + + monkeypatch.setattr(huggingface_hub.constants, "HF_HUB_OFFLINE", True) + + args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer") + assert args.model == "s3://bucket/model" + assert args.tokenizer == "s3://bucket/tokenizer" diff --git a/tests/test_config.py b/tests/test_config.py index 5c01d652a17a..c0bd4b14ff8d 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -754,6 +754,26 @@ def test_s3_url_different_models_create_different_directories(mock_pull_files): assert os.path.exists(config2.tokenizer) and os.path.isdir(config2.tokenizer) +@patch("vllm.transformers_utils.runai_utils.ObjectStorageModel.pull_files") +def test_s3_url_different_model_and_tokenizer(mock_pull_files): + """Test that when model and tokenizer are different cloud URIs, + pull_files receives the correct URI for each.""" + mock_pull_files.return_value = None + + model_url = "s3://bucket/model/" + tokenizer_url = "s3://bucket/tokenizer/" + + config = MockConfig(model=model_url, tokenizer=tokenizer_url) + ModelConfig.maybe_pull_model_tokenizer_for_runai(config, model_url, tokenizer_url) + + # pull_files should be called twice: once for model, once for tokenizer + assert mock_pull_files.call_count == 2 + # First call: model URI with allow_pattern + assert mock_pull_files.call_args_list[0][0][0] == model_url + # Second call: tokenizer URI with ignore_pattern + assert mock_pull_files.call_args_list[1][0][0] == tokenizer_url + + @pytest.mark.parametrize( ("model_id", "expected_attn_type", "expected_result", "reason"), [ diff --git a/vllm/config/model.py b/vllm/config/model.py index 547ec342f967..811798859b84 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -888,7 +888,7 @@ def maybe_pull_model_tokenizer_for_runai(self, model: str, tokenizer: str) -> No if is_runai_obj_uri(tokenizer): object_storage_tokenizer = ObjectStorageModel(url=tokenizer) object_storage_tokenizer.pull_files( - model, + tokenizer, ignore_pattern=["*.pt", "*.safetensors", "*.bin", "*.tensors", "*.pth"], ) self.tokenizer = object_storage_tokenizer.dir diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 79c6e0ae925d..0490cbc3e4b5 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -747,15 +747,20 @@ def __post_init__(self): load_general_plugins() # when use hf offline,replace model and tokenizer id to local model path if huggingface_hub.constants.HF_HUB_OFFLINE: - model_id = self.model - self.model = get_model_path(self.model, self.revision) - if model_id is not self.model: - logger.info( - "HF_HUB_OFFLINE is True, replace model_id [%s] to model_path [%s]", - model_id, - self.model, - ) - if self.tokenizer is not None: + # Skip cloud storage URIs (s3://, gs://, az://) — they are not + # HF repo IDs and will be resolved later by + # ModelConfig.maybe_pull_model_tokenizer_for_runai(). + if not is_cloud_storage(self.model): + model_id = self.model + self.model = get_model_path(self.model, self.revision) + if model_id is not self.model: + logger.info( + "HF_HUB_OFFLINE is True, replace model_id " + "[%s] to model_path [%s]", + model_id, + self.model, + ) + if self.tokenizer is not None and not is_cloud_storage(self.tokenizer): tokenizer_id = self.tokenizer self.tokenizer = get_model_path(self.tokenizer, self.tokenizer_revision) if tokenizer_id is not self.tokenizer: From 49a3510266469c59ba7cb36b7223765b4aa94c4e Mon Sep 17 00:00:00 2001 From: Chunyang Wen Date: Thu, 28 May 2026 00:09:58 +0800 Subject: [PATCH 13/27] [Docs] Fix the duplicate doc icon issue (#43546) Signed-off-by: chunyang.wen --- docs/mkdocs/hooks/url_schemes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/mkdocs/hooks/url_schemes.py b/docs/mkdocs/hooks/url_schemes.py index 4d5034990683..e6faf95cd469 100644 --- a/docs/mkdocs/hooks/url_schemes.py +++ b/docs/mkdocs/hooks/url_schemes.py @@ -106,8 +106,8 @@ def replace_github_link(match: re.Match) -> str: return f"[{gh_icon} {title}]({url})" markdown = "\n".join(lines) - markdown = relative_link.sub(replace_relative_link, markdown) markdown = github_link.sub(replace_github_link, markdown) + markdown = relative_link.sub(replace_relative_link, markdown) return markdown.split("\n") From 41688e2dc7f52b4f0c22ebe5470e340bbc7e0d6f Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 27 May 2026 17:30:11 +0100 Subject: [PATCH 14/27] Fix early CUDA init (#43791) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .../deepseek_v4/common/ops/cache_utils.py | 2 +- .../deepseek_v4/common/ops/fused_indexer_q.py | 4 ++-- vllm/models/deepseek_v4/compressor.py | 4 +++- vllm/models/deepseek_v4/nvidia/model.py | 2 +- .../models/deepseek_v4/nvidia/ops/__init__.py | 21 +++++-------------- 5 files changed, 12 insertions(+), 21 deletions(-) diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index efb936be6c07..ac66751e3111 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -366,7 +366,7 @@ def dequantize_and_gather_k_cache( ) -> None: if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. - from vllm.models.deepseek_v4.nvidia.ops import ( + from vllm.models.deepseek_v4.nvidia.ops.dequant_gather_k_cutedsl import ( dequantize_and_gather_k_cache_cutedsl, ) diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index e88fe1529cd6..d5aaf10feba4 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -346,7 +346,7 @@ def fused_indexer_q_rope_quant( ) if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. - from vllm.models.deepseek_v4.nvidia.ops import ( + from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( fused_indexer_q_rope_quant_mxfp4_cutedsl, ) @@ -400,7 +400,7 @@ def fused_indexer_q_rope_quant( index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn) if has_cutedsl(): # lazily import, otherwise some tests fail due to CUDA driver init failure. - from vllm.models.deepseek_v4.nvidia.ops import ( + from vllm.models.deepseek_v4.nvidia.ops.fused_indexer_q_cutedsl import ( fused_indexer_q_rope_quant_fp8_cutedsl, ) diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 3234faa5eb05..f36dc8f17629 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -338,7 +338,9 @@ def forward( if current_platform.is_cuda(): # NVIDIA GPUs. if self.head_dim == 512: - from .nvidia.ops import compress_norm_rope_store_cutedsl + from .nvidia.ops.sparse_attn_compress_cutedsl import ( + compress_norm_rope_store_cutedsl, + ) # Main compressor path. # Use a cutedsl kernel for better performance. diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 974593a8d390..25f0a730fdbd 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -59,7 +59,7 @@ DeepseekV4MLAModules, DeepseekV4MultiHeadLatentAttentionWrapper, ) -from vllm.models.deepseek_v4.nvidia.ops import prepare_megamoe_inputs +from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.torch_utils import direct_register_custom_op diff --git a/vllm/models/deepseek_v4/nvidia/ops/__init__.py b/vllm/models/deepseek_v4/nvidia/ops/__init__.py index dca25345ea6f..20752bc15489 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/__init__.py +++ b/vllm/models/deepseek_v4/nvidia/ops/__init__.py @@ -5,20 +5,9 @@ These modules import ``cutlass``/``cutedsl`` at module top level, so they must not be imported on non-CUDA platforms. Callers should gate on ``vllm.utils.import_utils.has_cutedsl()`` before importing from here. -""" - -from .dequant_gather_k_cutedsl import dequantize_and_gather_k_cache_cutedsl -from .fused_indexer_q_cutedsl import ( - fused_indexer_q_rope_quant_fp8_cutedsl, - fused_indexer_q_rope_quant_mxfp4_cutedsl, -) -from .prepare_megamoe import prepare_megamoe_inputs -from .sparse_attn_compress_cutedsl import compress_norm_rope_store_cutedsl -__all__ = [ - "compress_norm_rope_store_cutedsl", - "dequantize_and_gather_k_cache_cutedsl", - "fused_indexer_q_rope_quant_fp8_cutedsl", - "fused_indexer_q_rope_quant_mxfp4_cutedsl", - "prepare_megamoe_inputs", -] +This ``__init__`` deliberately imports nothing: re-exporting the cutedsl +modules here would eagerly ``import cutlass`` (initializing the CUDA driver) for +anyone who imports ``vllm.models.deepseek_v4``, breaking forked subprocesses. +Import the leaf modules directly under a ``has_cutedsl()``/``is_cuda()`` gate. +""" From 05c50c721e07035dd36169fa8bb2825a7b952555 Mon Sep 17 00:00:00 2001 From: jatseng-ai Date: Wed, 27 May 2026 09:33:32 -0700 Subject: [PATCH 15/27] [ROCm] mori: add InterNodeV1LL inter-node kernel selection via VLLM_MORI_INTERNODE_KERNEL (#41751) Signed-off-by: jatseng-ai --- tests/kernels/moe/modular_kernel_tools/common.py | 2 +- .../moe/modular_kernel_tools/mk_objects.py | 2 +- tests/kernels/moe/test_moe_layer.py | 8 +++++--- vllm/config/parallel.py | 9 ++++++--- vllm/distributed/device_communicators/all2all.py | 15 ++++++++++++--- .../device_communicators/cuda_communicator.py | 9 +++++++-- vllm/model_executor/layers/fused_moe/config.py | 5 ++++- 7 files changed, 36 insertions(+), 14 deletions(-) diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index f07d4c75e752..88aa9f2dc19c 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -249,7 +249,7 @@ def needs_aiter(self): def needs_mori(self): info = prepare_finalize_info(self.prepare_finalize_type) - return info.backend == "mori" + return info.backend in ("mori_high_throughput", "mori_low_latency") def all2all_backend(self): info = prepare_finalize_info(self.prepare_finalize_type) diff --git a/tests/kernels/moe/modular_kernel_tools/mk_objects.py b/tests/kernels/moe/modular_kernel_tools/mk_objects.py index 5d3572b7caa2..7c3bde2eafa1 100644 --- a/tests/kernels/moe/modular_kernel_tools/mk_objects.py +++ b/tests/kernels/moe/modular_kernel_tools/mk_objects.py @@ -232,7 +232,7 @@ def expert_info(kind) -> ExpertInfo: standard_format, fp8_types, blocked_quantization_support=True, - backend="mori", + backend="mori_high_throughput", supports_apply_weight_on_input=False, ) diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 7942f738a1e5..1bb899c1c891 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -92,7 +92,7 @@ BACKENDS = ["allgather_reducescatter"] if has_mori(): - BACKENDS += ["mori"] + BACKENDS += ["mori_high_throughput", "mori_low_latency"] if has_flashinfer_nvlink_two_sided(): BACKENDS += ["flashinfer_nvlink_two_sided"] @@ -118,7 +118,8 @@ # fmt: off BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = { "allgather_reducescatter": {None, "fp8", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501 - "mori": {None, "fp8", "modelopt_fp8"}, + "mori_high_throughput": {None, "fp8", "modelopt_fp8"}, + "mori_low_latency": {None, "fp8", "modelopt_fp8"}, "flashinfer_nvlink_two_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 "flashinfer_nvlink_one_sided": {None, "modelopt_fp4"}, # noqa: E501 "deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 @@ -129,7 +130,8 @@ # Map from backend -> (DP/EP support, DP support, TP support, SP support) BACKEND_EP_DP_TP_SUPPORT: dict[str, tuple[bool, bool, bool, bool]] = { "allgather_reducescatter": (True, True, True, True), - "mori": (True, False, False, True), + "mori_high_throughput": (True, False, False, True), + "mori_low_latency": (True, False, False, True), "flashinfer_nvlink_two_sided": (False, True, False, False), "flashinfer_nvlink_one_sided": (False, True, False, False), "deepep_low_latency": (True, False, False, True), diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 6e8f78e4ee30..960e1d343c05 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -42,7 +42,8 @@ "pplx", "deepep_high_throughput", "deepep_low_latency", - "mori", + "mori_high_throughput", + "mori_low_latency", "nixl_ep", "allgather_reducescatter", "flashinfer_all2allv", # temporary alias for flashinfer_nvlink_two_sided @@ -177,7 +178,8 @@ class ParallelConfig: - "allgather_reducescatter": All2all based on allgather and reducescatter - "deepep_high_throughput": Use deepep high-throughput kernels - "deepep_low_latency": Use deepep low-latency kernels - - "mori": Use mori kernels + - "mori_high_throughput": MoRI EP with InterNodeV1 for multi-node + - "mori_low_latency": MoRI EP with InterNodeV1LL for multi-node - "nixl_ep": Use nixl-ep kernels - "flashinfer_nvlink_two_sided": Use flashinfer two-sided kernels for mnnvl - "flashinfer_nvlink_one_sided": Use flashinfer high-throughput a2a kernels""" @@ -621,7 +623,8 @@ def use_sequence_parallel_moe(self) -> bool: "allgather_reducescatter", "deepep_high_throughput", "deepep_low_latency", - "mori", + "mori_high_throughput", + "mori_low_latency", "nixl_ep", ) and self.enable_expert_parallel diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 33ff55a64e66..c0055f4bbe1f 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -764,14 +764,19 @@ def cleanup(self): class MoriAll2AllManager(All2AllManagerBase): - def __init__(self, cpu_group): + def __init__(self, cpu_group, all2all_backend: str): assert has_mori(), ( "MoRI kernels not found. Please follow https://github.com/ROCm/mori/blob/main/README.md" " to install MoRI kernels." ) # noqa + assert all2all_backend in ( + "mori_high_throughput", + "mori_low_latency", + ), f"unsupported MoRI all2all backend: {all2all_backend!r}" import mori super().__init__(cpu_group) + self._all2all_backend = all2all_backend self.handle_cache = Cache() torch._C._distributed_c10d._register_process_group("mori", cpu_group) @@ -805,8 +810,12 @@ def _make_all2all_kwargs( warp_num_per_block = 16 block_num = 80 else: - # multi node - kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1 + # Multi-node: kernel follows --all2all-backend (mirrors deepep_* split). + # mori_low_latency → InterNodeV1LL; mori_high_throughput → V1. + if self._all2all_backend == "mori_low_latency": + kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1LL + else: + kernel_type = mori.ops.EpDispatchCombineKernelType.InterNodeV1 if on_gfx942(): warp_num_per_block = 16 block_num = 32 diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index ee81bc20f3c2..ff4929d8415d 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -137,10 +137,15 @@ def __init__( self.all2all_manager = DeepEPLLAll2AllManager( self.cpu_group, tcp_store_group ) - elif self.all2all_backend == "mori": + elif self.all2all_backend in ( + "mori_high_throughput", + "mori_low_latency", + ): from .all2all import MoriAll2AllManager - self.all2all_manager = MoriAll2AllManager(self.cpu_group) + self.all2all_manager = MoriAll2AllManager( + self.cpu_group, self.all2all_backend + ) elif self.all2all_backend == "nixl_ep": from .all2all import NixlEPAll2AllManager diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 17cd5592d1d6..8e92d0cfb674 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1091,7 +1091,10 @@ def use_ag_rs_all2all_kernels(self): @property def use_mori_kernels(self): - return self.use_all2all_kernels and self.all2all_backend == "mori" + return self.use_all2all_kernels and self.all2all_backend in ( + "mori_high_throughput", + "mori_low_latency", + ) @property def use_nixl_ep_kernels(self): From 284e6f543d462016fc80c055ccbf088832c63129 Mon Sep 17 00:00:00 2001 From: Chris Leonard Date: Wed, 27 May 2026 12:35:24 -0400 Subject: [PATCH 16/27] [8/n] Migrate merge_attn_states, mamba, sampler to torch stable ABI (continued) (#43361) Signed-off-by: Mikayla Gawarecki Signed-off-by: Chris Leonard Co-authored-by: Mikayla Gawarecki Co-authored-by: Shengqi Chen --- CMakeLists.txt | 10 +- .../attention/merge_attn_states.cu | 105 +++++---- .../mamba}/selective_scan.h | 11 +- .../mamba}/selective_scan_fwd.cu | 213 +++++++++--------- .../mamba}/static_switch.h | 0 csrc/libtorch_stable/ops.h | 53 +++++ csrc/{ => libtorch_stable}/sampler.cu | 115 +++++----- csrc/{ => libtorch_stable}/topk.cu | 147 ++++++------ csrc/libtorch_stable/torch_bindings.cpp | 62 +++++ csrc/ops.h | 43 ---- csrc/persistent_topk.cuh | 16 +- csrc/torch_bindings.cpp | 59 ----- pyproject.toml | 1 + 13 files changed, 432 insertions(+), 403 deletions(-) rename csrc/{ => libtorch_stable}/attention/merge_attn_states.cu (81%) rename csrc/{mamba/mamba_ssm => libtorch_stable/mamba}/selective_scan.h (95%) rename csrc/{mamba/mamba_ssm => libtorch_stable/mamba}/selective_scan_fwd.cu (82%) rename csrc/{mamba/mamba_ssm => libtorch_stable/mamba}/static_switch.h (100%) rename csrc/{ => libtorch_stable}/sampler.cu (87%) rename csrc/{ => libtorch_stable}/topk.cu (63%) diff --git a/CMakeLists.txt b/CMakeLists.txt index f0b1f53af831..e416d0529194 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -305,14 +305,10 @@ endif() # set(VLLM_EXT_SRC - "csrc/mamba/mamba_ssm/selective_scan_fwd.cu" "csrc/cache_kernels.cu" "csrc/cache_kernels_fused.cu" "csrc/attention/paged_attention_v1.cu" "csrc/attention/paged_attention_v2.cu" - "csrc/attention/merge_attn_states.cu" - "csrc/sampler.cu" - "csrc/topk.cu" "csrc/cuda_view.cu" "csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu" "csrc/quantization/activation_kernels.cu" @@ -633,7 +629,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/fused_qknorm_rope_kernel.cu" "csrc/libtorch_stable/layernorm_kernels.cu" "csrc/libtorch_stable/layernorm_quant_kernels.cu" - "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu") + "csrc/libtorch_stable/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu" + "csrc/libtorch_stable/attention/merge_attn_states.cu" + "csrc/libtorch_stable/sampler.cu" + "csrc/libtorch_stable/topk.cu" + "csrc/libtorch_stable/mamba/selective_scan_fwd.cu") if(VLLM_GPU_LANG STREQUAL "CUDA") list(APPEND VLLM_STABLE_EXT_SRC diff --git a/csrc/attention/merge_attn_states.cu b/csrc/libtorch_stable/attention/merge_attn_states.cu similarity index 81% rename from csrc/attention/merge_attn_states.cu rename to csrc/libtorch_stable/attention/merge_attn_states.cu index 75f066e80915..88e7723ad2f2 100644 --- a/csrc/attention/merge_attn_states.cu +++ b/csrc/libtorch_stable/attention/merge_attn_states.cu @@ -1,14 +1,14 @@ #include -#include -#include -#include #include #include -#include "attention_dtypes.h" -#include "attention_utils.cuh" -#include "../quantization/w8a8/fp8/common.cuh" +#include "../torch_utils.h" #include "../dispatch_utils.h" +#include + +#include "../../attention/attention_dtypes.h" +#include "../../attention/attention_utils.cuh" +#include "../../quantization/w8a8/fp8/common.cuh" namespace vllm { @@ -196,17 +196,17 @@ __global__ void merge_attn_states_kernel( // The following macro is used to dispatch the conversion function based on // the output data type. The FN is a macro that calls a function with // template. -#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn) \ - { \ - if (scalar_dtype == at::ScalarType::Float) { \ - fn(float); \ - } else if (scalar_dtype == at::ScalarType::Half) { \ - fn(uint16_t); \ - } else if (scalar_dtype == at::ScalarType::BFloat16) { \ - fn(__nv_bfloat16); \ - } else { \ - TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \ - } \ +#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn) \ + { \ + if (scalar_dtype == torch::headeronly::ScalarType::Float) { \ + fn(float); \ + } else if (scalar_dtype == torch::headeronly::ScalarType::Half) { \ + fn(uint16_t); \ + } else if (scalar_dtype == torch::headeronly::ScalarType::BFloat16) { \ + fn(__nv_bfloat16); \ + } else { \ + STD_TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \ + } \ } #define LAUNCH_MERGE_ATTN_STATES(scalar_t, output_t, NUM_THREADS, \ @@ -245,11 +245,14 @@ __global__ void merge_attn_states_kernel( */ template void merge_attn_states_launcher( - torch::Tensor& output, std::optional output_lse, - const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse, - const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse, + torch::stable::Tensor& output, + std::optional output_lse, + const torch::stable::Tensor& prefix_output, + const torch::stable::Tensor& prefix_lse, + const torch::stable::Tensor& suffix_output, + const torch::stable::Tensor& suffix_lse, const std::optional prefill_tokens_with_context, - const std::optional& output_scale) { + const std::optional& output_scale) { constexpr uint NUM_THREADS = 128; const uint num_tokens = output.size(0); const uint num_heads = output.size(1); @@ -258,23 +261,23 @@ void merge_attn_states_launcher( const uint output_head_stride = output.stride(1); // Thread mapping is based on input BF16 pack_size const uint pack_size = 16 / sizeof(scalar_t); - TORCH_CHECK(head_size % pack_size == 0, - "headsize must be multiple of pack_size:", pack_size); + STD_TORCH_CHECK(head_size % pack_size == 0, + "headsize must be multiple of pack_size:", pack_size); const uint prefix_num_tokens = prefill_tokens_with_context.has_value() ? static_cast(prefill_tokens_with_context.value()) : num_tokens; - TORCH_CHECK(prefix_num_tokens <= num_tokens, - "prefix_num_tokens must be <= num_tokens"); + STD_TORCH_CHECK(prefix_num_tokens <= num_tokens, + "prefix_num_tokens must be <= num_tokens"); float* output_lse_ptr = nullptr; if (output_lse.has_value()) { - output_lse_ptr = output_lse.value().data_ptr(); + output_lse_ptr = output_lse.value().mutable_data_ptr(); } float* output_scale_ptr = nullptr; if (output_scale.has_value()) { - output_scale_ptr = output_scale.value().data_ptr(); + output_scale_ptr = output_scale.value().mutable_data_ptr(); } // Process one pack elements per thread. for float, the // pack_size is 4 for half/bf16, the pack_size is 8. @@ -284,14 +287,15 @@ void merge_attn_states_launcher( dim3 block(NUM_THREADS); dim3 grid((total_threads + NUM_THREADS - 1) / NUM_THREADS); - const c10::cuda::OptionalCUDAGuard device_guard(prefix_output.device()); - auto stream = at::cuda::getCurrentCUDAStream(); + const torch::stable::accelerator::DeviceGuard device_guard( + prefix_output.get_device_index()); + auto stream = get_current_cuda_stream(); if (output_scale.has_value()) { // FP8 output path - dispatch on output FP8 type - VLLM_DISPATCH_FP8_TYPES(output.scalar_type(), "merge_attn_states_fp8", [&] { - LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true); - }); + VLLM_STABLE_DISPATCH_FP8_TYPES( + output.scalar_type(), "merge_attn_states_fp8", + [&] { LAUNCH_MERGE_ATTN_STATES(scalar_t, fp8_t, NUM_THREADS, true); }); } else { // Original BF16/FP16/FP32 output path LAUNCH_MERGE_ATTN_STATES(scalar_t, scalar_t, NUM_THREADS, false); @@ -305,26 +309,29 @@ void merge_attn_states_launcher( suffix_lse, prefill_tokens_with_context, output_scale); \ } -void merge_attn_states(torch::Tensor& output, - std::optional output_lse, - const torch::Tensor& prefix_output, - const torch::Tensor& prefix_lse, - const torch::Tensor& suffix_output, - const torch::Tensor& suffix_lse, - std::optional prefill_tokens_with_context, - const std::optional& output_scale) { +void merge_attn_states( + torch::stable::Tensor& output, + std::optional output_lse, + const torch::stable::Tensor& prefix_output, + const torch::stable::Tensor& prefix_lse, + const torch::stable::Tensor& suffix_output, + const torch::stable::Tensor& suffix_lse, + const std::optional prefill_tokens_with_context, + const std::optional& output_scale) { if (output_scale.has_value()) { - TORCH_CHECK(output.scalar_type() == at::ScalarType::Float8_e4m3fn || - output.scalar_type() == at::ScalarType::Float8_e4m3fnuz, - "output must be FP8 when output_scale is provided, got: ", - output.scalar_type()); + STD_TORCH_CHECK( + output.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn || + output.scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fnuz, + "output must be FP8 when output_scale is provided, got: ", + output.scalar_type()); } else { - TORCH_CHECK(output.scalar_type() == prefix_output.scalar_type(), - "output dtype (", output.scalar_type(), - ") must match prefix_output dtype (", - prefix_output.scalar_type(), ") when output_scale is not set"); + STD_TORCH_CHECK( + output.scalar_type() == prefix_output.scalar_type(), "output dtype (", + output.scalar_type(), ") must match prefix_output dtype (", + prefix_output.scalar_type(), ") when output_scale is not set"); } // Always dispatch on prefix_output (input) dtype - DISPATCH_BY_SCALAR_DTYPE(prefix_output.dtype(), + DISPATCH_BY_SCALAR_DTYPE(prefix_output.scalar_type(), CALL_MERGE_ATTN_STATES_LAUNCHER); } diff --git a/csrc/mamba/mamba_ssm/selective_scan.h b/csrc/libtorch_stable/mamba/selective_scan.h similarity index 95% rename from csrc/mamba/mamba_ssm/selective_scan.h rename to csrc/libtorch_stable/mamba/selective_scan.h index ff1d9528e0f6..d0cb2bc8272d 100644 --- a/csrc/mamba/mamba_ssm/selective_scan.h +++ b/csrc/libtorch_stable/mamba/selective_scan.h @@ -12,6 +12,9 @@ #include #endif #include + +#include +#include //////////////////////////////////////////////////////////////////////////////////////////////////// struct SSMParamsBase { @@ -159,8 +162,8 @@ struct Converter{ }; template -struct Converter{ - static inline __device__ void to_float(const at::Half (&src)[N], float (&dst)[N]) { +struct Converter{ + static inline __device__ void to_float(const torch::headeronly::Half (&src)[N], float (&dst)[N]) { static_assert(N % 2 == 0); auto &src2 = reinterpret_cast(src); auto &dst2 = reinterpret_cast(dst); @@ -171,8 +174,8 @@ struct Converter{ #if __CUDA_ARCH__ >= 800 template -struct Converter{ - static inline __device__ void to_float(const at::BFloat16 (&src)[N], float (&dst)[N]) { +struct Converter{ + static inline __device__ void to_float(const torch::headeronly::BFloat16 (&src)[N], float (&dst)[N]) { static_assert(N % 2 == 0); auto &src2 = reinterpret_cast(src); auto &dst2 = reinterpret_cast(dst); diff --git a/csrc/mamba/mamba_ssm/selective_scan_fwd.cu b/csrc/libtorch_stable/mamba/selective_scan_fwd.cu similarity index 82% rename from csrc/mamba/mamba_ssm/selective_scan_fwd.cu rename to csrc/libtorch_stable/mamba/selective_scan_fwd.cu index ba2f0cc61942..c17b06a79bd0 100644 --- a/csrc/mamba/mamba_ssm/selective_scan_fwd.cu +++ b/csrc/libtorch_stable/mamba/selective_scan_fwd.cu @@ -1,18 +1,9 @@ // clang-format off // adapted from https://github.com/state-spaces/mamba/blob/main/csrc/selective_scan/selective_scan_fwd_kernel.cuh -#include -#include -#include +#include "../torch_utils.h" +#include #include "selective_scan.h" -#include -#include -#ifdef USE_ROCM - #include // For C10_HIP_CHECK and C10_HIP_KERNEL_LAUNCH_CHECK -#else - #include // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK -#endif - #ifndef USE_ROCM #include #include @@ -416,15 +407,15 @@ void selective_scan_fwd_launch(SSMParamsBase ¶ms, cudaStream_t stream) { auto kernel = &selective_scan_fwd_kernel; if (kSmemSize >= 48 * 1024) { #ifdef USE_ROCM - C10_HIP_CHECK(hipFuncSetAttribute( + STD_CUDA_CHECK(hipFuncSetAttribute( reinterpret_cast(kernel), hipFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); #else - C10_CUDA_CHECK(cudaFuncSetAttribute( + STD_CUDA_CHECK(cudaFuncSetAttribute( kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize)); #endif } kernel<<>>(params); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + STD_CUDA_KERNEL_LAUNCH_CHECK(); }); }); }); @@ -462,46 +453,46 @@ void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream) { #endif } -template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); -template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); -template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); -template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); +template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); +template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); +template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); +template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); template void selective_scan_fwd_cuda(SSMParamsBase ¶ms, cudaStream_t stream); -#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")") +#define CHECK_SHAPE(x, ...) STD_TORCH_CHECK(x.sizes().equals(torch::headeronly::IntHeaderOnlyArrayRef({__VA_ARGS__})), #x " must have shape (" #__VA_ARGS__ ")") #define DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, STYPE, NAME, ...) \ - if (ITYPE == at::ScalarType::Half) { \ - using input_t = at::Half; \ + if (ITYPE == torch::headeronly::ScalarType::Half) { \ + using input_t = torch::headeronly::Half; \ using weight_t = float; \ - if (STYPE == at::ScalarType::Half) { \ - using state_t = at::Half; \ + if (STYPE == torch::headeronly::ScalarType::Half) { \ + using state_t = torch::headeronly::Half; \ __VA_ARGS__(); \ - } else if (STYPE == at::ScalarType::Float) { \ + } else if (STYPE == torch::headeronly::ScalarType::Float) { \ using state_t = float; \ __VA_ARGS__(); \ } else { \ - AT_ERROR(#NAME, " not implemented for state type '", toString(STYPE), "'"); \ + STD_TORCH_CHECK(false, #NAME " not implemented for state type '", STYPE, "'"); \ } \ - } else if (ITYPE == at::ScalarType::BFloat16) { \ - using input_t = at::BFloat16; \ + } else if (ITYPE == torch::headeronly::ScalarType::BFloat16) { \ + using input_t = torch::headeronly::BFloat16; \ using weight_t = float; \ - if (STYPE == at::ScalarType::BFloat16) { \ - using state_t = at::BFloat16; \ + if (STYPE == torch::headeronly::ScalarType::BFloat16) { \ + using state_t = torch::headeronly::BFloat16; \ __VA_ARGS__(); \ - } else if (STYPE == at::ScalarType::Float) { \ + } else if (STYPE == torch::headeronly::ScalarType::Float) { \ using state_t = float; \ __VA_ARGS__(); \ } else { \ - AT_ERROR(#NAME, " not implemented for state type '", toString(STYPE), "'"); \ + STD_TORCH_CHECK(false, #NAME " not implemented for state type '", STYPE, "'"); \ } \ - } else if (ITYPE == at::ScalarType::Float) { \ + } else if (ITYPE == torch::headeronly::ScalarType::Float) { \ using input_t = float; \ using weight_t = float; \ using state_t = float; \ __VA_ARGS__(); \ } else { \ - AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \ + STD_TORCH_CHECK(false, #NAME " not implemented for input type '", ITYPE, "'"); \ } @@ -518,30 +509,30 @@ void set_ssm_params_fwd(SSMParamsBase ¶ms, const bool is_variable_B, const bool is_variable_C, // device pointers - const torch::Tensor u, - const torch::Tensor delta, - const torch::Tensor A, - const torch::Tensor B, - const torch::Tensor C, - const torch::Tensor out, - const torch::Tensor z, - const torch::Tensor out_z, - const std::optional& D, - const std::optional& delta_bias, - const torch::Tensor ssm_states, + const torch::stable::Tensor u, + const torch::stable::Tensor delta, + const torch::stable::Tensor A, + const torch::stable::Tensor B, + const torch::stable::Tensor C, + const torch::stable::Tensor out, + const torch::stable::Tensor z, + const torch::stable::Tensor out_z, + const std::optional& D, + const std::optional& delta_bias, + const torch::stable::Tensor ssm_states, bool has_z, bool delta_softplus, - const std::optional& query_start_loc, - const std::optional& cache_indices, - const std::optional& has_initial_state, + const std::optional& query_start_loc, + const std::optional& cache_indices, + const std::optional& has_initial_state, bool varlen, int64_t null_block_id, int64_t block_size, - const std::optional &block_idx_first_scheduled_token, - const std::optional &block_idx_last_scheduled_token, - const std::optional &initial_state_idx, - const std::optional &cu_chunk_seqlen, - const std::optional &last_chunk_indices) { + const std::optional &block_idx_first_scheduled_token, + const std::optional &block_idx_last_scheduled_token, + const std::optional &initial_state_idx, + const std::optional &cu_chunk_seqlen, + const std::optional &last_chunk_indices) { // Reset the parameters memset(¶ms, 0, sizeof(params)); @@ -654,45 +645,45 @@ void set_ssm_params_fwd(SSMParamsBase ¶ms, } } -void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta, - const torch::Tensor &A, const torch::Tensor &B, const torch::Tensor &C, - const std::optional &D_, - const std::optional &z_, - const std::optional &delta_bias_, +void selective_scan_fwd(const torch::stable::Tensor &u, const torch::stable::Tensor &delta, + const torch::stable::Tensor &A, const torch::stable::Tensor &B, const torch::stable::Tensor &C, + const std::optional &D_, + const std::optional &z_, + const std::optional &delta_bias_, bool delta_softplus, - const std::optional &query_start_loc, - const std::optional &cache_indices, - const std::optional &has_initial_state, - const torch::Tensor &ssm_states, + const std::optional &query_start_loc, + const std::optional &cache_indices, + const std::optional &has_initial_state, + const torch::stable::Tensor &ssm_states, // used to identify padding entries if cache_indices provided // in case of padding, the kernel will return early int64_t null_block_id, int64_t block_size, - const std::optional &block_idx_first_scheduled_token, - const std::optional &block_idx_last_scheduled_token, - const std::optional &initial_state_idx, - const std::optional &cu_chunk_seqlen, - const std::optional &last_chunk_indices) { + const std::optional &block_idx_first_scheduled_token, + const std::optional &block_idx_last_scheduled_token, + const std::optional &initial_state_idx, + const std::optional &cu_chunk_seqlen, + const std::optional &last_chunk_indices) { auto input_type = u.scalar_type(); auto weight_type = A.scalar_type(); - TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16); - TORCH_CHECK(weight_type == at::ScalarType::Float); + STD_TORCH_CHECK(input_type == torch::headeronly::ScalarType::Float || input_type == torch::headeronly::ScalarType::Half || input_type == torch::headeronly::ScalarType::BFloat16); + STD_TORCH_CHECK(weight_type == torch::headeronly::ScalarType::Float); const bool is_variable_B = B.dim() >= 3; const bool is_variable_C = C.dim() >= 3; - TORCH_CHECK(delta.scalar_type() == input_type); - TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type)); - TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type)); + STD_TORCH_CHECK(delta.scalar_type() == input_type); + STD_TORCH_CHECK(B.scalar_type() == (!is_variable_B ? weight_type : input_type)); + STD_TORCH_CHECK(C.scalar_type() == (!is_variable_C ? weight_type : input_type)); - TORCH_CHECK(u.is_cuda()); - TORCH_CHECK(delta.is_cuda()); - TORCH_CHECK(A.is_cuda()); - TORCH_CHECK(B.is_cuda()); - TORCH_CHECK(C.is_cuda()); + STD_TORCH_CHECK(u.is_cuda()); + STD_TORCH_CHECK(delta.is_cuda()); + STD_TORCH_CHECK(A.is_cuda()); + STD_TORCH_CHECK(B.is_cuda()); + STD_TORCH_CHECK(C.is_cuda()); - TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1); - TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1); + STD_TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1); + STD_TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1); const auto sizes = u.sizes(); const bool varlen = query_start_loc.has_value(); @@ -702,7 +693,7 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta, const int dstate = A.size(1); const int n_groups = varlen ? B.size(0) : B.size(1); - TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256"); + STD_TORCH_CHECK(dstate <= 256, "selective_scan only supports state dimension <= 256"); if (varlen) { CHECK_SHAPE(u, dim, seqlen); @@ -712,94 +703,94 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta, CHECK_SHAPE(delta, batch_size, dim, seqlen); } CHECK_SHAPE(A, dim, dstate); - TORCH_CHECK(is_variable_B, "is_variable_B = False is disabled in favor of reduced binary size") + STD_TORCH_CHECK(is_variable_B, "is_variable_B = False is disabled in favor of reduced binary size"); if (varlen) { CHECK_SHAPE(B, n_groups, dstate, seqlen); } else { - CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen); + CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen); } - TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1); + STD_TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1); - TORCH_CHECK(is_variable_C, "is_variable_C = False is disabled in favor of reduced binary size") + STD_TORCH_CHECK(is_variable_C, "is_variable_C = False is disabled in favor of reduced binary size"); if (varlen) { CHECK_SHAPE(C, n_groups, dstate, seqlen); } else { - CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen); + CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen); } - TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1); + STD_TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1); if (D_.has_value()) { auto D = D_.value(); - TORCH_CHECK(D.scalar_type() == at::ScalarType::Float); - TORCH_CHECK(D.is_cuda()); - TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1); + STD_TORCH_CHECK(D.scalar_type() == torch::headeronly::ScalarType::Float); + STD_TORCH_CHECK(D.is_cuda()); + STD_TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1); CHECK_SHAPE(D, dim); } if (delta_bias_.has_value()) { auto delta_bias = delta_bias_.value(); - TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float); - TORCH_CHECK(delta_bias.is_cuda()); - TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1); + STD_TORCH_CHECK(delta_bias.scalar_type() == torch::headeronly::ScalarType::Float); + STD_TORCH_CHECK(delta_bias.is_cuda()); + STD_TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1); CHECK_SHAPE(delta_bias, dim); } if (has_initial_state.has_value()) { auto has_initial_state_ = has_initial_state.value(); - TORCH_CHECK(has_initial_state_.scalar_type() == at::ScalarType::Bool); - TORCH_CHECK(has_initial_state_.is_cuda()); + STD_TORCH_CHECK(has_initial_state_.scalar_type() == torch::headeronly::ScalarType::Bool); + STD_TORCH_CHECK(has_initial_state_.is_cuda()); CHECK_SHAPE(has_initial_state_, batch_size); } if (query_start_loc.has_value()) { auto query_start_loc_ = query_start_loc.value(); - TORCH_CHECK(query_start_loc_.scalar_type() == at::ScalarType::Int); - TORCH_CHECK(query_start_loc_.is_cuda()); + STD_TORCH_CHECK(query_start_loc_.scalar_type() == torch::headeronly::ScalarType::Int); + STD_TORCH_CHECK(query_start_loc_.is_cuda()); } if (cache_indices.has_value()) { auto cache_indices_ = cache_indices.value(); - TORCH_CHECK(cache_indices_.scalar_type() == at::ScalarType::Int); - TORCH_CHECK(cache_indices_.is_cuda()); + STD_TORCH_CHECK(cache_indices_.scalar_type() == torch::headeronly::ScalarType::Int); + STD_TORCH_CHECK(cache_indices_.is_cuda()); // cache_indices can be either 1D (batch_size,) for non-APC mode // or 2D (batch_size, max_positions) for APC mode const bool is_apc_mode = block_idx_first_scheduled_token.has_value(); if (is_apc_mode) { - TORCH_CHECK(cache_indices_.dim() == 2, "cache_indices must be 2D for APC mode"); - TORCH_CHECK(cache_indices_.size(0) == batch_size, "cache_indices first dimension must match batch_size"); + STD_TORCH_CHECK(cache_indices_.dim() == 2, "cache_indices must be 2D for APC mode"); + STD_TORCH_CHECK(cache_indices_.size(0) == batch_size, "cache_indices first dimension must match batch_size"); } else { CHECK_SHAPE(cache_indices_, batch_size); } } - - at::Tensor z, out_z; + + torch::stable::Tensor z, out_z; const bool has_z = z_.has_value(); if (has_z) { z = z_.value(); - TORCH_CHECK(z.scalar_type() == input_type); - TORCH_CHECK(z.is_cuda()); - TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1); + STD_TORCH_CHECK(z.scalar_type() == input_type); + STD_TORCH_CHECK(z.is_cuda()); + STD_TORCH_CHECK(z.stride(-1) == 1 || z.size(-1) == 1); if (varlen){ CHECK_SHAPE(z, dim, seqlen); } else { CHECK_SHAPE(z, batch_size, dim, seqlen); } - + out_z = z; } // Right now u has BHL layout and delta has HBL layout, and we want out to have HBL layout - at::Tensor out = delta; + torch::stable::Tensor out = delta; // ssm_states can now be either the same as input_type or float32 auto state_type = ssm_states.scalar_type(); - TORCH_CHECK(state_type == input_type || state_type == at::ScalarType::Float); - TORCH_CHECK(ssm_states.is_cuda()); - TORCH_CHECK(ssm_states.stride(-1) == 1); + STD_TORCH_CHECK(state_type == input_type || state_type == torch::headeronly::ScalarType::Float); + STD_TORCH_CHECK(ssm_states.is_cuda()); + STD_TORCH_CHECK(ssm_states.stride(-1) == 1); SSMParamsBase params; set_ssm_params_fwd(params, batch_size, dim, seqlen, dstate, n_groups, is_variable_B, is_variable_C, @@ -823,8 +814,8 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta, ); - const at::cuda::OptionalCUDAGuard device_guard(device_of(u)); - auto stream = at::cuda::getCurrentCUDAStream().stream(); + const torch::stable::accelerator::DeviceGuard device_guard(u.get_device_index()); + auto stream = get_current_cuda_stream(); DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), ssm_states.scalar_type(), "selective_scan_fwd", [&] { selective_scan_fwd_cuda(params, stream); }); diff --git a/csrc/mamba/mamba_ssm/static_switch.h b/csrc/libtorch_stable/mamba/static_switch.h similarity index 100% rename from csrc/mamba/mamba_ssm/static_switch.h rename to csrc/libtorch_stable/mamba/static_switch.h diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index f99ff1d1db5b..90082ac3d291 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -164,6 +164,17 @@ torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel, #endif +// Attention kernels (shared CUDA/ROCm) +void merge_attn_states( + torch::stable::Tensor& output, + std::optional output_lse, + const torch::stable::Tensor& prefix_output, + const torch::stable::Tensor& prefix_lse, + const torch::stable::Tensor& suffix_output, + const torch::stable::Tensor& suffix_lse, + const std::optional prefill_tokens_with_context, + const std::optional& output_scale = std::nullopt); + torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x, bool inplace); @@ -220,6 +231,48 @@ void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q, torch::stable::Tensor& position_ids, int64_t forced_token_heads_per_warp); +// Sampler kernels (shared CUDA/ROCm) +void apply_repetition_penalties_( + torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask, + const torch::stable::Tensor& output_mask, + const torch::stable::Tensor& repetition_penalties); + +void top_k_per_row_prefill(const torch::stable::Tensor& logits, + const torch::stable::Tensor& rowStarts, + const torch::stable::Tensor& rowEnds, + torch::stable::Tensor& indices, int64_t numRows, + int64_t stride0, int64_t stride1, int64_t topK); + +void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, + const torch::stable::Tensor& seqLens, + torch::stable::Tensor& indices, int64_t numRows, + int64_t stride0, int64_t stride1, int64_t topK); + +void persistent_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, + int64_t max_seq_len); + +void selective_scan_fwd( + const torch::stable::Tensor& u, const torch::stable::Tensor& delta, + const torch::stable::Tensor& A, const torch::stable::Tensor& B, + const torch::stable::Tensor& C, + const std::optional& D_, + const std::optional& z_, + const std::optional& delta_bias_, + bool delta_softplus, + const std::optional& query_start_loc, + const std::optional& cache_indices, + const std::optional& has_initial_state, + const torch::stable::Tensor& ssm_states, int64_t null_block_id, + int64_t block_size, + const std::optional& block_idx_first_scheduled_token, + const std::optional& block_idx_last_scheduled_token, + const std::optional& initial_state_idx, + const std::optional& cu_chunk_seqlen, + const std::optional& last_chunk_indices); + // Activation kernels (shared CUDA/ROCm) void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input); void silu_and_mul_clamp(torch::stable::Tensor& out, diff --git a/csrc/sampler.cu b/csrc/libtorch_stable/sampler.cu similarity index 87% rename from csrc/sampler.cu rename to csrc/libtorch_stable/sampler.cu index 14d84013c08d..68848b84566c 100644 --- a/csrc/sampler.cu +++ b/csrc/libtorch_stable/sampler.cu @@ -1,8 +1,6 @@ -#include "cuda_compat.h" +#include "../cuda_compat.h" #include "dispatch_utils.h" - -#include -#include +#include "torch_utils.h" #ifndef USE_ROCM #include @@ -618,14 +616,14 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode( } // namespace vllm void apply_repetition_penalties_( - torch::Tensor& logits, // [num_seqs, vocab_size], in-place - const torch::Tensor& prompt_mask, // [num_seqs, vocab_size] - const torch::Tensor& output_mask, // [num_seqs, vocab_size] - const torch::Tensor& repetition_penalties) { // [num_seqs] - TORCH_CHECK(logits.is_contiguous()); - TORCH_CHECK(prompt_mask.is_contiguous()); - TORCH_CHECK(output_mask.is_contiguous()); - TORCH_CHECK(repetition_penalties.is_contiguous()); + torch::stable::Tensor& logits, // [num_seqs, vocab_size], in-place + const torch::stable::Tensor& prompt_mask, // [num_seqs, vocab_size] + const torch::stable::Tensor& output_mask, // [num_seqs, vocab_size] + const torch::stable::Tensor& repetition_penalties) { // [num_seqs] + STD_TORCH_CHECK(logits.is_contiguous()); + STD_TORCH_CHECK(prompt_mask.is_contiguous()); + STD_TORCH_CHECK(output_mask.is_contiguous()); + STD_TORCH_CHECK(repetition_penalties.is_contiguous()); int vocab_size = logits.size(-1); int num_seqs = logits.size(0); @@ -635,7 +633,7 @@ void apply_repetition_penalties_( // Get number of SMs on the current device int sms = 0; cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, - logits.get_device()); + logits.get_device_index()); // Compute tile_num and tile_size int tile_num = @@ -645,27 +643,29 @@ void apply_repetition_penalties_( // Each block handles one sequence and a tile of vocab dim3 grid(num_seqs, tile_num); dim3 block(std::min(tile_size, 1024)); - const at::cuda::OptionalCUDAGuard device_guard(device_of(logits)); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - VLLM_DISPATCH_FLOATING_TYPES( + const torch::stable::accelerator::DeviceGuard device_guard( + logits.get_device_index()); + const cudaStream_t stream = get_current_cuda_stream(); + VLLM_STABLE_DISPATCH_FLOATING_TYPES( logits.scalar_type(), "apply_repetition_penalties_kernel", [&] { vllm::apply_repetition_penalties_kernel <<>>( - logits.data_ptr(), prompt_mask.data_ptr(), - output_mask.data_ptr(), - repetition_penalties.data_ptr(), num_seqs, vocab_size, - tile_size); + logits.mutable_data_ptr(), + prompt_mask.const_data_ptr(), + output_mask.const_data_ptr(), + repetition_penalties.const_data_ptr(), num_seqs, + vocab_size, tile_size); }); } -void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n, - const torch::Tensor& seqLens, torch::Tensor& indices, - int64_t numRows, int64_t stride0, int64_t stride1, - int64_t topK) { +void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n, + const torch::stable::Tensor& seqLens, + torch::stable::Tensor& indices, int64_t numRows, + int64_t stride0, int64_t stride1, int64_t topK) { constexpr int kSortingAlgorithmThreshold = 12288; constexpr int kSplitWorkThreshold = 200 * 1000; constexpr int kNumThreadsPerBlock = 512; - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const cudaStream_t stream = get_current_cuda_stream(); const auto numColumns = logits.size(1); // True if seqLens is 2D (B, next_n): each logit row has its own pre-computed @@ -677,73 +677,76 @@ void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n, // Use insertion sort vllm::topKPerRowDecode <<>>( - logits.data_ptr(), seqLens.data_ptr(), - indices.data_ptr(), static_cast(stride0), + logits.const_data_ptr(), seqLens.const_data_ptr(), + indices.mutable_data_ptr(), static_cast(stride0), static_cast(stride1), static_cast(topK), static_cast(next_n), seqLensIs2D); } else if (numColumns < kSplitWorkThreshold) { // From this threshold, use radix sort instead vllm::topKPerRowDecode <<>>( - logits.data_ptr(), seqLens.data_ptr(), - indices.data_ptr(), static_cast(stride0), + logits.const_data_ptr(), seqLens.const_data_ptr(), + indices.mutable_data_ptr(), static_cast(stride0), static_cast(stride1), static_cast(topK), static_cast(next_n), seqLensIs2D); } else { // Long sequences are run in two steps constexpr auto multipleBlocksPerRowConfig = 10; - const auto outIndicesAux = - torch::empty({numRows, multipleBlocksPerRowConfig, topK}, - torch::dtype(torch::kInt32).device(logits.device())); - const auto outLogitsAux = - torch::empty({numRows, multipleBlocksPerRowConfig, topK}, - torch::dtype(torch::kFloat).device(logits.device())); + const auto outIndicesAux = torch::stable::empty( + {numRows, multipleBlocksPerRowConfig, topK}, + torch::headeronly::ScalarType::Int, std::nullopt, logits.device()); + const auto outLogitsAux = torch::stable::empty( + {numRows, multipleBlocksPerRowConfig, topK}, + torch::headeronly::ScalarType::Float, std::nullopt, logits.device()); vllm::topKPerRowDecode <<>>( - logits.data_ptr(), seqLens.data_ptr(), - outIndicesAux.data_ptr(), static_cast(stride0), + logits.const_data_ptr(), seqLens.const_data_ptr(), + outIndicesAux.mutable_data_ptr(), static_cast(stride0), static_cast(stride1), static_cast(topK), static_cast(next_n), seqLensIs2D, - outLogitsAux.data_ptr()); + outLogitsAux.mutable_data_ptr()); constexpr int kNumThreadsPerBlockMerge = 1024; vllm::topKPerRowDecode <<>>( - outLogitsAux.data_ptr(), seqLens.data_ptr(), - indices.data_ptr(), multipleBlocksPerRowConfig * topK, 1, - static_cast(topK), static_cast(next_n), seqLensIs2D, - nullptr, multipleBlocksPerRowConfig, outIndicesAux.data_ptr()); + outLogitsAux.const_data_ptr(), seqLens.const_data_ptr(), + indices.mutable_data_ptr(), multipleBlocksPerRowConfig * topK, + 1, static_cast(topK), static_cast(next_n), seqLensIs2D, + nullptr, multipleBlocksPerRowConfig, + outIndicesAux.const_data_ptr()); } } -void top_k_per_row_prefill(const torch::Tensor& logits, - const torch::Tensor& rowStarts, - const torch::Tensor& rowEnds, torch::Tensor& indices, - int64_t numRows, int64_t stride0, int64_t stride1, - int64_t topK) { +void top_k_per_row_prefill(const torch::stable::Tensor& logits, + const torch::stable::Tensor& rowStarts, + const torch::stable::Tensor& rowEnds, + torch::stable::Tensor& indices, int64_t numRows, + int64_t stride0, int64_t stride1, int64_t topK) { constexpr int kSortingAlgorithmThreshold = 12288; constexpr int kNumThreadsPerBlock = 512; - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const cudaStream_t stream = get_current_cuda_stream(); int numInsertionBlocks = std::min(static_cast(numRows), kSortingAlgorithmThreshold); vllm::topKPerRowPrefill <<>>(logits.data_ptr(), rowStarts.data_ptr(), - rowEnds.data_ptr(), indices.data_ptr(), - static_cast(stride0), static_cast(stride1), - static_cast(topK), 0); + stream>>>(logits.const_data_ptr(), + rowStarts.const_data_ptr(), + rowEnds.const_data_ptr(), + indices.mutable_data_ptr(), static_cast(stride0), + static_cast(stride1), static_cast(topK), 0); if (numRows > kSortingAlgorithmThreshold) { int numRadixBlocks = numRows - kSortingAlgorithmThreshold; vllm::topKPerRowPrefill <<>>(logits.data_ptr(), rowStarts.data_ptr(), - rowEnds.data_ptr(), indices.data_ptr(), - static_cast(stride0), static_cast(stride1), - static_cast(topK), kSortingAlgorithmThreshold); + stream>>>( + logits.const_data_ptr(), rowStarts.const_data_ptr(), + rowEnds.const_data_ptr(), indices.mutable_data_ptr(), + static_cast(stride0), static_cast(stride1), + static_cast(topK), kSortingAlgorithmThreshold); } } \ No newline at end of file diff --git a/csrc/topk.cu b/csrc/libtorch_stable/topk.cu similarity index 63% rename from csrc/topk.cu rename to csrc/libtorch_stable/topk.cu index 9ca9aa8824d8..15af18118f30 100644 --- a/csrc/topk.cu +++ b/csrc/libtorch_stable/topk.cu @@ -1,49 +1,51 @@ // Persistent TopK kernel for DeepSeek V3 sparse attention indexer. // See persistent_topk.cuh for kernel implementation. -#include -#include #include #include +#include "torch_utils.h" + #ifndef USE_ROCM - #include "persistent_topk.cuh" + #include "../persistent_topk.cuh" #endif namespace { #ifndef USE_ROCM template -void launch_persistent_topk(const torch::Tensor& logits, - const torch::Tensor& lengths, torch::Tensor& output, - torch::Tensor& workspace, int64_t max_seq_len) { +void launch_persistent_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, + int64_t max_seq_len) { namespace P = vllm::persistent; const int64_t num_rows = logits.size(0); const int64_t stride = logits.stride(0); - cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const cudaStream_t stream = get_current_cuda_stream(); static int num_sms = 0; static int max_smem_per_block = 0; if (num_sms == 0) { - int device; - cudaGetDevice(&device); - cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device); - cudaDeviceGetAttribute(&max_smem_per_block, - cudaDevAttrMaxSharedMemoryPerBlockOptin, device); + const cudaDeviceProp* device_prop = get_device_prop(); + num_sms = device_prop->multiProcessorCount; + max_smem_per_block = device_prop->sharedMemPerBlockOptin; } if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { cudaError_t status = vllm::FilteredTopKRaggedTransform( - logits.data_ptr(), output.data_ptr(), - lengths.data_ptr(), static_cast(num_rows), + logits.const_data_ptr(), output.mutable_data_ptr(), + lengths.const_data_ptr(), static_cast(num_rows), static_cast(TopK), static_cast(stride), stream); - TORCH_CHECK(status == cudaSuccess, - "FilteredTopK failed: ", cudaGetErrorString(status)); + STD_TORCH_CHECK(status == cudaSuccess, + "FilteredTopK failed: ", cudaGetErrorString(status)); } else { - TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor"); - TORCH_CHECK(workspace.dtype() == torch::kUInt8, "workspace must be uint8"); + STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor"); + STD_TORCH_CHECK( + workspace.scalar_type() == torch::headeronly::ScalarType::Byte, + "workspace must be uint8"); int effective_max_smem; if (num_rows <= 4) { @@ -99,9 +101,9 @@ void launch_persistent_topk(const torch::Tensor& logits, &occupancy, P::persistent_topk_kernel, P::kThreadsPerBlock, smem_size); } - TORCH_CHECK(occ_err == cudaSuccess, - "persistent_topk occupancy query failed: ", - cudaGetErrorString(occ_err)); + STD_TORCH_CHECK(occ_err == cudaSuccess, + "persistent_topk occupancy query failed: ", + cudaGetErrorString(occ_err)); if (occupancy < 1) occupancy = 1; // The cooperative spin-wait barrier only runs when at least one row hits @@ -131,27 +133,29 @@ void launch_persistent_topk(const torch::Tensor& logits, // If the cooperative launch wouldn't fit, fall back to FilteredTopK // instead of deadlocking. Only relevant when needs_cooperative. if (needs_cooperative && total_ctas > hw_resident_cap) { - TORCH_CHECK(max_smem_per_block >= 128 * 1024, - "persistent_topk would oversubscribe and the FilteredTopK " - "fallback requires >=128KB smem per block (have ", - max_smem_per_block, "). total_ctas=", total_ctas, - " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK, - ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group, - ", smem=", smem_size, ")."); + STD_TORCH_CHECK( + max_smem_per_block >= 128 * 1024, + "persistent_topk would oversubscribe and the FilteredTopK " + "fallback requires >=128KB smem per block (have ", + max_smem_per_block, "). total_ctas=", total_ctas, + " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK, + ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group, + ", smem=", smem_size, ")."); cudaError_t status = vllm::FilteredTopKRaggedTransform( - logits.data_ptr(), output.data_ptr(), - lengths.data_ptr(), static_cast(num_rows), - static_cast(TopK), static_cast(stride), - stream); - TORCH_CHECK(status == cudaSuccess, - "FilteredTopK fallback failed: ", cudaGetErrorString(status)); + logits.const_data_ptr(), + output.mutable_data_ptr(), + lengths.const_data_ptr(), + static_cast(num_rows), static_cast(TopK), + static_cast(stride), stream); + STD_TORCH_CHECK(status == cudaSuccess, "FilteredTopK fallback failed: ", + cudaGetErrorString(status)); return; } size_t state_bytes = num_groups * sizeof(P::RadixRowState); - TORCH_CHECK(workspace.size(0) >= static_cast(state_bytes), - "workspace too small, need ", state_bytes, " bytes"); + STD_TORCH_CHECK(workspace.size(0) >= static_cast(state_bytes), + "workspace too small, need ", state_bytes, " bytes"); // Zero the per-group RadixRowState region before launch. // @@ -179,22 +183,22 @@ void launch_persistent_topk(const torch::Tensor& logits, // first red_release. cudaMemsetAsync is stream-ordered: the zero // is globally visible before any CTA runs. { - cudaError_t mz_err = cudaMemsetAsync(workspace.data_ptr(), 0, - state_bytes, stream); - TORCH_CHECK(mz_err == cudaSuccess, - "row_states memset failed: ", cudaGetErrorString(mz_err)); + cudaError_t mz_err = cudaMemsetAsync( + workspace.mutable_data_ptr(), 0, state_bytes, stream); + STD_TORCH_CHECK(mz_err == cudaSuccess, + "row_states memset failed: ", cudaGetErrorString(mz_err)); } P::PersistentTopKParams params; - params.input = logits.data_ptr(); - params.output = output.data_ptr(); - params.lengths = lengths.data_ptr(); + params.input = logits.const_data_ptr(); + params.output = output.mutable_data_ptr(); + params.lengths = lengths.const_data_ptr(); params.num_rows = static_cast(num_rows); params.stride = static_cast(stride); params.top_k = static_cast(TopK); params.chunk_size = chunk_size; - params.row_states = - reinterpret_cast(workspace.data_ptr()); + params.row_states = reinterpret_cast( + workspace.mutable_data_ptr()); params.ctas_per_group = ctas_per_group; params.max_seq_len = static_cast(max_seq_len); @@ -203,8 +207,8 @@ void launch_persistent_topk(const torch::Tensor& logits, auto kernel = &P::persistent_topk_kernel; \ cudaError_t err = cudaFuncSetAttribute( \ kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \ - TORCH_CHECK(err == cudaSuccess, \ - "Failed to set smem: ", cudaGetErrorString(err)); \ + STD_TORCH_CHECK(err == cudaSuccess, \ + "Failed to set smem: ", cudaGetErrorString(err)); \ kernel<<>>(params); \ } while (0) @@ -219,37 +223,42 @@ void launch_persistent_topk(const torch::Tensor& logits, } cudaError_t err = cudaGetLastError(); - TORCH_CHECK(err == cudaSuccess, - "persistent_topk failed: ", cudaGetErrorString(err)); + STD_TORCH_CHECK(err == cudaSuccess, + "persistent_topk failed: ", cudaGetErrorString(err)); } #endif } // anonymous namespace -void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths, - torch::Tensor& output, torch::Tensor& workspace, int64_t k, +void persistent_topk(const torch::stable::Tensor& logits, + const torch::stable::Tensor& lengths, + torch::stable::Tensor& output, + torch::stable::Tensor& workspace, int64_t k, int64_t max_seq_len) { #ifndef USE_ROCM - TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor"); - TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor"); - TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor"); - TORCH_CHECK(logits.dtype() == torch::kFloat32, "Only float32 supported"); - TORCH_CHECK(lengths.dtype() == torch::kInt32, "lengths must be int32"); - TORCH_CHECK(output.dtype() == torch::kInt32, "output must be int32"); - TORCH_CHECK(logits.dim() == 2, "logits must be 2D"); - TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2, - "lengths must be 1D or 2D"); - TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous"); - TORCH_CHECK(output.dim() == 2, "output must be 2D"); + STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor"); + STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor"); + STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor"); + STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float, + "Only float32 supported"); + STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int, + "lengths must be int32"); + STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int, + "output must be int32"); + STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D"); + STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2, + "lengths must be 1D or 2D"); + STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous"); + STD_TORCH_CHECK(output.dim() == 2, "output must be 2D"); const int64_t num_rows = logits.size(0); - const int64_t stride = logits.stride(0); - TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch"); - TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k, - "output size mismatch"); - TORCH_CHECK(k == 512 || k == 1024 || k == 2048, - "persistent_topk supports k=512, k=1024, or k=2048, got k=", k); + STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch"); + STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k, + "output size mismatch"); + STD_TORCH_CHECK( + k == 512 || k == 1024 || k == 2048, + "persistent_topk supports k=512, k=1024, or k=2048, got k=", k); if (k == 512) { launch_persistent_topk<512>(logits, lengths, output, workspace, @@ -262,6 +271,6 @@ void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths, max_seq_len); } #else - TORCH_CHECK(false, "persistent_topk is not supported on ROCm"); + STD_TORCH_CHECK(false, "persistent_topk is not supported on ROCm"); #endif } diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index 1601c3bd5bfa..239341932307 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -263,6 +263,20 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "CUBLAS_M_THRESHOLD, bool has_zp, bool n32k16_reorder) -> Tensor"); #endif + // Merge attn states + // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 + // can be used to combine partial attention results (in the split-KV case) + ops.def( + "merge_attn_states(" + " Tensor! output," + " Tensor!? output_lse," + " Tensor prefix_output," + " Tensor prefix_lse," + " Tensor suffix_output," + " Tensor suffix_lse," + " int!? prefill_tokens_with_context," + " Tensor? output_scale=None) -> ()"); + // Hadamard transforms // conditionally compiled so impl registration is in source file ops.def("hadacore_transform(Tensor! x, bool inplace) -> Tensor"); @@ -319,6 +333,26 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "bool is_neox, Tensor position_ids, " "int forced_token_heads_per_warp=-1) -> ()"); + // Apply repetition penalties to logits in-place. + ops.def( + "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " + "Tensor output_mask, Tensor repetition_penalties) -> ()"); + + // Optimized top-k per row operations. + ops.def( + "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, " + "Tensor! indices, int numRows, int stride0, " + "int stride1, int topK) -> ()"); + + ops.def( + "top_k_per_row_decode(Tensor logits, int next_n, " + "Tensor seq_lens, Tensor! indices, " + "int numRows, int stride0, int stride1, int topK) -> ()"); + + ops.def( + "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " + "Tensor workspace, int k, int max_seq_len) -> ()"); + // Activation ops // Activation function used in SwiGLU. ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()"); @@ -422,6 +456,24 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "int type, SymInt row, SymInt tokens) -> Tensor"); ops.def("ggml_moe_get_block_size(int type) -> int"); + + // Mamba selective scan kernel + ops.def( + "selective_scan_fwd(Tensor! u, Tensor! delta," + "Tensor! A, Tensor! B, Tensor! C," + "Tensor? D_, Tensor!? z_, Tensor? delta_bias_," + "bool delta_softplus," + "Tensor? query_start_loc," + "Tensor? cache_indices," + "Tensor? has_initial_state," + "Tensor! ssm_states," + "int null_block_id," + "int block_size," + "Tensor? block_idx_first_scheduled_token," + "Tensor? block_idx_last_scheduled_token," + "Tensor? initial_state_idx," + "Tensor? cu_chunk_seqlen," + "Tensor? last_chunk_indices) -> ()"); } STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { @@ -469,6 +521,8 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { // files (allspark_repack.cu and allspark_qgemm_w8a16.cu) #endif + ops.impl("merge_attn_states", TORCH_BOX(&merge_attn_states)); + // Layernorm kernels (shared CUDA/ROCm) ops.impl("rms_norm", TORCH_BOX(&rms_norm)); ops.impl("fused_add_rms_norm", TORCH_BOX(&fused_add_rms_norm)); @@ -487,6 +541,13 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("rotary_embedding", TORCH_BOX(&rotary_embedding)); ops.impl("fused_qk_norm_rope", TORCH_BOX(&fused_qk_norm_rope)); + // Sampler kernels (shared CUDA/ROCm) + ops.impl("apply_repetition_penalties_", + TORCH_BOX(&apply_repetition_penalties_)); + ops.impl("top_k_per_row_prefill", TORCH_BOX(&top_k_per_row_prefill)); + ops.impl("top_k_per_row_decode", TORCH_BOX(&top_k_per_row_decode)); + ops.impl("persistent_topk", TORCH_BOX(&persistent_topk)); + // Activation kernels (shared CUDA/ROCm) ops.impl("silu_and_mul", TORCH_BOX(&silu_and_mul)); ops.impl("mul_and_silu", TORCH_BOX(&mul_and_silu)); @@ -519,6 +580,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("ggml_mul_mat_a8", TORCH_BOX(&ggml_mul_mat_a8)); ops.impl("ggml_moe_a8", TORCH_BOX(&ggml_moe_a8)); ops.impl("ggml_moe_a8_vec", TORCH_BOX(&ggml_moe_a8_vec)); + ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); } // These capability-check functions take only primitive args (no tensors), so diff --git a/csrc/ops.h b/csrc/ops.h index 3c6b2b7b9bc2..34da56fec4bc 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -54,13 +54,6 @@ void paged_attention_v2( const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, const int64_t blocksparse_head_sliding_step); -void merge_attn_states( - torch::Tensor& output, std::optional output_lse, - const torch::Tensor& prefix_output, const torch::Tensor& prefix_lse, - const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse, - const std::optional prefill_tokens_with_context, - const std::optional& output_scale = std::nullopt); - // rms_norm and fused_add_rms_norm declarations also exist in // csrc/libtorch_stable/ops.h (torch::stable ABI for CUDA). They remain here // because the CPU build still uses these torch::Tensor declarations. @@ -76,26 +69,6 @@ torch::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( torch::Tensor const& cos_sin_cache, int64_t q_head_padded, double eps, int64_t cache_block_size); -void apply_repetition_penalties_(torch::Tensor& logits, - const torch::Tensor& prompt_mask, - const torch::Tensor& output_mask, - const torch::Tensor& repetition_penalties); - -void top_k_per_row_prefill(const torch::Tensor& logits, - const torch::Tensor& rowStarts, - const torch::Tensor& rowEnds, torch::Tensor& indices, - int64_t numRows, int64_t stride0, int64_t stride1, - int64_t topK); - -void top_k_per_row_decode(const torch::Tensor& logits, int64_t next_n, - const torch::Tensor& seqLens, torch::Tensor& indices, - int64_t numRows, int64_t stride0, int64_t stride1, - int64_t topK); - -void persistent_topk(const torch::Tensor& logits, const torch::Tensor& lengths, - torch::Tensor& output, torch::Tensor& workspace, int64_t k, - int64_t max_seq_len); - void silu_and_mul_per_block_quant(torch::Tensor& out, torch::Tensor const& input, torch::Tensor& scales, int64_t group_size, @@ -150,22 +123,6 @@ void dynamic_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, torch::Tensor& scales, std::optional const& azp); -void selective_scan_fwd( - const torch::Tensor& u, const torch::Tensor& delta, const torch::Tensor& A, - const torch::Tensor& B, const torch::Tensor& C, - const std::optional& D_, - const std::optional& z_, - const std::optional& delta_bias_, bool delta_softplus, - const std::optional& query_start_loc, - const std::optional& cache_indices, - const std::optional& has_initial_state, - const torch::Tensor& ssm_states, int64_t null_block_id, int64_t block_size, - const std::optional& block_idx_first_scheduled_token, - const std::optional& block_idx_last_scheduled_token, - const std::optional& initial_state_idx, - const std::optional& cu_chunk_seqlen, - const std::optional& last_chunk_indices); - torch::Tensor dynamic_4bit_int_moe_cpu( torch::Tensor x, torch::Tensor topk_ids, torch::Tensor topk_weights, torch::Tensor w13_packed, torch::Tensor w2_packed, int64_t H, int64_t I, diff --git a/csrc/persistent_topk.cuh b/csrc/persistent_topk.cuh index 8b9d10ff83dd..6b25dc9940e8 100644 --- a/csrc/persistent_topk.cuh +++ b/csrc/persistent_topk.cuh @@ -126,10 +126,10 @@ struct RadixRowState { // ============================================================================ struct PersistentTopKParams { - const float* __restrict__ input; // [num_rows, stride] - int32_t* __restrict__ output; // [num_rows, top_k] - int32_t* __restrict__ lengths; // [num_rows] - RadixRowState* row_states; // large path: per-group state + const float* __restrict__ input; // [num_rows, stride] + int32_t* __restrict__ output; // [num_rows, top_k] + const int32_t* __restrict__ lengths; // [num_rows] + RadixRowState* row_states; // large path: per-group state uint32_t num_rows; uint32_t stride; uint32_t top_k; // actual k value for output stride @@ -1269,9 +1269,11 @@ constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) { } template -cudaError_t FilteredTopKRaggedTransform(DType* input, IdType* output_indices, - IdType* lengths, uint32_t num_rows, - uint32_t top_k_val, uint32_t max_len, +cudaError_t FilteredTopKRaggedTransform(const DType* input, + IdType* output_indices, + const IdType* lengths, + uint32_t num_rows, uint32_t top_k_val, + uint32_t max_len, cudaStream_t stream = 0) { constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC; constexpr int MAX_VEC = 16 / sizeof(DType); diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 487ff8df9dd0..a2a5d8097459 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -62,21 +62,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { " int blocksparse_head_sliding_step) -> ()"); ops.impl("paged_attention_v2", torch::kCUDA, &paged_attention_v2); - // Merge attn states - // Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005 - // can be used to combine partial attention results (in the split-KV case) - ops.def( - "merge_attn_states(" - " Tensor! output," - " Tensor!? output_lse," - " Tensor prefix_output," - " Tensor prefix_lse," - " Tensor suffix_output," - " Tensor suffix_lse," - " int!? prefill_tokens_with_context," - " Tensor? output_scale=None) -> ()"); - ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states); - // Activation ops (quantized only — basic ops moved to _C_stable_libtorch) ops.def( "silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()"); @@ -105,31 +90,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert", torch::kCUDA, &fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert); - // Apply repetition penalties to logits in-place - ops.def( - "apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, " - "Tensor output_mask, Tensor repetition_penalties) -> ()"); - ops.impl("apply_repetition_penalties_", torch::kCUDA, - &apply_repetition_penalties_); - - // Optimized top-k per row operation - ops.def( - "top_k_per_row_prefill(Tensor logits, Tensor rowStarts, Tensor rowEnds, " - "Tensor! indices, int numRows, int stride0, " - "int stride1, int topK) -> ()"); - ops.impl("top_k_per_row_prefill", torch::kCUDA, &top_k_per_row_prefill); - - ops.def( - "top_k_per_row_decode(Tensor logits, int next_n, " - "Tensor seq_lens, Tensor! indices, " - "int numRows, int stride0, int stride1, int topK) -> ()"); - ops.impl("top_k_per_row_decode", torch::kCUDA, &top_k_per_row_decode); - - ops.def( - "persistent_topk(Tensor logits, Tensor lengths, Tensor! output, " - "Tensor workspace, int k, int max_seq_len) -> ()"); - ops.impl("persistent_topk", torch::kCUDA, &persistent_topk); - // Quantization ops #ifndef USE_ROCM @@ -230,25 +190,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { #endif - // Mamba selective scan kernel - ops.def( - "selective_scan_fwd(Tensor! u, Tensor! delta," - "Tensor! A, Tensor! B, Tensor! C," - "Tensor? D_, Tensor!? z_, Tensor? delta_bias_," - "bool delta_softplus," - "Tensor? query_start_loc," - "Tensor? cache_indices," - "Tensor? has_initial_state," - "Tensor! ssm_states," - "int null_block_id," - "int block_size," - "Tensor? block_idx_first_scheduled_token," - "Tensor? block_idx_last_scheduled_token," - "Tensor? initial_state_idx," - "Tensor? cu_chunk_seqlen," - "Tensor? last_chunk_indices) -> ()"); - ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd); - #ifndef USE_ROCM ops.def( "minimax_allreduce_rms(" diff --git a/pyproject.toml b/pyproject.toml index 700134211532..c782cc326bc1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -141,6 +141,7 @@ bbc5b7ede = "bbc5b7ede" NOOPs = "NOOPs" nin_shortcut = "nin_shortcut" cudaDevAttrMaxSharedMemoryPerBlockOptin = "cudaDevAttrMaxSharedMemoryPerBlockOptin" +sharedMemPerBlockOptin = "sharedMemPerBlockOptin" depthwise_seperable_out_channel = "depthwise_seperable_out_channel" pard_token = "pard_token" From 206b72c982db29eb10e475589cbf3fa8d6a37071 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Wed, 27 May 2026 19:51:56 +0200 Subject: [PATCH 17/27] [Quantization] Fix Humming RoutedExperts import (#43540) Signed-off-by: Minh Vu Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- vllm/model_executor/layers/quantization/utils/humming_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/utils/humming_utils.py b/vllm/model_executor/layers/quantization/utils/humming_utils.py index 63735708b6c6..3c01977f3b54 100644 --- a/vllm/model_executor/layers/quantization/utils/humming_utils.py +++ b/vllm/model_executor/layers/quantization/utils/humming_utils.py @@ -8,11 +8,11 @@ from humming.schema import BaseWeightSchema from vllm import envs +from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, FusedMoEQuantDesc, ) -from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts from vllm.model_executor.layers.linear import LinearBase from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape From 2616f67faaa735a3e0d9c17968fa91f242d36c56 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 27 May 2026 20:46:36 +0100 Subject: [PATCH 18/27] Remove Transformers forward/backward compatibility tests (#43785) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- .buildkite/test_areas/models_basic.yaml | 34 ------------------------- 1 file changed, 34 deletions(-) diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index ef4638c6c4f4..4e47cbb77948 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -58,37 +58,3 @@ steps: device: cpu-small commands: - pytest -v -s models/test_utils.py models/test_vision.py - -- label: Transformers Nightly Models - device: h200_35gb - key: transformers-nightly-models - working_dir: "/vllm-workspace/" - optional: true - soft_fail: true - commands: - - pip install --upgrade git+https://github.com/huggingface/transformers - - pytest -v -s tests/models/test_initialization.py - - pytest -v -s tests/models/test_transformers.py - - pytest -v -s tests/models/multimodal/processing/ - - pytest -v -s tests/models/multimodal/test_mapping.py - - python3 examples/basic/offline_inference/chat.py - - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl - # Whisper needs spawn method to avoid deadlock - - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper - -- label: Transformers Backward Compatibility Models Test - device: h200_35gb - key: transformers-backward-compatibility-models-test - working_dir: "/vllm-workspace/" - optional: true - soft_fail: true - commands: - - pip install transformers==4.57.5 - - pytest -v -s tests/models/test_initialization.py - - pytest -v -s tests/models/test_transformers.py - - pytest -v -s tests/models/multimodal/processing/ - - pytest -v -s tests/models/multimodal/test_mapping.py - - python3 examples/basic/offline_inference/chat.py - - python3 examples/generate/multimodal/vision_language_offline.py --model-type qwen2_5_vl - # Whisper needs spawn method to avoid deadlock - - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/generate/multimodal/audio_language_offline.py --model-type whisper From 2c2c966669032e863f94919e9225aa12378c9364 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Wed, 27 May 2026 22:14:49 +0100 Subject: [PATCH 19/27] Validate against some config fields being set to 0 (#43794) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- vllm/config/cache.py | 24 ++++++++++++++++-------- vllm/config/model.py | 2 +- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 4fa7e1cfcbb9..352ccec3202f 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Callable from dataclasses import field -from typing import ClassVar, Literal +from typing import Any, ClassVar, Literal -from pydantic import Field, SkipValidation, field_validator, model_validator +from pydantic import Field, field_validator, model_validator from vllm.config.utils import config from vllm.logger import init_logger @@ -44,14 +45,14 @@ class CacheConfig: DEFAULT_BLOCK_SIZE: ClassVar[int] = 16 - block_size: SkipValidation[int] = None # type: ignore[assignment] + block_size: int = Field(default=None, gt=0) # type: ignore[assignment] """Size of a contiguous cache block in number of tokens. Accepts None (meaning "use default"). After construction, always int.""" user_specified_block_size: bool = field(default=False, init=False) """Whether block_size was explicitly provided. Derived automatically.""" user_specified_mamba_block_size: bool = field(default=False, init=False) """Whether mamba_block_size was explicitly provided. Derived automatically.""" - hash_block_size: SkipValidation[int] | None = None # type: ignore + hash_block_size: int | None = Field(default=None, gt=0) """Block size (in tokens) used for computing Request's block_hashes. This can be set to a finer granularity than the physical KV cache block @@ -220,19 +221,26 @@ def metrics_info(self): _block_size_resolved: bool = field(default=False, init=False) """Guard against pydantic re-running _apply_block_size_default.""" + @field_validator("block_size", mode="wrap") + @classmethod + def _skip_none_validation(cls, value: Any, handler: Callable) -> Any: + if value is None: + return value + return handler(value) + @model_validator(mode="after") def _apply_block_size_default(self) -> "CacheConfig": # Pydantic re-runs validators when CacheConfig is nested inside # another pydantic model (e.g. VllmConfig). Guard against that. if self._block_size_resolved: return self - object.__setattr__(self, "_block_size_resolved", True) + self._block_size_resolved = True if self.block_size is None: - object.__setattr__(self, "block_size", self.DEFAULT_BLOCK_SIZE) + self.block_size = self.DEFAULT_BLOCK_SIZE else: - object.__setattr__(self, "user_specified_block_size", True) + self.user_specified_block_size = True if self.mamba_block_size is not None: - object.__setattr__(self, "user_specified_mamba_block_size", True) + self.user_specified_mamba_block_size = True return self @field_validator("calculate_kv_scales", mode="after") diff --git a/vllm/config/model.py b/vllm/config/model.py index 811798859b84..491bd1a415b2 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -791,7 +791,7 @@ def validate_model_config_after(self: "ModelConfig") -> "ModelConfig": f"{type(self.tokenizer).__name__}: {self.tokenizer!r}. " "Please provide a valid tokenizer path or HuggingFace model ID." ) - if not isinstance(self.max_model_len, int): + if not isinstance(self.max_model_len, int) or self.max_model_len < 1: raise ValueError( f"max_model_len must be a positive integer, " f"got {type(self.max_model_len).__name__}: {self.max_model_len!r}. " From 7fb9c0197a3173f2a2edcc9d64f6c0e73ef20717 Mon Sep 17 00:00:00 2001 From: Benjamin Chislett Date: Wed, 27 May 2026 17:45:34 -0400 Subject: [PATCH 20/27] [Bugfix][DFlash]allocate the proper number of lookahead slots (#43733) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Benjamin Chislett Signed-off-by: Benjamin Chislett Co-authored-by: Nicolò Lucchesi --- vllm/v1/core/sched/scheduler.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 5911859c9d7a..6ce9dcc07563 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -218,6 +218,11 @@ def __init__( self.num_lookahead_tokens = self.num_spec_tokens if speculative_config.uses_draft_model(): self.num_lookahead_tokens = self.num_spec_tokens + if speculative_config.use_dflash(): + # DFlash requires an extra lookahead slot since it uses in-fill-style + # decoding instead of standard next-token sampling, so it has a query + # for the last sampled token plus queries for each draft token. + self.num_lookahead_tokens = self.num_spec_tokens + 1 # Create the KV cache manager. if hash_block_size is None: @@ -702,8 +707,9 @@ def schedule(self) -> SchedulerOutput: # extra block gets allocated which # creates a mismatch between the number # of local and remote blocks. + limit_lookahead_tokens = load_kv_async and self.use_eagle effective_lookahead_tokens = ( - 0 if request.num_computed_tokens == 0 else self.num_lookahead_tokens + 0 if limit_lookahead_tokens else self.num_lookahead_tokens ) # Determine if we need to allocate cross-attention blocks. From 5963c194787d30ed4a49c1e2e01010d8dffe1e79 Mon Sep 17 00:00:00 2001 From: Dakai An <77474977+andakai@users.noreply.github.com> Date: Thu, 28 May 2026 06:34:08 +0800 Subject: [PATCH 21/27] Fix Qwen3-VL and Qwen3-omni-thinker accuracy degradation from deepstack inputs under torch.compile (#43617) Signed-off-by: Dakai An --- .../models/qwen3_omni_moe_thinker.py | 25 +++++++++++-------- vllm/model_executor/models/qwen3_vl.py | 25 +++++++++++-------- 2 files changed, 28 insertions(+), 22 deletions(-) diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 6d5123effa5d..bd8a87b7dce6 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -1778,8 +1778,8 @@ def _get_deepstack_input_embeds( ) -> IntermediateTensors | None: if not getattr(self, "deepstack_input_embeds", None): return None # If vision tower is skipped - if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0: - return None + if num_tokens > self.deepstack_input_embeds[0].size(0): + self._resize_deepstack_input_embeds(num_tokens) # get deepstack_input_embeds from buffer, and clear the buffer return IntermediateTensors( @@ -1791,6 +1791,17 @@ def _get_deepstack_input_embeds( } ) + def _resize_deepstack_input_embeds(self, num_tokens: int) -> None: + self.deepstack_input_embeds = [ + torch.zeros( + num_tokens, + self.config.text_config.hidden_size, + device=self.deepstack_input_embeds[0].device, + dtype=self.deepstack_input_embeds[0].dtype, + ) + for _ in range(self.deepstack_num_level) + ] + def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> None: if not getattr(self, "deepstack_input_embeds", None): return @@ -1798,15 +1809,7 @@ def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> N # set deepstack_input_embeds to buffer num_tokens = deepstack_input_embeds.size(1) if num_tokens > self.deepstack_input_embeds[0].size(0): - self.deepstack_input_embeds = [ - torch.zeros( - num_tokens, - self.config.text_config.hidden_size, - device=self.deepstack_input_embeds[0].device, - dtype=self.deepstack_input_embeds[0].dtype, - ) - for _ in range(self.deepstack_num_level) - ] + self._resize_deepstack_input_embeds(num_tokens) for idx in range(self.deepstack_num_level): self.deepstack_input_embeds[idx][:num_tokens].copy_( deepstack_input_embeds[idx] diff --git a/vllm/model_executor/models/qwen3_vl.py b/vllm/model_executor/models/qwen3_vl.py index d7765351e6dd..a474649cc93e 100644 --- a/vllm/model_executor/models/qwen3_vl.py +++ b/vllm/model_executor/models/qwen3_vl.py @@ -1715,8 +1715,8 @@ def _get_deepstack_input_embeds( ) -> IntermediateTensors | None: if not getattr(self, "deepstack_input_embeds", None): return None # If vision tower is skipped - if getattr(self, "deepstack_input_embeds_num_tokens", 0) == 0: - return None + if num_tokens > self.deepstack_input_embeds[0].size(0): + self._resize_deepstack_input_embeds(num_tokens) # get deepstack_input_embeds from buffer, and clear the buffer return IntermediateTensors( @@ -1728,6 +1728,17 @@ def _get_deepstack_input_embeds( } ) + def _resize_deepstack_input_embeds(self, num_tokens: int) -> None: + self.deepstack_input_embeds = [ + torch.zeros( + num_tokens, + self.config.text_config.hidden_size, + device=self.deepstack_input_embeds[0].device, + dtype=self.deepstack_input_embeds[0].dtype, + ) + for _ in range(self.deepstack_num_level) + ] + def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> None: if not getattr(self, "deepstack_input_embeds", None): return @@ -1735,15 +1746,7 @@ def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> N # set deepstack_input_embeds to buffer num_tokens = deepstack_input_embeds.size(1) if num_tokens > self.deepstack_input_embeds[0].size(0): - self.deepstack_input_embeds = [ - torch.zeros( - num_tokens, - self.config.text_config.hidden_size, - device=self.deepstack_input_embeds[0].device, - dtype=self.deepstack_input_embeds[0].dtype, - ) - for _ in range(self.deepstack_num_level) - ] + self._resize_deepstack_input_embeds(num_tokens) for idx in range(self.deepstack_num_level): self.deepstack_input_embeds[idx][:num_tokens].copy_( deepstack_input_embeds[idx] From 094124af15d9ac85d1162cb318379e109c1a6bff Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 27 May 2026 18:14:50 -0500 Subject: [PATCH 22/27] Add @AndreasKaratzas to CODEOWNERS (#43740) Signed-off-by: Andreas Karatzas --- .github/CODEOWNERS | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index d65c736d4ae8..540540e51328 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -80,13 +80,13 @@ /tests/distributed/test_multi_node_assignment.py @youkaichao /tests/distributed/test_pipeline_parallel.py @youkaichao /tests/distributed/test_same_node.py @youkaichao -/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche -/tests/evals @mgoin @vadiklyutiy -/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye +/tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche @AndreasKaratzas +/tests/evals @mgoin @vadiklyutiy @AndreasKaratzas +/tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye @AndreasKaratzas /tests/kernels/ir @ProExpertProg @tjtanaa -/tests/models @DarkLight1337 @ywang96 +/tests/models @DarkLight1337 @ywang96 @AndreasKaratzas /tests/multimodal @DarkLight1337 @ywang96 @NickLucche -/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye +/tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety @zyongye @AndreasKaratzas /tests/test_inputs.py @DarkLight1337 @ywang96 /tests/entrypoints/llm/test_struct_output_generate.py @mgoin @russellb @aarnphm /tests/v1/structured_output @mgoin @russellb @aarnphm @@ -171,20 +171,20 @@ mkdocs.yaml @hmellor # ROCm related: specify owner with write access to notify AMD folks for careful code review /vllm/**/*rocm* @tjtanaa @dllehr-amd -/docker/Dockerfile.rocm* @tjtanaa @dllehr-amd +/docker/Dockerfile.rocm* @tjtanaa @dllehr-amd @AndreasKaratzas /vllm/v1/attention/backends/rocm*.py @tjtanaa @dllehr-amd /vllm/v1/attention/backends/mla/rocm*.py @tjtanaa @dllehr-amd /vllm/v1/attention/ops/rocm*.py @tjtanaa @dllehr-amd /vllm/model_executor/layers/fused_moe/rocm*.py @tjtanaa @dllehr-amd /csrc/rocm @tjtanaa @dllehr-amd -/requirements/*rocm* @tjtanaa -/tests/**/*rocm* @tjtanaa +/requirements/*rocm* @tjtanaa @AndreasKaratzas +/tests/**/*rocm* @tjtanaa @AndreasKaratzas /docs/**/*rocm* @tjtanaa /vllm/**/*quark* @tjtanaa -/tests/**/*quark* @tjtanaa +/tests/**/*quark* @tjtanaa @AndreasKaratzas /docs/**/*quark* @tjtanaa -/vllm/**/*aiter* @tjtanaa -/tests/**/*aiter* @tjtanaa +/vllm/**/*aiter* @tjtanaa @AndreasKaratzas +/tests/**/*aiter* @tjtanaa @AndreasKaratzas # TPU /vllm/v1/worker/tpu* @NickLucche From 381edde1b9bf6e4efce1ef2008d68ecdd47f415e Mon Sep 17 00:00:00 2001 From: amitz-nv <203509407+amitz-nv@users.noreply.github.com> Date: Thu, 28 May 2026 03:36:21 +0300 Subject: [PATCH 23/27] [Bugfix][Kernel] TRTLLM NVFP4 MoE chunking (#43599) Signed-off-by: amitz-nv <203509407+amitz-nv@users.noreply.github.com> --- .../fused_moe/experts/trtllm_bf16_moe.py | 3 - .../fused_moe/experts/trtllm_fp8_moe.py | 3 - .../fused_moe/experts/trtllm_mxfp4_moe.py | 3 - .../fused_moe/experts/trtllm_nvfp4_moe.py | 86 ++++++++++++++++--- 4 files changed, 74 insertions(+), 21 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py index 0b679b78c929..d11c7a3767d4 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py @@ -99,9 +99,6 @@ def _supports_router_logits_dtype( ) -> bool: return True - def supports_chunking(self) -> bool: - return False - def supports_expert_map(self) -> bool: return False diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index 31af4a32bae2..046d3e006af2 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -88,9 +88,6 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo or moe_parallel_config.use_ag_rs_all2all_kernels ) and not moe_parallel_config.enable_eplb - def supports_chunking(self) -> bool: - return False - def supports_expert_map(self) -> bool: return False diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index 64c163d73eb1..1e2fff8eb66c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -113,9 +113,6 @@ def _supports_activation(activation: MoEActivation) -> bool: def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard - def supports_chunking(self) -> bool: - return False - def supports_expert_map(self) -> bool: return False diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 3ddcdef77575..ee56b50acd28 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -157,8 +157,27 @@ def _supports_shape(hidden_dim: int) -> bool: def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard - def supports_chunking(self) -> bool: - return False + def _get_chunk_size(self) -> int: + MAX_GRID_Y = 65535 + MAX_TILE_TOKENS_DIM = 128 + + def _calc_max_supported_tokens(top_k: int, num_experts: int) -> int: + """Calculates the max number of supported tokens, so the CUDA grid.Y limit + won't be reached. + Based on getMaxNumCtasInBatchDim function in flashinfer's TRTLLM MoE runner: + https://github.com/flashinfer-ai/flashinfer/blob/719ee23fd82cb220d51ad118ca60198718f6c9d1/include/flashinfer/trtllm/fused_moe/runner.h#L97 + Which given numTokens, topK, numExperts, tileTokensDim calculates maxNumCtas + which is used as the CUDA grid.Y dimension, which we want to + be <= MAX_GRID_Y. Solving for numTokens gives the formula below. + """ + return ( + num_experts + (MAX_GRID_Y - num_experts + 1) * MAX_TILE_TOKENS_DIM - 1 + ) // top_k + + # Using 305k or more causes IMA error in the kernel, so limit to 300k. + return min( + 300000, _calc_max_supported_tokens(self.topk, self.moe_config.num_experts) + ) def supports_expert_map(self) -> bool: return False @@ -199,7 +218,7 @@ def workspace_shapes( def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: return TopKWeightAndReduceNoOP() - def apply( + def _invoke_kernel( self, output: torch.Tensor, hidden_states: torch.Tensor, @@ -209,18 +228,10 @@ def apply( topk_ids: torch.Tensor, activation: MoEActivation, global_num_experts: int, - expert_map: torch.Tensor | None, - a1q_scale: torch.Tensor | None, - a2_scale: torch.Tensor | None, - workspace13: torch.Tensor, - workspace2: torch.Tensor, - expert_tokens_meta: mk.ExpertTokensMetadata | None, - apply_router_weight_on_input: bool, + a1q_scale: torch.Tensor, ): import flashinfer - assert self._supports_activation(activation) - assert a1q_scale is not None assert self.quant_config.w1_scale is not None assert self.quant_config.w2_scale is not None @@ -262,6 +273,57 @@ def apply( output=output, ) + def apply( + self, + output: torch.Tensor, + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int, + expert_map: torch.Tensor | None, + a1q_scale: torch.Tensor | None, + a2_scale: torch.Tensor | None, + workspace13: torch.Tensor, + workspace2: torch.Tensor, + expert_tokens_meta: mk.ExpertTokensMetadata | None, + apply_router_weight_on_input: bool, + ): + assert self._supports_activation(activation) + assert a1q_scale is not None + + M = hidden_states.shape[0] + chunk_size = self._get_chunk_size() + + if chunk_size >= M: + self._invoke_kernel( + output, + hidden_states, + w1, + w2, + topk_weights, + topk_ids, + activation, + global_num_experts, + a1q_scale, + ) + else: + for start in range(0, M, chunk_size): + end = min(start + chunk_size, M) + self._invoke_kernel( + output[start:end], + hidden_states[start:end], + w1, + w2, + topk_weights[start:end], + topk_ids[start:end], + activation, + global_num_experts, + a1q_scale[start:end], + ) + class TrtLlmNvFp4ExpertsMonolithic( TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsMonolithic From 97c342af98bc938022a58829d6a741433d55aa3e Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Wed, 27 May 2026 19:41:49 -0500 Subject: [PATCH 24/27] add adaptation citation Signed-off-by: tjtanaa --- vllm/_tilelang_ops.py | 2 ++ vllm/platforms/cuda.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/vllm/_tilelang_ops.py b/vllm/_tilelang_ops.py index 5cc91a470a31..272cd256939a 100644 --- a/vllm/_tilelang_ops.py +++ b/vllm/_tilelang_ops.py @@ -24,6 +24,8 @@ tilelang = None # type: ignore[assignment] T = None # type: ignore[assignment] +# Conditions adapted from +# https://github.com/sgl-project/sglang/blob/0abe6a85a51f2b7f1c3ca0e8f78944b609b94344/python/sglang/srt/layers/mhc.py#L33 # noqa: E501 ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda() diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 58cef2ec976e..1a6a59535fb2 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -594,6 +594,8 @@ def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConf @classmethod def is_arch_support_pdl(cls) -> bool: + # Conditions adapted from + # https://github.com/sgl-project/sglang/blob/0abe6a85a51f2b7f1c3ca0e8f78944b609b94344/sgl-kernel/python/sgl_kernel/utils.py#L61 # noqa: E501 try: device = torch.cuda.current_device() major, _ = torch.cuda.get_device_capability(device) From 5bc61a45a1ed5a5f068c7e1246e1e0333ace2e98 Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Wed, 27 May 2026 19:47:01 -0500 Subject: [PATCH 25/27] remove unnecessary citation Signed-off-by: tjtanaa --- vllm/_tilelang_ops.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/_tilelang_ops.py b/vllm/_tilelang_ops.py index 272cd256939a..5cc91a470a31 100644 --- a/vllm/_tilelang_ops.py +++ b/vllm/_tilelang_ops.py @@ -24,8 +24,6 @@ tilelang = None # type: ignore[assignment] T = None # type: ignore[assignment] -# Conditions adapted from -# https://github.com/sgl-project/sglang/blob/0abe6a85a51f2b7f1c3ca0e8f78944b609b94344/python/sglang/srt/layers/mhc.py#L33 # noqa: E501 ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda() From 685bede28efa744f9c3644682c33284acd10ccf6 Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Wed, 27 May 2026 19:47:18 -0500 Subject: [PATCH 26/27] remove unnecessary citation Signed-off-by: tjtanaa --- vllm/platforms/cuda.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 1a6a59535fb2..58cef2ec976e 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -594,8 +594,6 @@ def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConf @classmethod def is_arch_support_pdl(cls) -> bool: - # Conditions adapted from - # https://github.com/sgl-project/sglang/blob/0abe6a85a51f2b7f1c3ca0e8f78944b609b94344/sgl-kernel/python/sgl_kernel/utils.py#L61 # noqa: E501 try: device = torch.cuda.current_device() major, _ = torch.cuda.get_device_capability(device) From 8f40bc10951c3637e399b3d79608a67b21f8dc7f Mon Sep 17 00:00:00 2001 From: tjtanaa Date: Wed, 27 May 2026 21:09:28 -0500 Subject: [PATCH 27/27] pin tilelang version rather than relax Signed-off-by: tjtanaa --- requirements/build/rocm.txt | 2 +- requirements/rocm.txt | 2 +- requirements/test/rocm.in | 1 + requirements/test/rocm.txt | 23 +++++++++++++++++++++-- 4 files changed, 24 insertions(+), 4 deletions(-) diff --git a/requirements/build/rocm.txt b/requirements/build/rocm.txt index 752fb1db786a..e09bdc078bf5 100644 --- a/requirements/build/rocm.txt +++ b/requirements/build/rocm.txt @@ -16,4 +16,4 @@ wheel jinja2>=3.1.6 amdsmi==7.0.2 timm>=1.0.17 -tilelang>=0.1.10 +tilelang==0.1.10 diff --git a/requirements/rocm.txt b/requirements/rocm.txt index 37bdf1afd15a..0520f4ca1e91 100644 --- a/requirements/rocm.txt +++ b/requirements/rocm.txt @@ -22,4 +22,4 @@ timm>=1.0.17 # amd-quark: required for Quark quantization on ROCm # To be consistent with test_quark.py amd-quark>=0.8.99 -tilelang>=0.1.10 +tilelang==0.1.10 diff --git a/requirements/test/rocm.in b/requirements/test/rocm.in index 812fb736b570..97e0658fb106 100644 --- a/requirements/test/rocm.in +++ b/requirements/test/rocm.in @@ -43,6 +43,7 @@ schemathesis>=3.39.15 # Required for openai schema test # quantization bitsandbytes==0.49.2 buildkite-test-collector==0.1.9 +tilelang==0.1.10 genai_perf>=0.0.8 tritonclient>=2.51.0 diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index b15e00edf1dd..c39f268709b5 100644 --- a/requirements/test/rocm.txt +++ b/requirements/test/rocm.txt @@ -43,7 +43,9 @@ anyio==4.13.0 # starlette # watchfiles apache-tvm-ffi==0.1.10 - # via xgrammar + # via + # tilelang + # xgrammar arctic-inference==0.1.1 # via -r requirements/test/rocm.in argcomplete==3.6.3 @@ -129,7 +131,9 @@ click==8.3.1 # typer # uvicorn cloudpickle==3.1.2 - # via -r requirements/test/../common.txt + # via + # -r requirements/test/../common.txt + # tilelang colorama==0.4.6 # via # perceptron @@ -511,6 +515,8 @@ mistral-common==1.11.2 # -c requirements/common.txt # -r requirements/test/../common.txt # -r requirements/test/rocm.in +ml-dtypes==0.5.4 + # via tilelang model-hosting-container-standards==0.1.14 # via # -c requirements/common.txt @@ -587,6 +593,7 @@ numpy==2.2.6 # lm-eval # matplotlib # mistral-common + # ml-dtypes # mteb # numba # opencv-python-headless @@ -610,6 +617,7 @@ numpy==2.2.6 # statsmodels # tensorizer # tifffile + # tilelang # torchvision # transformers # tritonclient @@ -811,6 +819,7 @@ psutil==7.2.2 # accelerate # peft # tensorizer + # tilelang py==1.11.0 # via pytest-forked py-cpuinfo==9.0.0 @@ -1192,6 +1201,10 @@ tiktoken==0.12.0 # gpt-oss # lm-eval # mistral-common +tilelang==0.1.10 + # via + # -c requirements/rocm.txt + # -r requirements/test/rocm.in timm==1.0.17 # via # -c requirements/rocm.txt @@ -1208,6 +1221,8 @@ tomli==2.4.0 # via schemathesis tomli-w==1.2.0 # via schemathesis +torch-c-dlpack-ext==0.1.5 + # via tilelang tqdm==4.67.3 # via # -r requirements/test/../common.txt @@ -1225,6 +1240,7 @@ tqdm==4.67.3 # pqdm # segmentation-models-pytorch # sentence-transformers + # tilelang # transformers transformers==5.5.3 # via @@ -1293,6 +1309,7 @@ typing-extensions==4.15.0 # sentence-transformers # sqlalchemy # starlette + # tilelang # torch # typeguard # typing-inspection @@ -1359,6 +1376,8 @@ yarl==1.23.0 # via # aiohttp # schemathesis +z3-solver==4.15.4.0 + # via tilelang zipp==3.23.0 # via importlib-metadata