forked from Karylab-cklius/vllm
[Kernel][Helion][1/N] Add Helion kernel for rms_norm_per_block_quant (#36895)
Signed-off-by: Sean Chen <seachen@redhat.com> Co-authored-by: Yanan Cao <gmagogsfm@gmail.com>
This commit is contained in:
Xiaohong (Sean) Chen
co-authored by
Yanan Cao
Yanan Cao
parent
bf5149b516
commit
ce3ef17bec
@@ -0,0 +1,298 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the rms_norm_per_block_quant helion kernel
|
||||
|
||||
Run `pytest tests/kernels/helion/test_rms_norm_per_block_quant.py`.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import FakeTensorMode
|
||||
|
||||
from tests.kernels.helion.utils import skip_if_platform_unsupported
|
||||
from tests.kernels.quant_utils import FP8_DTYPE
|
||||
from vllm.kernels.helion.case_key import CaseKey
|
||||
from vllm.kernels.helion.config_manager import ConfigManager
|
||||
from vllm.kernels.helion.ops.rms_norm_per_block_quant import (
|
||||
_pick_cache,
|
||||
baseline,
|
||||
pick_config,
|
||||
rms_norm_per_block_quant,
|
||||
)
|
||||
from vllm.utils.import_utils import has_helion
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not has_helion():
|
||||
pytest.skip(
|
||||
"Helion is not installed. Install with: pip install vllm[helion]",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
def _generate_fake_input(
|
||||
num_tokens: int, hidden_size: int, group_size: int
|
||||
) -> tuple[Any, ...]:
|
||||
with FakeTensorMode():
|
||||
input = torch.randn(
|
||||
(num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
result = torch.empty(input.shape, device=input.device, dtype=FP8_DTYPE)
|
||||
scale = torch.empty(
|
||||
(num_tokens, hidden_size // group_size),
|
||||
device=input.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
scale_ub = torch.mean(input).to(scale.dtype)
|
||||
residual = torch.randn_like(input)
|
||||
weight = torch.normal(
|
||||
mean=1.0,
|
||||
std=1.0,
|
||||
size=(hidden_size,),
|
||||
dtype=input.dtype,
|
||||
device=input.device,
|
||||
)
|
||||
epsilon = 1e-6
|
||||
args = (
|
||||
result,
|
||||
input,
|
||||
weight,
|
||||
scale,
|
||||
epsilon,
|
||||
scale_ub,
|
||||
residual,
|
||||
group_size,
|
||||
False,
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_config_manager_singleton():
|
||||
ConfigManager.reset_instance()
|
||||
ConfigManager()
|
||||
yield
|
||||
ConfigManager.reset_instance()
|
||||
|
||||
|
||||
class TestRmsNormPerBlockQuantConfigPicker:
|
||||
def setup_method(self):
|
||||
_pick_cache.clear()
|
||||
|
||||
def test_config_picker_exact_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"hidden_size": 4096, "group_size": 128, "num_tokens": 16}
|
||||
)
|
||||
|
||||
def test_config_picker_closest_match(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(20, 3000, 70)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"hidden_size": 2048, "group_size": 64, "num_tokens": 32}
|
||||
)
|
||||
|
||||
def test_config_picker_no_configs(self):
|
||||
config_keys: list[dict] = []
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key is None
|
||||
|
||||
def test_config_picker_fallback_to_largest(self):
|
||||
config_keys = [
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
|
||||
CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}),
|
||||
]
|
||||
|
||||
args = _generate_fake_input(64, 8192, 256)
|
||||
selected_key = pick_config(args, config_keys)
|
||||
assert selected_key == CaseKey(
|
||||
{"hidden_size": 4096, "group_size": 128, "num_tokens": 32}
|
||||
)
|
||||
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float]
|
||||
QUANT_DTYPES = [torch.int8, FP8_DTYPE]
|
||||
VEC_HIDDEN_SIZES = [64, 1024]
|
||||
# Avoid combinatorial explosion with full Cartesian product
|
||||
NUM_TOKENS_HIDDEN_SIZES = [
|
||||
*[(1, i) for i in [64, 128, 1024, 5120]],
|
||||
*[(2048, i) for i in [64, 1024]],
|
||||
*[(4096, i) for i in [64]],
|
||||
]
|
||||
|
||||
ADD_RESIDUAL = [False, True]
|
||||
SCALE_UBS = [True, False]
|
||||
GROUP_SIZES = [64, 128]
|
||||
TMA_ALIGNMENTS = [0, 4]
|
||||
SEEDS = [0]
|
||||
EPS = 1e-6
|
||||
|
||||
|
||||
class TestRmsNormPerBlockQuantCorrectness:
|
||||
@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
|
||||
@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
|
||||
@pytest.mark.parametrize("is_scale_transposed", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"group_size, tma_alignment",
|
||||
[*itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
|
||||
)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
def test_rms_norm_per_block_quant(
|
||||
self,
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
add_residual: bool,
|
||||
has_scale_ub: bool,
|
||||
dtype: torch.dtype,
|
||||
quant_dtype: torch.dtype,
|
||||
is_scale_transposed: bool,
|
||||
group_size: int,
|
||||
tma_alignment: int,
|
||||
seed: int,
|
||||
) -> None:
|
||||
skip_if_platform_unsupported("rms_norm_per_block_quant")
|
||||
|
||||
set_random_seed(seed)
|
||||
|
||||
if hidden_size % group_size != 0:
|
||||
# skip
|
||||
return
|
||||
|
||||
if tma_alignment != 0 and hidden_size // group_size % tma_alignment == 0:
|
||||
# Skip tests where TMA alignment doesn't create extra padding to save time
|
||||
return
|
||||
|
||||
if has_scale_ub and quant_dtype != FP8_DTYPE:
|
||||
# skip
|
||||
return
|
||||
|
||||
scale = 1 / (hidden_size)
|
||||
input = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda") * scale
|
||||
weight = torch.normal(
|
||||
mean=1.0, std=1.0, size=(hidden_size,), dtype=dtype, device=input.device
|
||||
)
|
||||
residual = torch.randn_like(input) * scale if add_residual else None
|
||||
scale_ub = (
|
||||
torch.mean(input).to(dtype=torch.float32, device="cuda")
|
||||
if has_scale_ub
|
||||
else None
|
||||
)
|
||||
groups_per_row = hidden_size // group_size
|
||||
|
||||
ref_residual = residual.clone() if residual is not None else None
|
||||
ops_residual = residual.clone() if residual is not None else None
|
||||
ref_out = torch.empty(input.shape, device=input.device, dtype=quant_dtype)
|
||||
ops_out = ref_out.clone()
|
||||
|
||||
if is_scale_transposed:
|
||||
if tma_alignment == 0:
|
||||
ref_scales = torch.empty(
|
||||
(groups_per_row, num_tokens),
|
||||
device=input.device,
|
||||
dtype=torch.float32,
|
||||
).transpose(0, 1)
|
||||
else:
|
||||
tma_aligned_m = (
|
||||
(num_tokens + tma_alignment - 1) // tma_alignment * tma_alignment
|
||||
)
|
||||
shape = (num_tokens, groups_per_row)
|
||||
stride = (1, tma_aligned_m)
|
||||
ref_scales = torch.empty_strided(
|
||||
shape, stride, device=input.device, dtype=torch.float32
|
||||
)
|
||||
else:
|
||||
ref_scales = torch.empty(
|
||||
(num_tokens, groups_per_row), device=input.device, dtype=torch.float32
|
||||
)
|
||||
|
||||
ops_scales = ref_scales.clone()
|
||||
|
||||
baseline(
|
||||
ref_out,
|
||||
input,
|
||||
weight,
|
||||
ref_scales,
|
||||
EPS,
|
||||
scale_ub,
|
||||
ref_residual,
|
||||
group_size,
|
||||
is_scale_transposed,
|
||||
)
|
||||
ref_scales = ref_scales.contiguous()
|
||||
|
||||
rms_norm_per_block_quant(
|
||||
ops_out,
|
||||
input,
|
||||
weight,
|
||||
ops_scales,
|
||||
EPS,
|
||||
scale_ub,
|
||||
ops_residual,
|
||||
group_size,
|
||||
is_scale_transposed,
|
||||
)
|
||||
ops_scales = ops_scales.contiguous()
|
||||
|
||||
torch.testing.assert_close(ref_scales, ops_scales)
|
||||
# allow 1 ULP difference
|
||||
assert (
|
||||
ref_out.view(torch.uint8).to(torch.int16)
|
||||
- ops_out.view(torch.uint8).to(torch.int16)
|
||||
).abs().max() <= 1
|
||||
|
||||
if add_residual:
|
||||
torch.testing.assert_close(ref_residual, ops_residual)
|
||||
|
||||
|
||||
class TestRmsNormPerBlockQuantIntegration:
|
||||
def test_kernel_registration_integration(self):
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
assert "rms_norm_per_block_quant" in registered_kernels
|
||||
|
||||
kernel_wrapper = registered_kernels["rms_norm_per_block_quant"]
|
||||
assert kernel_wrapper.op_name == "rms_norm_per_block_quant"
|
||||
assert kernel_wrapper._config_picker is not None
|
||||
assert kernel_wrapper._mutates_args == ["result", "scale", "residual"]
|
||||
|
||||
def test_fake_impl_functionality(self):
|
||||
skip_if_platform_unsupported("rms_norm_per_block_quant")
|
||||
from vllm.kernels.helion.register import get_registered_kernels
|
||||
|
||||
registered_kernels = get_registered_kernels()
|
||||
kernel_wrapper = registered_kernels["rms_norm_per_block_quant"]
|
||||
fake_impl = kernel_wrapper._fake_impl
|
||||
|
||||
args = _generate_fake_input(16, 4096, 128)
|
||||
assert fake_impl(*args) is None
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,299 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from itertools import product
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.kernels.helion.case_key import CaseKey
|
||||
from vllm.kernels.helion.utils import (
|
||||
get_fp8_dtype,
|
||||
get_int8_min_max,
|
||||
get_int8_min_scaling_factor,
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
get_fp8_min_max,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.import_utils import has_helion
|
||||
|
||||
if not has_helion():
|
||||
raise ImportError(
|
||||
"Helion kernel requires helion to be installed. "
|
||||
"Install it with: pip install helion"
|
||||
)
|
||||
|
||||
import helion
|
||||
import helion.language as hl
|
||||
|
||||
from vllm.kernels.helion.register import register_kernel
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]:
|
||||
# TODO(xiaohongchen1991): it is difficult for kernel author to cover all
|
||||
# input property combination. Currently, dtypes are fixed. We need
|
||||
# optimization to bucket/skip some combinations
|
||||
num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
|
||||
hidden_size_list = [2048, 4096, 5120]
|
||||
group_size_list = [128]
|
||||
in_dtype: torch.dtype = torch.bfloat16
|
||||
out_dtype: torch.dtype = current_platform.fp8_dtype()
|
||||
scale_dtype: torch.dtype = torch.float32
|
||||
inputs = {}
|
||||
|
||||
for hidden_size, group_size, num_tokens in product(
|
||||
hidden_size_list, group_size_list, num_tokens_list
|
||||
):
|
||||
input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype)
|
||||
result = torch.empty(input.shape, device=input.device, dtype=out_dtype)
|
||||
scale = torch.empty(
|
||||
(num_tokens, hidden_size // group_size),
|
||||
device=input.device,
|
||||
dtype=scale_dtype,
|
||||
)
|
||||
scale_ub = torch.mean(input).to(scale_dtype)
|
||||
residual = torch.randn_like(input)
|
||||
weight = torch.normal(
|
||||
mean=1.0,
|
||||
std=1.0,
|
||||
size=(hidden_size,),
|
||||
dtype=input.dtype,
|
||||
device=input.device,
|
||||
)
|
||||
epsilon = 1e-6
|
||||
|
||||
config_key = CaseKey(
|
||||
{
|
||||
"hidden_size": hidden_size,
|
||||
"group_size": group_size,
|
||||
"num_tokens": num_tokens,
|
||||
}
|
||||
)
|
||||
inputs[config_key] = (
|
||||
result,
|
||||
input,
|
||||
weight,
|
||||
scale,
|
||||
epsilon,
|
||||
scale_ub,
|
||||
residual,
|
||||
group_size,
|
||||
False,
|
||||
)
|
||||
|
||||
return inputs
|
||||
|
||||
|
||||
_pick_cache: dict[tuple[int, int, int], CaseKey | None] = {}
|
||||
|
||||
|
||||
def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None:
|
||||
"""Pick the best pre-tuned config for the given input shape.
|
||||
|
||||
Selection strategy:
|
||||
1. Find the closest hidden_size among available configs
|
||||
(exact match preferred).
|
||||
2. Find the closest group_size among available configs
|
||||
(exact match preferred).
|
||||
2. Among the num_tokens values tuned for that hidden_size and group_size, pick
|
||||
the smallest num_tokens >= the input's num_tokens. If the input is
|
||||
larger than all available num_tokens, fall back to the largest.
|
||||
"""
|
||||
|
||||
if not config_keys:
|
||||
return None
|
||||
|
||||
_, input, _, _, _, _, _, group_size, *_ = args
|
||||
num_tokens, hidden_size = input.shape
|
||||
|
||||
cache_key = (num_tokens, group_size, hidden_size)
|
||||
cached = _pick_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
configs: dict[int, dict[int, list[int]]] = {}
|
||||
for key in config_keys:
|
||||
if key.is_default():
|
||||
continue
|
||||
configs.setdefault(key["hidden_size"], {}).setdefault(
|
||||
key["group_size"], []
|
||||
).append(key["num_tokens"])
|
||||
|
||||
if not configs:
|
||||
return None
|
||||
|
||||
best_hidden_size = min(configs, key=lambda s: abs(s - hidden_size))
|
||||
best_group_size = min(configs[best_hidden_size], key=lambda s: abs(s - group_size))
|
||||
available_num_tokens = sorted(configs[best_hidden_size][best_group_size])
|
||||
best_num_tokens = next(
|
||||
(n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1]
|
||||
)
|
||||
|
||||
result = CaseKey(
|
||||
{
|
||||
"hidden_size": best_hidden_size,
|
||||
"group_size": best_group_size,
|
||||
"num_tokens": best_num_tokens,
|
||||
}
|
||||
)
|
||||
_pick_cache[cache_key] = result
|
||||
return result
|
||||
|
||||
|
||||
def fake_impl(
|
||||
result: torch.Tensor, # [num_tokens, hidden_size]
|
||||
input: torch.Tensor, # [num_tokens, hidden_size]
|
||||
weight: torch.Tensor, # [hidden_size]
|
||||
scale: torch.Tensor, # [num_tokens, groups_per_row]
|
||||
epsilon: float,
|
||||
scale_ub: torch.Tensor | None, # []
|
||||
residual: torch.Tensor | None, # [num_tokens, hidden_size]
|
||||
group_size: int,
|
||||
is_scale_transposed: bool, # dummy
|
||||
) -> None:
|
||||
return
|
||||
|
||||
|
||||
def baseline(
|
||||
result: torch.Tensor, # [num_tokens, hidden_size]
|
||||
input: torch.Tensor, # [num_tokens, hidden_size]
|
||||
weight: torch.Tensor, # [hidden_size]
|
||||
scale: torch.Tensor, # [num_tokens, groups_per_row]
|
||||
epsilon: float,
|
||||
scale_ub: torch.Tensor | None, # []
|
||||
residual: torch.Tensor | None, # [num_tokens, hidden_size]
|
||||
group_size: int,
|
||||
is_scale_transposed: bool,
|
||||
) -> None:
|
||||
torch.ops._C.rms_norm_per_block_quant(
|
||||
result,
|
||||
input,
|
||||
weight,
|
||||
scale,
|
||||
epsilon,
|
||||
scale_ub,
|
||||
residual,
|
||||
group_size,
|
||||
is_scale_transposed,
|
||||
)
|
||||
|
||||
|
||||
@register_kernel(
|
||||
mutates_args=["result", "scale", "residual"],
|
||||
config_picker=pick_config,
|
||||
input_generator=generate_inputs,
|
||||
fake_impl=fake_impl,
|
||||
helion_settings=helion.Settings(
|
||||
autotune_baseline_fn=baseline,
|
||||
ignore_warnings=[helion.exc.TensorOperationInWrapper],
|
||||
),
|
||||
) # type: ignore[misc]
|
||||
def rms_norm_per_block_quant(
|
||||
result: torch.Tensor, # [num_tokens, hidden_size]
|
||||
input: torch.Tensor, # [num_tokens, hidden_size]
|
||||
weight: torch.Tensor, # [hidden_size]
|
||||
scale: torch.Tensor, # [num_tokens, groups_per_row]
|
||||
epsilon: float,
|
||||
scale_ub: torch.Tensor | None, # []
|
||||
residual: torch.Tensor | None, # [num_tokens, hidden_size]
|
||||
group_size: int,
|
||||
is_scale_transposed: bool, # dummy
|
||||
) -> None:
|
||||
# This code assumes batch_dim and num_tokens are flattened
|
||||
assert input.ndim == 2
|
||||
num_tokens, hidden_size = input.shape
|
||||
hl.specialize(hidden_size)
|
||||
hl.specialize(group_size)
|
||||
|
||||
groups_per_row = scale.shape[1]
|
||||
hl.specialize(groups_per_row)
|
||||
assert hidden_size % group_size == 0 and hidden_size // group_size == groups_per_row
|
||||
assert scale.shape[0] == num_tokens
|
||||
assert scale.dtype == torch.float32
|
||||
if scale.stride(1) > 1:
|
||||
assert is_scale_transposed
|
||||
|
||||
fp8_dtype = get_fp8_dtype()
|
||||
assert result.dtype in [fp8_dtype, torch.int8]
|
||||
assert result.is_contiguous() and input.is_contiguous()
|
||||
|
||||
if scale_ub is not None:
|
||||
assert result.dtype == fp8_dtype
|
||||
assert scale_ub.dtype == torch.float32
|
||||
|
||||
assert input.dtype == weight.dtype
|
||||
|
||||
if residual is not None:
|
||||
assert residual.dtype == input.dtype
|
||||
|
||||
assert group_size in [64, 128]
|
||||
|
||||
quant_dtype = result.dtype
|
||||
qtype_traits_min: int | float
|
||||
qtype_traits_max: int | float
|
||||
if quant_dtype == torch.int8:
|
||||
qtype_traits_min, qtype_traits_max = get_int8_min_max()
|
||||
min_scaling_factor = get_int8_min_scaling_factor()
|
||||
else:
|
||||
qtype_traits_min, qtype_traits_max = get_fp8_min_max()
|
||||
min_scaling_factor = 1.0 / (qtype_traits_max * 512.0)
|
||||
|
||||
qtype_max = float(qtype_traits_max)
|
||||
|
||||
for tile_m in hl.tile(num_tokens, block_size=1):
|
||||
rms = hl.zeros([tile_m], dtype=torch.float32)
|
||||
for tile_n in hl.tile(hidden_size):
|
||||
x_blk = input[tile_m, tile_n].to(torch.float32)
|
||||
if residual is not None:
|
||||
x_blk = x_blk + residual[tile_m, tile_n]
|
||||
rms = rms + x_blk.pow(2).sum(dim=-1)
|
||||
|
||||
rms = torch.rsqrt(rms * (1.0 / hidden_size) + epsilon)
|
||||
|
||||
m_idx = tile_m.begin + hl.arange(tile_m.block_size)
|
||||
m_blk = m_idx[:, None, None]
|
||||
for tile_gn, tile_n in hl.tile(
|
||||
[groups_per_row, group_size], block_size=[None, group_size]
|
||||
):
|
||||
gn_idx = tile_gn.index
|
||||
n_offset = tile_n.index
|
||||
n_idx = gn_idx[:, None] * group_size + n_offset[None, :]
|
||||
n_blk = n_idx[None, :, :]
|
||||
mask = (gn_idx < groups_per_row)[None, :, None]
|
||||
|
||||
x_blk = hl.load(input, [m_blk, n_blk], extra_mask=mask).to(
|
||||
dtype=torch.float32
|
||||
)
|
||||
if residual is not None:
|
||||
r_blk = hl.load(residual, [m_blk, n_blk], extra_mask=mask)
|
||||
x_blk = x_blk + r_blk
|
||||
|
||||
w_blk = hl.load(weight, [n_blk], extra_mask=mask)
|
||||
x_norm_blk = (x_blk * rms[:, None, None]).to(input.dtype) * w_blk
|
||||
s_blk = torch.amax(torch.abs(x_norm_blk), dim=-1).to(torch.float32)
|
||||
|
||||
if scale_ub is not None:
|
||||
scale_ub_s = hl.load(scale_ub, [])
|
||||
s_blk = s_blk.clamp(max=scale_ub_s)
|
||||
|
||||
s_blk = s_blk * (1.0 / qtype_max)
|
||||
s_blk = s_blk.clamp(min=min_scaling_factor)
|
||||
|
||||
scale[tile_m, tile_gn] = s_blk
|
||||
|
||||
if quant_dtype == torch.int8:
|
||||
y_blk = (x_norm_blk * (1.0 / s_blk[:, :, None])).round()
|
||||
else:
|
||||
y_blk = x_norm_blk / s_blk[:, :, None]
|
||||
|
||||
y_blk = y_blk.clamp(qtype_traits_min, qtype_traits_max).to(result.dtype)
|
||||
hl.store(result, [m_blk, n_blk], y_blk, extra_mask=mask)
|
||||
|
||||
if residual is not None:
|
||||
hl.store(
|
||||
residual, [m_blk, n_blk], x_blk.to(residual.dtype), extra_mask=mask
|
||||
)
|
||||
@@ -2,6 +2,8 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Utility functions for Helion kernel management."""
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
@@ -78,3 +80,16 @@ def canonicalize_gpu_name(name: str) -> str:
|
||||
|
||||
def get_canonical_gpu_name(device_id: int | None = None) -> str:
|
||||
return canonicalize_gpu_name(get_gpu_name(device_id))
|
||||
|
||||
|
||||
def get_fp8_dtype() -> torch.dtype:
|
||||
return current_platform.fp8_dtype()
|
||||
|
||||
|
||||
def get_int8_min_max() -> tuple[int, int]:
|
||||
qtype_traits = torch.iinfo(torch.int8)
|
||||
return qtype_traits.min, qtype_traits.max
|
||||
|
||||
|
||||
def get_int8_min_scaling_factor() -> float:
|
||||
return torch.finfo(torch.float32).eps
|
||||
|
||||
Reference in New Issue
Block a user