diff --git a/tests/kernels/helion/test_silu_and_mul_per_block_quant.py b/tests/kernels/helion/test_silu_and_mul_per_block_quant.py new file mode 100644 index 00000000000..b8fcd9c8a67 --- /dev/null +++ b/tests/kernels/helion/test_silu_and_mul_per_block_quant.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the silu_and_mul_per_block_quant helion kernel +Run `pytest tests/kernels/helion/test_silu_and_mul_per_block_quant.py`. +""" + +from typing import Any + +import pytest +import torch +from torch._subclasses.fake_tensor import FakeTensorMode + +from tests.kernels.helion.utils import skip_if_platform_unsupported +from tests.kernels.quant_utils import FP8_DTYPE +from vllm.kernels.helion.case_key import CaseKey +from vllm.kernels.helion.config_manager import ConfigManager +from vllm.kernels.helion.ops.silu_and_mul_per_block_quant import ( + _pick_cache, + baseline, + pick_config, + silu_and_mul_per_block_quant, +) +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_helion +from vllm.utils.torch_utils import set_random_seed + +if not has_helion(): + pytest.skip( + "Helion is not installed. Install with: pip install vllm[helion]", + allow_module_level=True, + ) + + +def _generate_fake_input( + num_tokens: int, intermediate_size: int, group_size: int +) -> tuple[Any, ...]: + with FakeTensorMode(): + in_dtype: torch.dtype = torch.bfloat16 + out_dtype: torch.dtype = current_platform.fp8_dtype() + scale_dtype: torch.dtype = torch.float32 + input = torch.randn( + num_tokens, 2 * intermediate_size, device="cuda", dtype=in_dtype + ) + result = torch.empty( + num_tokens, intermediate_size, device=input.device, dtype=out_dtype + ) + scale = torch.empty( + (num_tokens, intermediate_size // group_size), + device=input.device, + dtype=scale_dtype, + ) + scale_ub = torch.mean(input).to(scale_dtype) + args = ( + result, + input, + scale, + group_size, + scale_ub, + False, + ) + return args + + +class TestSiluAndMulPerBlockQuantConfigPicker: + def setup_method(self): + _pick_cache.clear() + + def test_config_picker_exact_match(self): + config_keys = [ + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}), + ] + + args = _generate_fake_input(16, 4096, 128) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey( + {"intermediate_size": 4096, "group_size": 128, "num_tokens": 16} + ) + + def test_config_picker_closest_match(self): + config_keys = [ + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}), + CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}), + CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}), + CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}), + CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}), + CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}), + ] + + args = _generate_fake_input(20, 3000, 70) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey( + {"intermediate_size": 2048, "group_size": 64, "num_tokens": 32} + ) + + def test_config_picker_no_configs(self): + config_keys: list[dict] = [] + + args = _generate_fake_input(16, 4096, 128) + selected_key = pick_config(args, config_keys) + assert selected_key is None + + def test_config_picker_fallback_to_largest(self): + config_keys = [ + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}), + CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}), + CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}), + CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}), + CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}), + CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}), + CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}), + ] + + args = _generate_fake_input(64, 8192, 256) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey( + {"intermediate_size": 4096, "group_size": 128, "num_tokens": 32} + ) + + +@pytest.fixture(autouse=True) +def reset_config_manager_singleton(): + ConfigManager.reset_instance() + ConfigManager() + yield + ConfigManager.reset_instance() + + +class TestSiluAndMulPerBlockQuantCorrectness: + @pytest.mark.parametrize("num_tokens", [1, 7, 4096]) + @pytest.mark.parametrize("hidden_size", [1024, 2048, 5120]) + @pytest.mark.parametrize("group_size", [64, 128]) + @pytest.mark.parametrize("is_scale_transposed", [False, True]) + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + @pytest.mark.parametrize("quant_dtype", [current_platform.fp8_dtype(), torch.int8]) + @pytest.mark.parametrize("has_scale_ub", [True, False]) + @pytest.mark.parametrize("seed", [0]) + def test_silu_and_mul_per_block_quant( + self, + num_tokens: int, + hidden_size: int, + group_size: int, + is_scale_transposed: bool, + dtype: torch.dtype, + quant_dtype: torch.dtype, + has_scale_ub: bool, + seed: int, + ) -> None: + skip_if_platform_unsupported("silu_and_mul_per_block_quant") + set_random_seed(seed) + + if hidden_size % group_size != 0: + return + + if has_scale_ub and quant_dtype != FP8_DTYPE: + # skip + return + + scale = 1 / hidden_size + x = torch.randn(num_tokens, 2 * hidden_size, dtype=dtype, device="cuda") * scale + + if has_scale_ub: + act = torch.nn.functional.silu(x[:, :hidden_size]) * x[:, hidden_size:] + act_abs = act.abs().float() + scale_ub = 0.5 * (act_abs.mean() + act_abs.amax()) + else: + scale_ub = None + + ref_out = torch.empty(num_tokens, hidden_size, device="cuda", dtype=quant_dtype) + + if is_scale_transposed: + ref_scales = torch.empty( + (hidden_size // group_size, x.shape[0]), + device="cuda", + dtype=torch.float32, + ).t() + else: + ref_scales = torch.empty( + (x.shape[0], hidden_size // group_size), + device="cuda", + dtype=torch.float32, + ) + + ops_out = ref_out.clone() + ops_scales = ref_scales.clone() + + baseline(ref_out, x, ref_scales, group_size, scale_ub, is_scale_transposed) + silu_and_mul_per_block_quant( + ops_out, x, ops_scales, group_size, scale_ub, is_scale_transposed + ) + + torch.testing.assert_close(ref_scales, ops_scales) + # allow 1 ULP difference + assert ( + ref_out.view(torch.uint8).to(torch.int16) + - ops_out.view(torch.uint8).to(torch.int16) + ).abs().max() <= 1 + + +class TestSiluAndMulPerBlockQuantIntegration: + def test_kernel_registration_integration(self): + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + assert "silu_and_mul_per_block_quant" in registered_kernels + + kernel_wrapper = registered_kernels["silu_and_mul_per_block_quant"] + assert kernel_wrapper.op_name == "silu_and_mul_per_block_quant" + assert kernel_wrapper._config_picker is not None + assert kernel_wrapper._mutates_args == ["out", "scales"] + + def test_fake_impl_functionality(self): + skip_if_platform_unsupported("silu_and_mul_per_block_quant") + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + kernel_wrapper = registered_kernels["silu_and_mul_per_block_quant"] + fake_impl = kernel_wrapper._fake_impl + + args = _generate_fake_input(16, 4096, 128) + assert fake_impl(*args) is None diff --git a/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_b200.json b/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_b200.json new file mode 100644 index 00000000000..b50cabb0da9 --- /dev/null +++ b/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_b200.json @@ -0,0 +1,2099 @@ +[ + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 0, + 2 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 4 + ], + "range_warp_specializes": [ + null + ], + "range_multi_buffers": [ + false + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 4, + "maxnreg": 128 + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "first" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "last" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 2, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [ + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_h100.json b/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_h100.json new file mode 100644 index 00000000000..ca411c69470 --- /dev/null +++ b/vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_h100.json @@ -0,0 +1,2207 @@ +[ + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 0, + 1, + 2 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 8, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "last", + "last" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "last", + "first" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 0, + 2, + 1 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 2 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 4 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 1 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8 + ], + "loop_orders": [ + [ + 2, + 1, + 0 + ] + ], + "l2_groupings": [ + 8 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 6144, + "group_size": 128, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 64 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "", + "first", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 12288, + "group_size": 128, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 4 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 32 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "first", + "", + "last" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "intermediate_size": 25600, + "group_size": 128, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 16 + ], + "loop_orders": [ + [ + 1, + 2, + 0 + ] + ], + "l2_groupings": [ + 16 + ], + "range_unroll_factors": [ + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null + ], + "range_flattens": [ + null + ], + "load_eviction_policies": [ + "last", + "", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py new file mode 100644 index 00000000000..f3aaf226b04 --- /dev/null +++ b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from itertools import product +from typing import Any + +import torch + +from vllm.kernels.helion.case_key import CaseKey +from vllm.kernels.helion.utils import ( + get_fp8_dtype, + get_int8_min_max, + get_int8_min_scaling_factor, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + get_fp8_min_max, +) +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_helion + +if not has_helion(): + raise ImportError( + "Helion kernel requires helion to be installed. " + "Install it with: pip install helion" + ) + +import helion +import helion.language as hl + +from vllm.kernels.helion.register import register_kernel + +logger = init_logger(__name__) + + +def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: + # TODO(xiaohongchen1991): it is difficult for kernel author to cover all input + # property combination. Currently, dtypes are fixed. We need optimization to + # bucket/skip some combinations + num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] + intermediate_size_list = [6144, 12288, 25600] + + in_dtype: torch.dtype = torch.bfloat16 + out_dtype: torch.dtype = current_platform.fp8_dtype() + scale_dtype: torch.dtype = torch.float32 + group_size_list = [128] + inputs = {} + for intermediate_size, group_size, num_tokens in product( + intermediate_size_list, group_size_list, num_tokens_list + ): + input = torch.randn( + num_tokens, 2 * intermediate_size, device="cuda", dtype=in_dtype + ) + result = torch.empty( + num_tokens, intermediate_size, device=input.device, dtype=out_dtype + ) + scale = torch.empty( + (num_tokens, intermediate_size // group_size), + device=input.device, + dtype=scale_dtype, + ) + scale_ub = torch.mean(input).to(scale_dtype) + + config_key = CaseKey( + { + "intermediate_size": intermediate_size, + "group_size": group_size, + "num_tokens": num_tokens, + } + ) + inputs[config_key] = (result, input, scale, group_size, scale_ub, False) + + return inputs + + +_pick_cache: dict[tuple[int, int, int], CaseKey | None] = {} + + +def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None: + """Pick the best pre-tuned config for the given input shape. + + Selection strategy: + 1. Find the closest intermediate_size among available configs + (exact match preferred). + 2. Find the closest group_size among available configs + (exact match preferred). + 3. Among the num_tokens values tuned for that intermediate_size and group_size, + pick the smallest num_tokens >= the input's num_tokens. If the input is + larger than all available num_tokens, fall back to the largest. + """ + + if not config_keys: + return None + + result, _, _, group_size, *_ = args + num_tokens, intermediate_size = result.shape + + cache_key = (num_tokens, group_size, intermediate_size) + cached = _pick_cache.get(cache_key) + if cached is not None: + return cached + + configs: dict[int, dict[int, list[int]]] = {} + for key in config_keys: + if key.is_default(): + continue + configs.setdefault(key["intermediate_size"], {}).setdefault( + key["group_size"], [] + ).append(key["num_tokens"]) + + if not configs: + return None + + best_intermediate_size = min(configs, key=lambda s: abs(s - intermediate_size)) + best_group_size = min( + configs[best_intermediate_size], key=lambda s: abs(s - group_size) + ) + available_num_tokens = sorted(configs[best_intermediate_size][best_group_size]) + best_num_tokens = next( + (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1] + ) + + result = CaseKey( + { + "intermediate_size": best_intermediate_size, + "group_size": best_group_size, + "num_tokens": best_num_tokens, + } + ) + _pick_cache[cache_key] = result + return result + + +def fake_impl( + out: torch.Tensor, # [num_tokens, intermediate_size] + input: torch.Tensor, # [num_tokens, 2 * intermediate_size] + scales: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + scale_ub: torch.Tensor | None = None, # scalar tensor + is_scale_transposed: bool = False, +) -> None: + return + + +def baseline( + out: torch.Tensor, # [num_tokens, intermediate_size] + input: torch.Tensor, # [num_tokens, 2 * intermediate_size] + scales: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + scale_ub: torch.Tensor | None = None, # scalar tensor + is_scale_transposed: bool = False, +) -> None: + torch.ops._C.silu_and_mul_per_block_quant( + out, input, scales, group_size, scale_ub, is_scale_transposed + ) + + +@register_kernel( + mutates_args=["out", "scales"], + config_picker=pick_config, + input_generator=generate_inputs, + fake_impl=fake_impl, + helion_settings=helion.Settings( + autotune_baseline_fn=baseline, + ignore_warnings=[helion.exc.TensorOperationInWrapper], + ), +) # type: ignore[misc] +def silu_and_mul_per_block_quant( + out: torch.Tensor, # [num_tokens, intermediate_size] + input: torch.Tensor, # [num_tokens, 2 * intermediate_size] + scales: torch.Tensor, # [num_tokens, groups_per_row] + group_size: int, + scale_ub: torch.Tensor | None = None, # scalar tensor + is_scale_transposed: bool = False, # dummy +) -> None: + # This code assumes batch_dim and num_tokens are flattened + assert input.ndim == 2 + num_tokens, two_intermediate_size = input.shape + hl.specialize(two_intermediate_size) + + assert two_intermediate_size % 2 == 0 + intermediate_size = two_intermediate_size // 2 + + assert out.shape[0] == num_tokens + assert out.shape[1] == intermediate_size + fp8_dtype = get_fp8_dtype() + assert out.dtype in [fp8_dtype, torch.int8] + + if scale_ub is not None: + assert out.dtype == fp8_dtype + assert scale_ub.dtype == torch.float32 + + assert scales.ndim == 2 and scales.dtype == torch.float32 + + assert scales.shape[0] == num_tokens + groups_per_row = scales.shape[1] + hl.specialize(groups_per_row) + assert ( + intermediate_size % group_size == 0 + and intermediate_size // group_size == groups_per_row + ) + + assert group_size in [64, 128] + hl.specialize(group_size) + + assert input.stride()[-1] == 1 + assert out.stride()[-1] == 1 + + quant_dtype = out.dtype + qtype_traits_min: int | float + qtype_traits_max: int | float + if quant_dtype == torch.int8: + qtype_traits_min, qtype_traits_max = get_int8_min_max() + min_scaling_factor = get_int8_min_scaling_factor() + else: + qtype_traits_min, qtype_traits_max = get_fp8_min_max() + min_scaling_factor = 1.0 / (qtype_traits_max * 512.0) + + qtype_max = float(qtype_traits_max) + + input = input.view(num_tokens, -1, group_size) + out = out.view(num_tokens, -1, group_size) + + for tile_m, tile_gn, tile_n in hl.tile( + [num_tokens, groups_per_row, group_size], block_size=[1, None, group_size] + ): + x_a_blk = input[tile_m, tile_gn, tile_n].to(torch.float32) + x_b_blk = hl.load( + input, + [tile_m, tile_gn.index + groups_per_row, tile_n], + extra_mask=(tile_gn.index + groups_per_row < 2 * groups_per_row)[ + None, :, None + ], + ).to(torch.float32) + x_blk = x_a_blk * torch.sigmoid(x_a_blk) * x_b_blk + s_blk = torch.amax(torch.abs(x_blk), dim=-1).to(torch.float32) + + if scale_ub is not None: + scale_ub_s = hl.load(scale_ub, []) + s_blk = s_blk.clamp(max=scale_ub_s) + s_blk = s_blk * (1.0 / qtype_max) + s_blk = s_blk.clamp(min=min_scaling_factor) + + scales[tile_m, tile_gn] = s_blk + if quant_dtype == torch.int8: + y_blk = (x_blk * (1.0 / s_blk[:, :, None])).round() + else: + y_blk = x_blk / s_blk[:, :, None] + + out[tile_m, tile_gn, tile_n] = y_blk.clamp( + qtype_traits_min, qtype_traits_max + ).to(out.dtype)