From ce3ef17bec12cd3fc4fbecf2902ea54f22d3a953 Mon Sep 17 00:00:00 2001 From: "Xiaohong (Sean) Chen" Date: Tue, 16 Jun 2026 10:09:52 -0400 Subject: [PATCH] [Kernel][Helion][1/N] Add Helion kernel for rms_norm_per_block_quant (#36895) Signed-off-by: Sean Chen Co-authored-by: Yanan Cao --- .../helion/test_rms_norm_per_block_quant.py | 298 + .../rms_norm_per_block_quant/nvidia_b200.json | 2944 ++++++++++ .../rms_norm_per_block_quant/nvidia_h100.json | 5017 +++++++++++++++++ .../helion/ops/rms_norm_per_block_quant.py | 299 + vllm/kernels/helion/utils.py | 15 + 5 files changed, 8573 insertions(+) create mode 100644 tests/kernels/helion/test_rms_norm_per_block_quant.py create mode 100644 vllm/kernels/helion/configs/rms_norm_per_block_quant/nvidia_b200.json create mode 100644 vllm/kernels/helion/configs/rms_norm_per_block_quant/nvidia_h100.json create mode 100644 vllm/kernels/helion/ops/rms_norm_per_block_quant.py diff --git a/tests/kernels/helion/test_rms_norm_per_block_quant.py b/tests/kernels/helion/test_rms_norm_per_block_quant.py new file mode 100644 index 00000000000..4cb18d77598 --- /dev/null +++ b/tests/kernels/helion/test_rms_norm_per_block_quant.py @@ -0,0 +1,298 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the rms_norm_per_block_quant helion kernel + +Run `pytest tests/kernels/helion/test_rms_norm_per_block_quant.py`. +""" + +import itertools +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.rms_norm_per_block_quant import ( + _pick_cache, + baseline, + pick_config, + rms_norm_per_block_quant, +) +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, hidden_size: int, group_size: int +) -> tuple[Any, ...]: + with FakeTensorMode(): + input = torch.randn( + (num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16 + ) + result = torch.empty(input.shape, device=input.device, dtype=FP8_DTYPE) + scale = torch.empty( + (num_tokens, hidden_size // group_size), + device=input.device, + dtype=torch.float32, + ) + scale_ub = torch.mean(input).to(scale.dtype) + residual = torch.randn_like(input) + weight = torch.normal( + mean=1.0, + std=1.0, + size=(hidden_size,), + dtype=input.dtype, + device=input.device, + ) + epsilon = 1e-6 + args = ( + result, + input, + weight, + scale, + epsilon, + scale_ub, + residual, + group_size, + False, + ) + return args + + +@pytest.fixture(autouse=True) +def reset_config_manager_singleton(): + ConfigManager.reset_instance() + ConfigManager() + yield + ConfigManager.reset_instance() + + +class TestRmsNormPerBlockQuantConfigPicker: + def setup_method(self): + _pick_cache.clear() + + def test_config_picker_exact_match(self): + config_keys = [ + CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"hidden_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( + {"hidden_size": 4096, "group_size": 128, "num_tokens": 16} + ) + + def test_config_picker_closest_match(self): + config_keys = [ + CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}), + CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}), + CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}), + CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}), + CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}), + CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}), + CaseKey({"hidden_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( + {"hidden_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({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}), + CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}), + CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}), + CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}), + CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}), + CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}), + CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}), + CaseKey({"hidden_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( + {"hidden_size": 4096, "group_size": 128, "num_tokens": 32} + ) + + +DTYPES = [torch.bfloat16, torch.float] +QUANT_DTYPES = [torch.int8, FP8_DTYPE] +VEC_HIDDEN_SIZES = [64, 1024] +# Avoid combinatorial explosion with full Cartesian product +NUM_TOKENS_HIDDEN_SIZES = [ + *[(1, i) for i in [64, 128, 1024, 5120]], + *[(2048, i) for i in [64, 1024]], + *[(4096, i) for i in [64]], +] + +ADD_RESIDUAL = [False, True] +SCALE_UBS = [True, False] +GROUP_SIZES = [64, 128] +TMA_ALIGNMENTS = [0, 4] +SEEDS = [0] +EPS = 1e-6 + + +class TestRmsNormPerBlockQuantCorrectness: + @pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES) + @pytest.mark.parametrize("add_residual", ADD_RESIDUAL) + @pytest.mark.parametrize("has_scale_ub", SCALE_UBS) + @pytest.mark.parametrize("dtype", DTYPES) + @pytest.mark.parametrize("quant_dtype", QUANT_DTYPES) + @pytest.mark.parametrize("is_scale_transposed", [False, True]) + @pytest.mark.parametrize( + "group_size, tma_alignment", + [*itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)], + ) + @pytest.mark.parametrize("seed", SEEDS) + def test_rms_norm_per_block_quant( + self, + num_tokens: int, + hidden_size: int, + add_residual: bool, + has_scale_ub: bool, + dtype: torch.dtype, + quant_dtype: torch.dtype, + is_scale_transposed: bool, + group_size: int, + tma_alignment: int, + seed: int, + ) -> None: + skip_if_platform_unsupported("rms_norm_per_block_quant") + + set_random_seed(seed) + + if hidden_size % group_size != 0: + # skip + return + + if tma_alignment != 0 and hidden_size // group_size % tma_alignment == 0: + # Skip tests where TMA alignment doesn't create extra padding to save time + return + + if has_scale_ub and quant_dtype != FP8_DTYPE: + # skip + return + + scale = 1 / (hidden_size) + input = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda") * scale + weight = torch.normal( + mean=1.0, std=1.0, size=(hidden_size,), dtype=dtype, device=input.device + ) + residual = torch.randn_like(input) * scale if add_residual else None + scale_ub = ( + torch.mean(input).to(dtype=torch.float32, device="cuda") + if has_scale_ub + else None + ) + groups_per_row = hidden_size // group_size + + ref_residual = residual.clone() if residual is not None else None + ops_residual = residual.clone() if residual is not None else None + ref_out = torch.empty(input.shape, device=input.device, dtype=quant_dtype) + ops_out = ref_out.clone() + + if is_scale_transposed: + if tma_alignment == 0: + ref_scales = torch.empty( + (groups_per_row, num_tokens), + device=input.device, + dtype=torch.float32, + ).transpose(0, 1) + else: + tma_aligned_m = ( + (num_tokens + tma_alignment - 1) // tma_alignment * tma_alignment + ) + shape = (num_tokens, groups_per_row) + stride = (1, tma_aligned_m) + ref_scales = torch.empty_strided( + shape, stride, device=input.device, dtype=torch.float32 + ) + else: + ref_scales = torch.empty( + (num_tokens, groups_per_row), device=input.device, dtype=torch.float32 + ) + + ops_scales = ref_scales.clone() + + baseline( + ref_out, + input, + weight, + ref_scales, + EPS, + scale_ub, + ref_residual, + group_size, + is_scale_transposed, + ) + ref_scales = ref_scales.contiguous() + + rms_norm_per_block_quant( + ops_out, + input, + weight, + ops_scales, + EPS, + scale_ub, + ops_residual, + group_size, + is_scale_transposed, + ) + ops_scales = ops_scales.contiguous() + + 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 + + if add_residual: + torch.testing.assert_close(ref_residual, ops_residual) + + +class TestRmsNormPerBlockQuantIntegration: + def test_kernel_registration_integration(self): + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + assert "rms_norm_per_block_quant" in registered_kernels + + kernel_wrapper = registered_kernels["rms_norm_per_block_quant"] + assert kernel_wrapper.op_name == "rms_norm_per_block_quant" + assert kernel_wrapper._config_picker is not None + assert kernel_wrapper._mutates_args == ["result", "scale", "residual"] + + def test_fake_impl_functionality(self): + skip_if_platform_unsupported("rms_norm_per_block_quant") + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + kernel_wrapper = registered_kernels["rms_norm_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/rms_norm_per_block_quant/nvidia_b200.json b/vllm/kernels/helion/configs/rms_norm_per_block_quant/nvidia_b200.json new file mode 100644 index 00000000000..6acf8f29dab --- /dev/null +++ b/vllm/kernels/helion/configs/rms_norm_per_block_quant/nvidia_b200.json @@ -0,0 +1,2944 @@ +[ + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 0, + 0 + ], + "range_warp_specializes": [ + null, + false, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false, + null + ], + "range_flattens": [ + null, + true, + null, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "last", + "", + "", + "last", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 0, + 0 + ], + "range_warp_specializes": [ + null, + false, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false, + null + ], + "range_flattens": [ + null, + true, + null, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "last", + "", + "", + "last", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 0, + 0 + ], + "range_warp_specializes": [ + null, + false, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false, + null + ], + "range_flattens": [ + null, + true, + null, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "last", + "", + "", + "last", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 0, + 1, + 0 + ], + "range_warp_specializes": [ + null, + true, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null, + null + ], + "range_flattens": [ + null, + true, + false, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "first", + "", + "first", + "first", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 3, + 4, + 2 + ], + "range_warp_specializes": [ + null, + null, + null, + false + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false, + null + ], + "range_flattens": [ + null, + true, + null, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "last", + "first", + "last", + "last", + "", + "last" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 3, + 2, + 0 + ], + "range_warp_specializes": [ + null, + null, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null, + null + ], + "range_flattens": [ + null, + false, + null, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "last", + "first", + "last", + "last", + "", + "" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 0, + 0 + ], + "range_warp_specializes": [ + null, + false, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false, + null + ], + "range_flattens": [ + null, + true, + null, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "last", + "", + "", + "last", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 0, + 0 + ], + "range_warp_specializes": [ + null, + false, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false, + null + ], + "range_flattens": [ + null, + true, + null, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "last", + "", + "", + "last", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 1, + 0, + 2 + ], + "range_warp_specializes": [ + null, + null, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false, + true + ], + "range_flattens": [ + null, + true, + true, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "", + "last", + "first", + "", + "first", + "" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 2048, + 8 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 3, + 0, + 0 + ], + "range_warp_specializes": [ + null, + null, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null, + null + ], + "range_flattens": [ + null, + true, + null, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "last", + "first", + "last", + "last", + "", + "" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 3, + 0, + 0 + ], + "range_warp_specializes": [ + null, + null, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true, + null + ], + "range_flattens": [ + null, + true, + null, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "last", + "first", + "last", + "last", + "", + "" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 4, + 0 + ], + "range_warp_specializes": [ + null, + false, + false, + false + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null, + true + ], + "range_flattens": [ + null, + true, + true, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "", + "", + "", + "last", + "last", + "" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 4, + 4, + 0 + ], + "range_warp_specializes": [ + null, + false, + false, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null, + null + ], + "range_flattens": [ + null, + null, + true, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "", + "", + "", + "first", + "first", + "" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 2048, + 8 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 4, + 4, + 2 + ], + "range_warp_specializes": [ + null, + false, + false, + false + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true, + true + ], + "range_flattens": [ + null, + false, + false, + null + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "", + "", + "", + "first", + "", + "" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 3, + 0 + ], + "range_warp_specializes": [ + null, + false, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null, + null + ], + "range_flattens": [ + null, + null, + false, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "", + "" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 3, + 0 + ], + "range_warp_specializes": [ + null, + false, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null, + null + ], + "range_flattens": [ + null, + null, + false, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "", + "" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 0, + 1, + 4 + ], + "range_warp_specializes": [ + null, + null, + false, + false + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null, + null + ], + "range_flattens": [ + null, + null, + null, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "first", + "", + "first", + "first", + "", + "" + ], + "num_warps": 32, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 1, + 0 + ], + "range_warp_specializes": [ + null, + false, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false, + null + ], + "range_flattens": [ + null, + null, + true, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "first", + "last", + "first", + "first", + "first", + "last" + ], + "num_warps": 32, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 3, + 0 + ], + "range_warp_specializes": [ + null, + null, + false, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false, + null + ], + "range_flattens": [ + null, + true, + false, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "last", + "first", + "", + "", + "", + "first" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 0, + 1, + 0 + ], + "range_warp_specializes": [ + null, + true, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false, + null + ], + "range_flattens": [ + null, + true, + false, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "first", + "", + "first", + "first", + "first", + "" + ], + "num_warps": 32, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 4, + 4, + 0 + ], + "range_warp_specializes": [ + null, + null, + false, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true, + null + ], + "range_flattens": [ + null, + null, + false, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "", + "last", + "first", + "first", + "", + "first" + ], + "num_warps": 32, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 0, + 1, + 0 + ], + "range_warp_specializes": [ + null, + false, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true, + null + ], + "range_flattens": [ + null, + null, + null, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "last", + "last" + ], + "num_warps": 32, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 0, + 3, + 2 + ], + "range_warp_specializes": [ + null, + true, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true, + true + ], + "range_flattens": [ + null, + null, + null, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "first", + "", + "first", + "", + "last", + "last" + ], + "num_warps": 32, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 0, + 4, + 2 + ], + "range_warp_specializes": [ + null, + true, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false, + true + ], + "range_flattens": [ + null, + true, + true, + null + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "last", + "first", + "", + "first", + "last", + "" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4096, + 16 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 0, + 0, + 0 + ], + "range_warp_specializes": [ + null, + false, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null, + null + ], + "range_flattens": [ + null, + false, + null, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "first", + "", + "", + "first", + "first", + "last" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 4096, + 16 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 2, + 0 + ], + "range_warp_specializes": [ + null, + false, + null, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null, + null + ], + "range_flattens": [ + null, + false, + false, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "first", + "last", + "" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 4096, + 16 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 4, + 0, + 1 + ], + "range_warp_specializes": [ + null, + null, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true, + false + ], + "range_flattens": [ + null, + false, + false, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "last", + "last", + "first", + "", + "last", + "" + ], + "num_warps": 2, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 4096, + 16 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 2, + 0 + ], + "range_warp_specializes": [ + null, + false, + null, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null, + null + ], + "range_flattens": [ + null, + false, + false, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "first", + "last", + "" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 3, + 1, + 1, + 0 + ], + "range_warp_specializes": [ + false, + false, + null, + null + ], + "range_multi_buffers": [ + true, + true, + null, + null + ], + "range_flattens": [ + null, + null, + true, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "first", + "last", + "", + "last", + "", + "" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 8, + "maxnreg": 128 + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 0, + 0, + 4 + ], + "range_warp_specializes": [ + null, + null, + null, + false + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true, + null + ], + "range_flattens": [ + null, + false, + false, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "", + "first", + "first", + "", + "first", + "" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 3, + 4, + 0 + ], + "range_warp_specializes": [ + null, + null, + false, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null, + null + ], + "range_flattens": [ + null, + null, + null, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "last", + "first", + "last", + "last", + "", + "" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 2, + 0, + 0 + ], + "range_warp_specializes": [ + null, + null, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false, + null + ], + "range_flattens": [ + null, + null, + false, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "last", + "first", + "last", + "first", + "", + "" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "loop_orders": [ + [ + 1, + 0 + ] + ], + "range_unroll_factors": [ + 0, + 0, + 0, + 4 + ], + "range_warp_specializes": [ + null, + null, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true, + null + ], + "range_flattens": [ + null, + false, + false, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "", + "first", + "", + "", + "first", + "" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 3, + 1, + 1, + 0 + ], + "range_warp_specializes": [ + false, + false, + null, + null + ], + "range_multi_buffers": [ + true, + true, + null, + null + ], + "range_flattens": [ + null, + null, + true, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "first", + "last", + "", + "last", + "", + "" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 8, + "maxnreg": 128 + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 3, + 0, + 0 + ], + "range_warp_specializes": [ + null, + null, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false, + null + ], + "range_flattens": [ + null, + null, + true, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "last", + "first", + "last", + "last", + "", + "" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 1, + 3, + 0 + ], + "range_warp_specializes": [ + null, + false, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null, + null + ], + "range_flattens": [ + null, + false, + true, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "first", + "last", + "", + "last", + "", + "" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8192, + 8 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 1, + 4, + 0 + ], + "range_warp_specializes": [ + null, + true, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null, + null + ], + "range_flattens": [ + null, + true, + false, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "first", + "first", + "last", + "", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8192, + 8 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 3, + 0, + 0 + ], + "range_warp_specializes": [ + null, + null, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null, + null + ], + "range_flattens": [ + null, + null, + true, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "last", + "first", + "first", + "first", + "", + "last" + ], + "num_warps": 1, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8192, + 8 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 4, + 0, + 0 + ], + "range_warp_specializes": [ + null, + null, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false, + null + ], + "range_flattens": [ + null, + true, + true, + null + ], + "static_ranges": [ + true + ], + "load_eviction_policies": [ + "", + "", + "first", + "last", + "last", + "" + ], + "num_warps": 1, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 2048, + 8 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 4, + 1, + 3 + ], + "range_warp_specializes": [ + null, + null, + false, + false + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true, + false + ], + "range_flattens": [ + null, + true, + true, + null + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "first", + "", + "last" + ], + "num_warps": 2, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8192, + 8 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 3, + 0, + 2 + ], + "range_warp_specializes": [ + null, + null, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null, + null + ], + "range_flattens": [ + null, + null, + false, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "last", + "", + "last" + ], + "num_warps": 1, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8192, + 8 + ], + "loop_orders": [ + [ + 0, + 1 + ] + ], + "range_unroll_factors": [ + 0, + 3, + 0, + 2 + ], + "range_warp_specializes": [ + null, + null, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null, + null + ], + "range_flattens": [ + null, + null, + false, + false + ], + "static_ranges": [ + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "last", + "", + "last" + ], + "num_warps": 1, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/configs/rms_norm_per_block_quant/nvidia_h100.json b/vllm/kernels/helion/configs/rms_norm_per_block_quant/nvidia_h100.json new file mode 100644 index 00000000000..f7c148a407c --- /dev/null +++ b/vllm/kernels/helion/configs/rms_norm_per_block_quant/nvidia_h100.json @@ -0,0 +1,5017 @@ +[ + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 2048, + 32 + ], + "range_unroll_factors": [ + 0, + 2, + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + true, + null, + false + ], + "load_eviction_policies": [ + "last", + "first", + "last", + "last", + "last", + "last" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 4, + "maxnreg": 256 + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 1024, + 32 + ], + "range_unroll_factors": [ + 0, + 3, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "first", + "last", + "", + "last", + "last" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 2048, + 32 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "first", + "first", + "", + "last", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 2048, + 32 + ], + "range_unroll_factors": [ + 0, + 1, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "", + "", + "", + "", + "last" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 2048, + 32 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "last", + "", + "", + "last", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "last", + "first", + "" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 2048, + 32 + ], + "range_unroll_factors": [ + 0, + 4, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "first", + "", + "first", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 2048, + 32 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "first", + "first", + "last", + "first", + "first", + "last" + ], + "num_warps": 8, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 2048, + 32 + ], + "range_unroll_factors": [ + 0, + 1, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "last", + "", + "first", + "last", + "first" + ], + "num_warps": 8, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 2048, + 32 + ], + "range_unroll_factors": [ + 0, + 2, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "first", + "", + "" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 2048, + 32 + ], + "range_unroll_factors": [ + 0, + 2, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "first", + "last", + "", + "", + "first", + "last" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 2048, + 32 + ], + "range_unroll_factors": [ + 0, + 3, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "last", + "", + "last", + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 64, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 2048, + 32 + ], + "range_unroll_factors": [ + 0, + 3, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "", + "first", + "first", + "last", + "last", + "first" + ], + "num_warps": 8, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4096, + 64 + ], + "range_unroll_factors": [ + 1, + 2, + 0 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "first", + "", + "last", + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 4, + "maxnreg": 128 + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "", + "", + "last", + "last", + "last" + ], + "num_warps": 8, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4096, + 64 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "first", + "first", + "first", + "", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4096, + 64 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "", + "", + "", + "last", + "" + ], + "num_warps": 16, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4096, + 64 + ], + "range_unroll_factors": [ + 0, + 2, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "last", + "first", + "last", + "last", + "last", + "" + ], + "num_warps": 16, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 4096, + 64 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "", + "first", + "first", + "", + "" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4096, + 64 + ], + "range_unroll_factors": [ + 0, + 1, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "", + "last", + "last", + "last", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4096, + 64 + ], + "range_unroll_factors": [ + 0, + 1, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "last", + "last", + "", + "first", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4096, + 64 + ], + "range_unroll_factors": [ + 0, + 2, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "last", + "first", + "", + "", + "last", + "first" + ], + "num_warps": 16, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 4096, + 64 + ], + "range_unroll_factors": [ + 0, + 2, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "first", + "", + "last" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 2048, + 64 + ], + "range_unroll_factors": [ + 0, + 1, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "", + "first", + "last", + "last", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4096, + 64 + ], + "range_unroll_factors": [ + 0, + 1, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "last", + "last", + "last", + "last", + "" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 64, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4096, + 64 + ], + "range_unroll_factors": [ + 0, + 2, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "first", + "first", + "last", + "first", + "" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 8192, + 128 + ], + "range_unroll_factors": [ + 1, + 0, + 0 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + true, + false, + true + ], + "load_eviction_policies": [ + "last", + "last", + "", + "last", + "", + "" + ], + "num_warps": 32, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 2 + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 3, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "first", + "last", + "last" + ], + "num_warps": 8, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8192, + 128 + ], + "range_unroll_factors": [ + 0, + 2, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "", + "last", + "first", + "first", + "" + ], + "num_warps": 32, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8192, + 128 + ], + "range_unroll_factors": [ + 0, + 3, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "first", + "", + "", + "", + "" + ], + "num_warps": 32, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 8192, + 128 + ], + "range_unroll_factors": [ + 0, + 3, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "first", + "first", + "", + "first", + "" + ], + "num_warps": 32, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 0, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "last", + "last", + "first", + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8192, + 128 + ], + "range_unroll_factors": [ + 0, + 2, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "last", + "", + "first", + "", + "last", + "last" + ], + "num_warps": 32, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8192, + 128 + ], + "range_unroll_factors": [ + 0, + 3, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "last", + "last", + "first", + "", + "last", + "last" + ], + "num_warps": 32, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8192, + 128 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "first", + "first", + "first", + "", + "", + "" + ], + "num_warps": 32, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "", + "", + "" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8192, + 32 + ], + "range_unroll_factors": [ + 0, + 4, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "first", + "", + "" + ], + "num_warps": 8, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8192, + 128 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "last", + "last", + "", + "", + "first", + "first" + ], + "num_warps": 32, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 64, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8192, + 128 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "last", + "", + "", + "", + "last", + "" + ], + "num_warps": 32, + "num_stages": 2, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "first", + "last", + "first", + "last", + "first", + "" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "last", + "first", + "last", + "", + "first" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "first", + "", + "first", + "last", + "first", + "last" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "", + "", + "first" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "", + "", + "first" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "last", + "last", + "last" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 2, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "first", + "", + "", + "", + "" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 1, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "", + "", + "last", + "last", + "" + ], + "num_warps": 8, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "", + "", + "first" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "last", + "last", + "last" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 0, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "last", + "first", + "last", + "", + "first" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "", + "", + "first" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 4, + 0 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + true, + null, + true + ], + "load_eviction_policies": [ + "first", + "last", + "first", + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 8, + "maxnreg": 128 + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8192, + 16 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "", + "first", + "first", + "", + "" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 0, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "first", + "last", + "last", + "", + "", + "first" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "last", + "", + "first", + "", + "", + "last" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 0, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "", + "first", + "first", + "first" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 4096, + 16 + ], + "range_unroll_factors": [ + 0, + 1, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "last", + "", + "" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "first", + "", + "", + "last", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 0, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "", + "first", + "first", + "first" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 2, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "first", + "last", + "first", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8192, + 32 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "last", + "last", + "" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "first", + "", + "", + "last", + "", + "" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 3, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "first", + "first", + "last", + "first", + "", + "" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 0, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "", + "first", + "first", + "first" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "last", + "", + "last", + "first", + "first", + "first" + ], + "num_warps": 32, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8192, + 32 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "last", + "last", + "" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "last", + "last", + "", + "first", + "last", + "first" + ], + "num_warps": 32, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 2, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "first", + "", + "", + "", + "", + "" + ], + "num_warps": 32, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 0, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "first", + "", + "", + "first", + "", + "first" + ], + "num_warps": 32, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8192, + 32 + ], + "range_unroll_factors": [ + 0, + 3, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "first", + "first", + "", + "", + "last" + ], + "num_warps": 16, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "last", + "", + "first", + "", + "", + "first" + ], + "num_warps": 32, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 1, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "last", + "last", + "first", + "", + "last" + ], + "num_warps": 32, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "last", + "first", + "last", + "last" + ], + "num_warps": 32, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8192, + 32 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "", + "last", + "" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 2, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "first", + "first", + "first", + "first", + "last", + "last" + ], + "num_warps": 8, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "", + "last", + "", + "", + "" + ], + "num_warps": 32, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 4, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "last", + "", + "last", + "first", + "first", + "first" + ], + "num_warps": 32, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "last", + "last", + "last" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "group_size": 128, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 2048, + 16 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "", + "", + "first" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8192, + 32 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "first", + "last", + "last", + "" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "group_size": 128, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 4096, + 32 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "last", + "last", + "", + "", + "", + "last" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 4, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "last", + "last", + "last" + ], + "num_warps": 32, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 4, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "last", + "last", + "last" + ], + "num_warps": 32, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 1, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "first", + "", + "first", + "first", + "first" + ], + "num_warps": 32, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 4, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "first", + "", + "first", + "first", + "last", + "first" + ], + "num_warps": 32, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "last", + "", + "", + "last", + "last", + "first" + ], + "num_warps": 32, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 2, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "first", + "", + "first", + "", + "first" + ], + "num_warps": 32, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "last", + "", + "", + "last", + "last", + "first" + ], + "num_warps": 32, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 4, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "last", + "", + "last", + "last", + "last" + ], + "num_warps": 32, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8192, + 16 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "", + "first", + "first", + "", + "" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8192, + 8 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "first", + "first", + "last", + "", + "last" + ], + "num_warps": 4, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 2048, + 8 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "", + "first", + "", + "first", + "" + ], + "num_warps": 4, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8192, + 16 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "", + "first", + "first", + "", + "" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "last", + "last", + "first", + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "group_size": 128, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8192, + 64 + ], + "range_unroll_factors": [ + 0, + 0, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "last", + "last", + "first", + "first", + "", + "" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py new file mode 100644 index 00000000000..e7df42f4dbe --- /dev/null +++ b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py @@ -0,0 +1,299 @@ +# 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] + hidden_size_list = [2048, 4096, 5120] + group_size_list = [128] + in_dtype: torch.dtype = torch.bfloat16 + out_dtype: torch.dtype = current_platform.fp8_dtype() + scale_dtype: torch.dtype = torch.float32 + inputs = {} + + for hidden_size, group_size, num_tokens in product( + hidden_size_list, group_size_list, num_tokens_list + ): + input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype) + result = torch.empty(input.shape, device=input.device, dtype=out_dtype) + scale = torch.empty( + (num_tokens, hidden_size // group_size), + device=input.device, + dtype=scale_dtype, + ) + scale_ub = torch.mean(input).to(scale_dtype) + residual = torch.randn_like(input) + weight = torch.normal( + mean=1.0, + std=1.0, + size=(hidden_size,), + dtype=input.dtype, + device=input.device, + ) + epsilon = 1e-6 + + config_key = CaseKey( + { + "hidden_size": hidden_size, + "group_size": group_size, + "num_tokens": num_tokens, + } + ) + inputs[config_key] = ( + result, + input, + weight, + scale, + epsilon, + scale_ub, + residual, + group_size, + 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 hidden_size among available configs + (exact match preferred). + 2. Find the closest group_size among available configs + (exact match preferred). + 2. Among the num_tokens values tuned for that hidden_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 + + _, input, _, _, _, _, _, group_size, *_ = args + num_tokens, hidden_size = input.shape + + cache_key = (num_tokens, group_size, hidden_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["hidden_size"], {}).setdefault( + key["group_size"], [] + ).append(key["num_tokens"]) + + if not configs: + return None + + best_hidden_size = min(configs, key=lambda s: abs(s - hidden_size)) + best_group_size = min(configs[best_hidden_size], key=lambda s: abs(s - group_size)) + available_num_tokens = sorted(configs[best_hidden_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( + { + "hidden_size": best_hidden_size, + "group_size": best_group_size, + "num_tokens": best_num_tokens, + } + ) + _pick_cache[cache_key] = result + return result + + +def fake_impl( + result: torch.Tensor, # [num_tokens, hidden_size] + input: torch.Tensor, # [num_tokens, hidden_size] + weight: torch.Tensor, # [hidden_size] + scale: torch.Tensor, # [num_tokens, groups_per_row] + epsilon: float, + scale_ub: torch.Tensor | None, # [] + residual: torch.Tensor | None, # [num_tokens, hidden_size] + group_size: int, + is_scale_transposed: bool, # dummy +) -> None: + return + + +def baseline( + result: torch.Tensor, # [num_tokens, hidden_size] + input: torch.Tensor, # [num_tokens, hidden_size] + weight: torch.Tensor, # [hidden_size] + scale: torch.Tensor, # [num_tokens, groups_per_row] + epsilon: float, + scale_ub: torch.Tensor | None, # [] + residual: torch.Tensor | None, # [num_tokens, hidden_size] + group_size: int, + is_scale_transposed: bool, +) -> None: + torch.ops._C.rms_norm_per_block_quant( + result, + input, + weight, + scale, + epsilon, + scale_ub, + residual, + group_size, + is_scale_transposed, + ) + + +@register_kernel( + mutates_args=["result", "scale", "residual"], + 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 rms_norm_per_block_quant( + result: torch.Tensor, # [num_tokens, hidden_size] + input: torch.Tensor, # [num_tokens, hidden_size] + weight: torch.Tensor, # [hidden_size] + scale: torch.Tensor, # [num_tokens, groups_per_row] + epsilon: float, + scale_ub: torch.Tensor | None, # [] + residual: torch.Tensor | None, # [num_tokens, hidden_size] + group_size: int, + is_scale_transposed: bool, # dummy +) -> None: + # This code assumes batch_dim and num_tokens are flattened + assert input.ndim == 2 + num_tokens, hidden_size = input.shape + hl.specialize(hidden_size) + hl.specialize(group_size) + + groups_per_row = scale.shape[1] + hl.specialize(groups_per_row) + assert hidden_size % group_size == 0 and hidden_size // group_size == groups_per_row + assert scale.shape[0] == num_tokens + assert scale.dtype == torch.float32 + if scale.stride(1) > 1: + assert is_scale_transposed + + fp8_dtype = get_fp8_dtype() + assert result.dtype in [fp8_dtype, torch.int8] + assert result.is_contiguous() and input.is_contiguous() + + if scale_ub is not None: + assert result.dtype == fp8_dtype + assert scale_ub.dtype == torch.float32 + + assert input.dtype == weight.dtype + + if residual is not None: + assert residual.dtype == input.dtype + + assert group_size in [64, 128] + + quant_dtype = result.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) + + for tile_m in hl.tile(num_tokens, block_size=1): + rms = hl.zeros([tile_m], dtype=torch.float32) + for tile_n in hl.tile(hidden_size): + x_blk = input[tile_m, tile_n].to(torch.float32) + if residual is not None: + x_blk = x_blk + residual[tile_m, tile_n] + rms = rms + x_blk.pow(2).sum(dim=-1) + + rms = torch.rsqrt(rms * (1.0 / hidden_size) + epsilon) + + m_idx = tile_m.begin + hl.arange(tile_m.block_size) + m_blk = m_idx[:, None, None] + for tile_gn, tile_n in hl.tile( + [groups_per_row, group_size], block_size=[None, group_size] + ): + gn_idx = tile_gn.index + n_offset = tile_n.index + n_idx = gn_idx[:, None] * group_size + n_offset[None, :] + n_blk = n_idx[None, :, :] + mask = (gn_idx < groups_per_row)[None, :, None] + + x_blk = hl.load(input, [m_blk, n_blk], extra_mask=mask).to( + dtype=torch.float32 + ) + if residual is not None: + r_blk = hl.load(residual, [m_blk, n_blk], extra_mask=mask) + x_blk = x_blk + r_blk + + w_blk = hl.load(weight, [n_blk], extra_mask=mask) + x_norm_blk = (x_blk * rms[:, None, None]).to(input.dtype) * w_blk + s_blk = torch.amax(torch.abs(x_norm_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) + + scale[tile_m, tile_gn] = s_blk + + if quant_dtype == torch.int8: + y_blk = (x_norm_blk * (1.0 / s_blk[:, :, None])).round() + else: + y_blk = x_norm_blk / s_blk[:, :, None] + + y_blk = y_blk.clamp(qtype_traits_min, qtype_traits_max).to(result.dtype) + hl.store(result, [m_blk, n_blk], y_blk, extra_mask=mask) + + if residual is not None: + hl.store( + residual, [m_blk, n_blk], x_blk.to(residual.dtype), extra_mask=mask + ) diff --git a/vllm/kernels/helion/utils.py b/vllm/kernels/helion/utils.py index 130d79093b7..460fcc85065 100644 --- a/vllm/kernels/helion/utils.py +++ b/vllm/kernels/helion/utils.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility functions for Helion kernel management.""" +import torch + from vllm.logger import init_logger from vllm.platforms import current_platform @@ -78,3 +80,16 @@ def canonicalize_gpu_name(name: str) -> str: def get_canonical_gpu_name(device_id: int | None = None) -> str: return canonicalize_gpu_name(get_gpu_name(device_id)) + + +def get_fp8_dtype() -> torch.dtype: + return current_platform.fp8_dtype() + + +def get_int8_min_max() -> tuple[int, int]: + qtype_traits = torch.iinfo(torch.int8) + return qtype_traits.min, qtype_traits.max + + +def get_int8_min_scaling_factor() -> float: + return torch.finfo(torch.float32).eps