diff --git a/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py b/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py new file mode 100644 index 0000000000..50fd9b70d2 --- /dev/null +++ b/tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the dynamic_per_token_scaled_fp8_quant helion kernel + +Run `pytest tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py`. +""" + +from typing import Any + +import pytest +import torch +from torch._subclasses.fake_tensor import FakeTensorMode + +from tests.kernels.helion.utils import skip_if_platform_unsupported +from tests.kernels.quant_utils import FP8_DTYPE +from vllm.kernels.helion.case_key import CaseKey +from vllm.kernels.helion.config_manager import ConfigManager +from vllm.kernels.helion.ops.dynamic_per_token_scaled_fp8_quant import ( + _pick_cache, + baseline, + dynamic_per_token_scaled_fp8_quant, + pick_config, +) +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_helion +from vllm.utils.torch_utils import set_random_seed + +if not has_helion(): + pytest.skip( + "Helion is not installed. Install with: pip install vllm[helion]", + allow_module_level=True, + ) + + +def _generate_fake_input(num_tokens: int, hidden_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=current_platform.fp8_dtype() + ) + scale = torch.empty((num_tokens, 1), device=input.device, dtype=torch.float32) + scale_ub = torch.mean(input).to(torch.float32) + args = (result, input, scale, scale_ub) + return args + + +@pytest.fixture(autouse=True) +def reset_config_manager_singleton(): + ConfigManager.reset_instance() + ConfigManager() + yield + ConfigManager.reset_instance() + + +class TestDynamicPerTokenScaledFp8QuantConfigPicker: + def setup_method(self): + _pick_cache.clear() + + def test_config_picker_exact_match(self): + config_keys = [ + CaseKey({"hidden_size": 2048, "num_tokens": 16}), + CaseKey({"hidden_size": 4096, "num_tokens": 16}), + ] + + args = _generate_fake_input(16, 4096) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16}) + + def test_config_picker_closest_match(self): + config_keys = [ + CaseKey({"hidden_size": 2048, "num_tokens": 16}), + CaseKey({"hidden_size": 2048, "num_tokens": 32}), + CaseKey({"hidden_size": 4096, "num_tokens": 16}), + CaseKey({"hidden_size": 4096, "num_tokens": 32}), + ] + + args = _generate_fake_input(20, 3000) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey({"hidden_size": 2048, "num_tokens": 32}) + + def test_config_picker_no_configs(self): + config_keys: list[dict] = [] + + args = _generate_fake_input(16, 4096) + 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, "num_tokens": 16}), + CaseKey({"hidden_size": 4096, "num_tokens": 16}), + ] + + args = _generate_fake_input(32, 8192) + selected_key = pick_config(args, config_keys) + assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16}) + + +class TestDynamicPerTokenScaledFp8QuantCorrectness: + @pytest.mark.parametrize("num_tokens", [1, 7, 4096]) + @pytest.mark.parametrize("hidden_size", [17, 1024, 1025, 1026, 5137, 8193]) + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float]) + @pytest.mark.parametrize("has_scale_ub", [True, False]) + @pytest.mark.parametrize("seed", [0]) + def test_dynamic_per_token_fp8_quant( + self, + num_tokens: int, + hidden_size: int, + dtype: torch.dtype, + has_scale_ub: bool, + seed: int, + ) -> None: + skip_if_platform_unsupported("dynamic_per_token_scaled_fp8_quant") + set_random_seed(seed) + + x = ( + torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") + 1e-6 + ) # avoid nans + + scale_ub = ( + torch.mean(x).to(dtype=torch.float32, device="cuda") + if has_scale_ub + else None + ) + + ref_out = torch.empty(x.shape, device="cuda", dtype=FP8_DTYPE) + ref_scales = torch.empty((x.shape[0], 1), device="cuda", dtype=torch.float32) + baseline(ref_out, x, ref_scales, scale_ub) + + ops_out = torch.empty(x.shape, device="cuda", dtype=FP8_DTYPE) + ops_scales = torch.empty((x.shape[0], 1), device="cuda", dtype=torch.float32) + dynamic_per_token_scaled_fp8_quant(ops_out, x, ops_scales, scale_ub) + + torch.testing.assert_close(ref_scales, ops_scales) + # allow 1 ULP difference + assert ( + ref_out.view(torch.uint8).to(torch.int16) + - ops_out.view(torch.uint8).to(torch.int16) + ).abs().max() <= 1 + + +class TestDynamicPerTokenScaledFp8QuantIntegration: + def test_kernel_registration_integration(self): + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + assert "dynamic_per_token_scaled_fp8_quant" in registered_kernels + + kernel_wrapper = registered_kernels["dynamic_per_token_scaled_fp8_quant"] + assert kernel_wrapper.op_name == "dynamic_per_token_scaled_fp8_quant" + assert kernel_wrapper._config_picker is not None + assert kernel_wrapper._mutates_args == ["result", "scale"] + + def test_fake_impl_functionality(self): + skip_if_platform_unsupported("dynamic_per_token_scaled_fp8_quant") + from vllm.kernels.helion.register import get_registered_kernels + + registered_kernels = get_registered_kernels() + kernel_wrapper = registered_kernels["dynamic_per_token_scaled_fp8_quant"] + fake_impl = kernel_wrapper._fake_impl + + args = _generate_fake_input(16, 4096) + assert fake_impl(*args) is None diff --git a/vllm/kernels/helion/configs/dynamic_per_token_scaled_fp8_quant/nvidia_b200.json b/vllm/kernels/helion/configs/dynamic_per_token_scaled_fp8_quant/nvidia_b200.json new file mode 100644 index 0000000000..eb45fd7e61 --- /dev/null +++ b/vllm/kernels/helion/configs/dynamic_per_token_scaled_fp8_quant/nvidia_b200.json @@ -0,0 +1,2025 @@ +[ + { + "key": { + "hidden_size": 2048, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4096, + 4096 + ], + "range_unroll_factors": [ + 0, + 4, + 3 + ], + "range_warp_specializes": [ + null, + false, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "first", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 3, + 2 + ], + "range_warp_specializes": [ + null, + false, + false + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "last", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 2048, + 1024 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [ + null, + null, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "range_unroll_factors": [ + 0, + 2, + 4 + ], + "range_warp_specializes": [ + null, + null, + false + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "first", + "", + "first" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8192, + 512 + ], + "range_unroll_factors": [ + 0, + 2, + 0 + ], + "range_warp_specializes": [ + null, + false, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "first", + "first", + "" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 2048, + 1024 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [ + null, + null, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [ + null, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "first", + "last", + "first" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 3, + 2 + ], + "range_warp_specializes": [ + null, + false, + false + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "last", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 2048, + 1024 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [ + null, + null, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4096, + 4096 + ], + "range_unroll_factors": [ + 0, + 0, + 4 + ], + "range_warp_specializes": [ + false, + null, + false + ], + "range_multi_buffers": [ + true, + true, + true + ], + "range_flattens": [ + false, + true, + true + ], + "load_eviction_policies": [ + "last", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 1, + "maxnreg": 128 + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 3, + 4, + 1 + ], + "range_warp_specializes": [ + false, + false, + true + ], + "range_multi_buffers": [ + false, + false, + null + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "last", + "", + "first" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 8, + "maxnreg": 64 + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 2048, + 1024 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [ + null, + null, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4096, + 4096 + ], + "range_unroll_factors": [ + 0, + 0, + 4 + ], + "range_warp_specializes": [ + false, + null, + false + ], + "range_multi_buffers": [ + true, + true, + true + ], + "range_flattens": [ + false, + true, + true + ], + "load_eviction_policies": [ + "last", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 1, + "maxnreg": 128 + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 3, + 4, + 1 + ], + "range_warp_specializes": [ + false, + false, + true + ], + "range_multi_buffers": [ + false, + false, + null + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "last", + "", + "first" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 8, + "maxnreg": 64 + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 2048, + 1024 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [ + null, + null, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "last", + "last", + "first" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 3, + 2 + ], + "range_warp_specializes": [ + null, + false, + false + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "last", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 2048, + 1024 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [ + null, + null, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "last", + "last", + "first" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 3, + 2 + ], + "range_warp_specializes": [ + null, + false, + false + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "last", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 2048, + 1024 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [ + null, + null, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "last", + "last", + "first" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 0, + 4 + ], + "range_warp_specializes": [ + null, + true, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 2048, + 1024 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [ + null, + null, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "first", + "first", + "first" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 1, + 0 + ], + "range_warp_specializes": [ + null, + false, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "first", + "last", + "first" + ], + "num_warps": 1, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 2048, + 1024 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [ + null, + null, + true + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [ + null, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "first", + "last", + "first" + ], + "num_warps": 4, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8192, + 8192 + ], + "range_unroll_factors": [ + 0, + 3, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 2048, + 1024 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [ + null, + false, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8192, + 8192 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [ + null, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8192, + 8192 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [ + null, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 512, + 4096 + ], + "range_unroll_factors": [ + 0, + 0, + 4 + ], + "range_warp_specializes": [ + null, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [ + null, + false, + false + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8192, + 8192 + ], + "range_unroll_factors": [ + 0, + 3, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 512, + 512 + ], + "range_unroll_factors": [ + 0, + 0, + 4 + ], + "range_warp_specializes": [ + null, + false, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "last", + "" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 512, + 4096 + ], + "range_unroll_factors": [ + 0, + 0, + 4 + ], + "range_warp_specializes": [ + null, + null, + null + ], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 1024, + 512 + ], + "range_unroll_factors": [ + 2, + 3, + 2 + ], + "range_warp_specializes": [ + false, + false, + null + ], + "range_multi_buffers": [ + false, + null, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 64, + "maxnreg": 128 + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/configs/dynamic_per_token_scaled_fp8_quant/nvidia_h100.json b/vllm/kernels/helion/configs/dynamic_per_token_scaled_fp8_quant/nvidia_h100.json new file mode 100644 index 0000000000..217a293568 --- /dev/null +++ b/vllm/kernels/helion/configs/dynamic_per_token_scaled_fp8_quant/nvidia_h100.json @@ -0,0 +1,5185 @@ +[ + { + "key": { + "hidden_size": 512, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 256, + 512 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "last", + "", + "" + ], + "num_warps": 1, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 2048, + 2048 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "first", + "" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 4096, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "first", + "first", + "last" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 0, + 1, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 8192, + 4096 + ], + "range_unroll_factors": [ + 4, + 1, + 0 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + true, + null, + true + ], + "range_flattens": [ + false, + false, + true + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 1 + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "first", + "last" + ], + "num_warps": 32, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 512, + 512 + ], + "range_unroll_factors": [ + 0, + 1, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "first", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 2048, + 2048 + ], + "range_unroll_factors": [ + 0, + 2, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "last", + "", + "" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "range_unroll_factors": [ + 0, + 4, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "last", + "" + ], + "num_warps": 8, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 8192, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 16, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "last", + "last", + "first" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "last", + "last", + "" + ], + "num_warps": 16, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 512, + 512 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "last", + "first", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 2048, + 2048 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "last", + "last", + "first" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "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", + "", + "" + ], + "num_warps": 8, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 0, + 1, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 8, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8192, + 4096 + ], + "range_unroll_factors": [ + 0, + 0, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "first", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "last", + "first", + "" + ], + "num_warps": 16, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 256, + 512 + ], + "range_unroll_factors": [ + 0, + 4, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "last", + "first", + "" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 2048, + 2048 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "range_unroll_factors": [ + 0, + 4, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "last", + "first", + "" + ], + "num_warps": 8, + "num_stages": 4, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "last", + "last", + "first" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8192, + 4096 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "last", + "last", + "first" + ], + "num_warps": 16, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 0, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 16, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 512, + 128 + ], + "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": [ + "", + "last", + "last" + ], + "num_warps": 1, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 2048, + 2048 + ], + "range_unroll_factors": [ + 0, + 2, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "last", + "", + "" + ], + "num_warps": 16, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "", + "last" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "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", + "last" + ], + "num_warps": 8, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8192, + 4096 + ], + "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", + "last" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 0, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 16, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 512, + 256 + ], + "range_unroll_factors": [ + 0, + 3, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "first", + "last", + "last" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 2048, + 2048 + ], + "range_unroll_factors": [ + 0, + 0, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 0, + 1, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "last", + "", + "" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8192, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "last", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 16, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 512, + 256 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "first", + "first", + "last" + ], + "num_warps": 1, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 2048, + 2048 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "first", + "last" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 4096, + 4096 + ], + "range_unroll_factors": [ + 0, + 2, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "first", + "first" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 0, + 2, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 8, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8192, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "first", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 4, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 16, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 256, + 128 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 2048, + 1024 + ], + "range_unroll_factors": [ + 0, + 3, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "", + "last" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "range_unroll_factors": [ + 0, + 1, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "last", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 0, + 3, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "last", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 16, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "first", + "first", + "first" + ], + "num_warps": 16, + "num_stages": 3, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "first", + "" + ], + "num_warps": 32, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 512, + 512 + ], + "range_unroll_factors": [ + 0, + 4, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "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", + "last" + ], + "num_warps": 8, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "range_unroll_factors": [ + 0, + 3, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "first", + "last", + "" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "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", + "" + ], + "num_warps": 8, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 0, + 2, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 4, + "num_stages": 7, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 32768, + 16384 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 512, + 128 + ], + "range_unroll_factors": [ + 0, + 3, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "first", + "", + "" + ], + "num_warps": 1, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 512, + 1024 + ], + "range_unroll_factors": [ + 0, + 2, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "first", + "first", + "last" + ], + "num_warps": 4, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8192, + 4096 + ], + "range_unroll_factors": [ + 0, + 4, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + null, + null + ], + "load_eviction_policies": [ + "", + "", + "" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8192, + 8192 + ], + "range_unroll_factors": [ + 0, + 3, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "last", + "last", + "" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 32, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 512, + 512 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "last", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8192, + 8192 + ], + "range_unroll_factors": [ + 0, + 3, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 4096, + 2048 + ], + "range_unroll_factors": [ + 2, + 2, + 4 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false, + null, + true + ], + "range_flattens": [ + false, + null, + true + ], + "load_eviction_policies": [ + "", + "first", + "first" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 4, + "maxnreg": 32 + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8192, + 8192 + ], + "range_unroll_factors": [ + 0, + 3, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8192, + 8192 + ], + "range_unroll_factors": [ + 1, + 2, + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + true, + false, + null + ], + "range_flattens": [ + true, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 4, + "maxnreg": 128 + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 32, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 512, + 512 + ], + "range_unroll_factors": [ + 0, + 1, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "last", + "first", + "last" + ], + "num_warps": 1, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 1024, + 1024 + ], + "range_unroll_factors": [ + 0, + 2, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "last" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 4096, + 4096 + ], + "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", + "" + ], + "num_warps": 8, + "num_stages": 7, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 2, + 3, + 0 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false, + true, + true + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 8, + "num_stages": 8, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 8, + "maxnreg": 32 + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 0, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 16, + "num_stages": 1, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 32, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 512, + 256 + ], + "range_unroll_factors": [ + 1, + 1, + 4 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false, + false, + true + ], + "range_flattens": [ + false, + null, + false + ], + "load_eviction_policies": [ + "", + "first", + "first" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 16, + "maxnreg": 128 + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 2048, + 2048 + ], + "range_unroll_factors": [ + 3, + 2, + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + true, + false, + false + ], + "range_flattens": [ + false, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 16 + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 4096, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 8, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 0, + 0, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8192, + 8192 + ], + "range_unroll_factors": [ + 0, + 1, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 16, + "num_stages": 4, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 16, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 2, + 2, + 3 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false, + null, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "last", + "", + "last" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 32, + "maxnreg": 64 + } + }, + { + "key": { + "hidden_size": 512, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 512, + 512 + ], + "range_unroll_factors": [ + 3, + 4, + 0 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + false, + true, + null + ], + "range_flattens": [ + false, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 1, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 32, + "maxnreg": 128 + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 512, + 512 + ], + "range_unroll_factors": [ + 0, + 1, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "first", + "last" + ], + "num_warps": 1, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 4096, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 8, + "num_stages": 3, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 6144, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 0, + 4, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 8192, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8192, + 8192 + ], + "range_unroll_factors": [ + 0, + 4, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + false, + null + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 16, + "num_stages": 1, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 12288, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 16384, + 4096 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "last", + "last", + "first" + ], + "num_warps": 16, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 28672, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 32768, + 4096 + ], + "range_unroll_factors": [ + 1, + 0, + 1 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + true, + null, + false + ], + "range_flattens": [ + false, + false, + null + ], + "load_eviction_policies": [ + "last", + "", + "first" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_blocked", + "num_sm_multiplier": 64, + "maxnreg": 128 + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 1 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "last", + "" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 2 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "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", + "last" + ], + "num_warps": 8, + "num_stages": 5, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 4 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 1, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 8 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 4, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 16 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 0, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + false + ], + "range_flattens": [ + null, + null, + false + ], + "load_eviction_policies": [ + "", + "last", + "" + ], + "num_warps": 16, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 32 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 1, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 64 + }, + "config": { + "block_sizes": [ + 8192, + 2048 + ], + "range_unroll_factors": [ + 0, + 2, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "last", + "last" + ], + "num_warps": 16, + "num_stages": 8, + "indexing": [ + "pointer", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 128 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 1, + 1 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + null + ], + "range_flattens": [ + null, + null, + true + ], + "load_eviction_policies": [ + "", + "first", + "last" + ], + "num_warps": 16, + "num_stages": 5, + "indexing": [ + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 256 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 1, + 3 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "first", + "first", + "last" + ], + "num_warps": 8, + "num_stages": 2, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 512 + }, + "config": { + "block_sizes": [ + 8192, + 4096 + ], + "range_unroll_factors": [ + 0, + 0, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + false + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "", + "last", + "" + ], + "num_warps": 2, + "num_stages": 1, + "indexing": [ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 1024 + }, + "config": { + "block_sizes": [ + 8192, + 8192 + ], + "range_unroll_factors": [ + 0, + 3, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 2048 + }, + "config": { + "block_sizes": [ + 8192, + 8192 + ], + "range_unroll_factors": [ + 0, + 3, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 2, + "num_stages": 8, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 4096 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 3, + 4, + 2 + ], + "range_warp_specializes": [], + "range_multi_buffers": [ + true, + false, + true + ], + "range_flattens": [ + false, + true, + true + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 4, + "num_stages": 6, + "indexing": [ + "pointer", + "pointer", + "pointer", + "pointer", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "persistent_interleaved", + "num_sm_multiplier": 32, + "maxnreg": 256 + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 8192 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 1, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + false, + true + ], + "range_flattens": [ + null, + false, + false + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 5120, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 8192, + 1024 + ], + "range_unroll_factors": [ + 0, + 2, + 4 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + null, + true + ], + "range_flattens": [ + null, + true, + null + ], + "load_eviction_policies": [ + "last", + "", + "first" + ], + "num_warps": 8, + "num_stages": 5, + "indexing": [ + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 2048, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 2048, + 2048 + ], + "range_unroll_factors": [ + 0, + 0, + 0 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + false + ], + "range_flattens": [ + null, + false, + true + ], + "load_eviction_policies": [ + "", + "last", + "first" + ], + "num_warps": 4, + "num_stages": 2, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "pointer" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + }, + { + "key": { + "hidden_size": 4096, + "num_tokens": 16384 + }, + "config": { + "block_sizes": [ + 4096, + 4096 + ], + "range_unroll_factors": [ + 0, + 3, + 2 + ], + "range_warp_specializes": [], + "range_num_stages": [], + "range_multi_buffers": [ + null, + true, + true + ], + "range_flattens": [ + null, + true, + false + ], + "load_eviction_policies": [ + "", + "", + "first" + ], + "num_warps": 8, + "num_stages": 6, + "indexing": [ + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "tensor_descriptor", + "tensor_descriptor" + ], + "atomic_indexing": [], + "pid_type": "flat" + } + } +] \ No newline at end of file diff --git a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py new file mode 100644 index 0000000000..eef262dcfe --- /dev/null +++ b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py @@ -0,0 +1,165 @@ +# 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.register import register_kernel +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 + +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] + in_dtype: torch.dtype = torch.bfloat16 + out_dtype: torch.dtype = current_platform.fp8_dtype() + scale_dtype: torch.dtype = torch.float32 + inputs = {} + for num_tokens, hidden_size in product(num_tokens_list, hidden_size_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, 1), device=input.device, dtype=scale_dtype) + scale_ub = torch.mean(input).to(scale_dtype) + + config_key = CaseKey({"hidden_size": hidden_size, "num_tokens": num_tokens}) + inputs[config_key] = (result, input, scale, scale_ub) + + return inputs + + +_pick_cache: dict[tuple[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. Among the num_tokens values tuned for that hidden_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, *_ = args + num_tokens, hidden_size = input.shape + + cache_key = (num_tokens, hidden_size) + cached = _pick_cache.get(cache_key) + if cached is not None: + return cached + + configs: dict[int, list[int]] = {} + for key in config_keys: + if key.is_default(): + continue + configs.setdefault(key["hidden_size"], []).append(key["num_tokens"]) + + if not configs: + return None + + best_hidden_size = min(configs, key=lambda s: abs(s - hidden_size)) + available_num_tokens = sorted(configs[best_hidden_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, "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] + scale: torch.Tensor, # [num_tokens, 1] + scale_ub: torch.Tensor | None = None, # scalar tensor +) -> None: + return + + +def baseline( + result: torch.Tensor, # [num_tokens, hidden_size] + input: torch.Tensor, # [num_tokens, hidden_size] + scale: torch.Tensor, # [num_tokens, 1] + scale_ub: torch.Tensor | None = None, # scalar tensor +) -> None: + torch.ops._C.dynamic_per_token_scaled_fp8_quant(result, input, scale, scale_ub) + + +# Overwrite autotune_baseline_atol and autotune_baseline_rtol +# if too many configs failed due to baseline check during autotuning +@register_kernel( + mutates_args=["result", "scale"], + 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], + ), +) +def dynamic_per_token_scaled_fp8_quant( + result: torch.Tensor, # [num_tokens, hidden_size] + input: torch.Tensor, # [num_tokens, hidden_size] + scale: torch.Tensor, # [num_tokens, 1] + scale_ub: torch.Tensor | None = None, # scalar tensor +) -> 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) + + assert result.shape == input.shape + assert scale.shape[0] == num_tokens + assert scale.dtype == torch.float32 + assert input.stride()[-1] == 1 + assert result.stride()[-1] == 1 + + fp8_min, fp8_max = get_fp8_min_max() + min_scaling_factor = 1.0 / (fp8_max * 512.0) + + for tile_m in hl.tile(num_tokens, block_size=1): + s_blk = hl.zeros([tile_m], dtype=torch.float32) + for tile_n in hl.tile(hidden_size): + x_blk = input[tile_m, tile_n].to(dtype=torch.float32) + tmp_blk = torch.amax(torch.abs(x_blk), dim=-1) + s_blk = torch.maximum(s_blk, tmp_blk) + + 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 / fp8_max) + s_blk = s_blk.clamp(min=min_scaling_factor) + scale[tile_m, 0] = s_blk + + for tile_n in hl.tile(hidden_size): + x_blk = input[tile_m, tile_n].to(torch.float32) + y_blk = x_blk * (1.0 / s_blk[:, None]) + + result[tile_m, tile_n] = y_blk.clamp(fp8_min, fp8_max).to(result.dtype)