diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 152f64101..02c9644c5 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -643,7 +643,7 @@ if(WITH_NVIDIA) endif() if(WITH_NVIDIA AND WITH_NINETOOTHED) - list(APPEND _infini_ops_smoke_ops fused_add_rms_norm) + list(APPEND _infini_ops_smoke_ops fused_add_rms_norm merged_qk_rms_norm) endif() if(((WITH_NVIDIA OR WITH_MOORE) AND WITH_TORCH) OR WITH_ILUVATAR) diff --git a/src/base/merged_qk_rms_norm.h b/src/base/merged_qk_rms_norm.h new file mode 100644 index 000000000..6f522f5d2 --- /dev/null +++ b/src/base/merged_qk_rms_norm.h @@ -0,0 +1,57 @@ +#ifndef INFINI_OPS_BASE_MERGED_QK_RMS_NORM_H_ +#define INFINI_OPS_BASE_MERGED_QK_RMS_NORM_H_ + +#include +#include + +#include "operator.h" +#include "tensor.h" + +namespace infini::ops { + +class MergedQkRmsNorm : public Operator { + public: + MergedQkRmsNorm(const Tensor input, const Tensor weight, float epsilon, + Tensor output) + : input_shape_{input.shape()}, + output_shape_{output.shape()}, + weight_shape_{weight.shape()}, + epsilon_{epsilon}, + num_heads_{input.size(1)}, + normalized_width_{input.size(2)} { + assert(input.ndim() == 3 && output.ndim() == 3 && + "`MergedQkRmsNorm` requires rank-3 packed Q/K tensors"); + assert(input.shape() == output.shape() && + "`MergedQkRmsNorm` requires matching input and output shapes"); + assert(input.stride(2) == 1 && output.stride(2) == 1 && + "`MergedQkRmsNorm` requires contiguous head dimensions"); + assert(weight.ndim() == 2 && weight.size(0) == num_heads_ && + weight.size(1) == normalized_width_ && + "`MergedQkRmsNorm` requires one weight row per packed head"); + assert(weight.stride(1) == 1 && + "`MergedQkRmsNorm` requires contiguous weight rows"); + assert(input.dtype() == output.dtype() && + output.dtype() == weight.dtype() && + "`MergedQkRmsNorm` requires matching tensor dtypes"); + } + + virtual void operator()(const Tensor input, const Tensor weight, + float epsilon, Tensor output) const = 0; + + protected: + Tensor::Shape input_shape_; + + Tensor::Shape output_shape_; + + Tensor::Shape weight_shape_; + + float epsilon_{}; + + Tensor::Size num_heads_{0}; + + Tensor::Size normalized_width_{0}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_BASE_MERGED_QK_RMS_NORM_H_ diff --git a/src/ninetoothed/ops/merged_qk_rms_norm/build.py b/src/ninetoothed/ops/merged_qk_rms_norm/build.py new file mode 100644 index 000000000..d8adf1926 --- /dev/null +++ b/src/ninetoothed/ops/merged_qk_rms_norm/build.py @@ -0,0 +1,49 @@ +import ninetoothed +import ninetoothed.language as ntl + +_CONFIGS = (((), {}, {}),) + + +def premake(): + tile_shape = (1, 1, 128) + + def arrangement(input, weight, epsilon, output, num_normalized_elements): + return ( + input.tile(tile_shape), + weight.tile(tile_shape), + epsilon, + output.tile(tile_shape), + num_normalized_elements, + ) + + def application(input, weight, epsilon, output, num_normalized_elements): + input_f32 = ntl.cast(input, ntl.float32) + scale = ntl.rsqrt( + ntl.sum(input_f32 * input_f32) / num_normalized_elements + epsilon + ) + output = input_f32 * ntl.cast(weight, ntl.float32) * scale # noqa: F841 + + return ( + arrangement, + application, + ( + ninetoothed.Tensor(3, other=0, dtype="bfloat16"), + ninetoothed.Tensor(3, dtype="bfloat16"), + ninetoothed.Tensor(0, dtype=ninetoothed.float64), + ninetoothed.Tensor(3, dtype="bfloat16"), + ninetoothed.Tensor(0, dtype=ninetoothed.int64), + ), + ) + + +def build(output_dir): + variant_dir = output_dir / "merged_qk_rms_norm" + variant_dir.mkdir(parents=True, exist_ok=True) + ninetoothed.build( + premake, + _CONFIGS, + caller="cuda", + kernel_name="infini_ops_ninetoothed_merged_qk_rms_norm", + output_dir=variant_dir, + lazy=False, + ) diff --git a/src/ninetoothed/ops/merged_qk_rms_norm/merged_qk_rms_norm.h b/src/ninetoothed/ops/merged_qk_rms_norm/merged_qk_rms_norm.h new file mode 100644 index 000000000..2281fcbea --- /dev/null +++ b/src/ninetoothed/ops/merged_qk_rms_norm/merged_qk_rms_norm.h @@ -0,0 +1,68 @@ +#ifndef INFINI_OPS_NINETOOTHED_MERGED_QK_RMS_NORM_H_ +#define INFINI_OPS_NINETOOTHED_MERGED_QK_RMS_NORM_H_ + +#include +#include +#include + +#include "base/merged_qk_rms_norm.h" +#include "data_type.h" +#include "merged_qk_rms_norm/infini_ops_ninetoothed_merged_qk_rms_norm.h" +#include "ninetoothed/tensor.h" + +namespace infini::ops { + +template <> +class Operator + : public MergedQkRmsNorm { + public: + using MergedQkRmsNorm::MergedQkRmsNorm; + using MergedQkRmsNorm::operator(); + + static bool supports(const Tensor input, const Tensor weight, float, + Tensor output) { + return input.ndim() == 3 && output.ndim() == 3 && + input.shape() == output.shape() && input.size(2) == 128 && + input.stride(1) == 128 && input.stride(2) == 1 && + output.stride(1) == 128 && output.stride(2) == 1 && + input.dtype() == DataType::kBFloat16 && + output.dtype() == input.dtype() && weight.dtype() == input.dtype() && + weight.ndim() == 2 && weight.size(0) == input.size(1) && + weight.size(1) == 128 && weight.stride(0) == 128 && + weight.stride(1) == 1; + } + + void operator()(const Tensor input, const Tensor weight, float epsilon, + Tensor output) const override { + std::vector weight_sizes{ + static_cast(input.size(0)), + static_cast(input.size(1)), + static_cast(input.size(2))}; + std::vector weight_strides{ + 0, static_cast(weight.stride(0)), + static_cast(weight.stride(1))}; + std::uint64_t scalar_shape[1] = {}; + std::int64_t scalar_strides[1] = {}; + double epsilon_value = static_cast(epsilon); + std::int64_t normalized_width = static_cast(input.size(2)); + + ninetoothed::Tensor input_tensor(input); + ninetoothed::Tensor weight_tensor(const_cast(weight.data()), + weight_sizes.data(), + weight_strides.data()); + ninetoothed::Tensor epsilon_tensor(epsilon_value, scalar_shape, + scalar_strides); + ninetoothed::Tensor output_tensor(output); + ninetoothed::Tensor width_tensor(normalized_width, scalar_shape, + scalar_strides); + + const auto result = launch_infini_ops_ninetoothed_merged_qk_rms_norm( + static_cast(stream_), input_tensor, weight_tensor, + epsilon_tensor, output_tensor, width_tensor); + assert(result == 0 && "NineToothed `MergedQkRmsNorm` launch failed"); + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NINETOOTHED_MERGED_QK_RMS_NORM_H_ diff --git a/tests/conftest.py b/tests/conftest.py index cfc2691c1..83b70afb3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -245,6 +245,7 @@ def _is_smoke_item(item): "tests/test_cutlass_scaled_mm.py": _is_smoke_cutlass_scaled_mm_case, "tests/test_flash_attn_varlen_func.py": _is_smoke_flash_attn_varlen_func_case, "tests/test_fused_add_rms_norm.py": _is_smoke_fused_add_rms_norm_case, + "tests/test_merged_qk_rms_norm.py": _is_smoke_merged_qk_rms_norm_case, "tests/test_swiglu.py": _is_smoke_swiglu_case, "tests/test_torch_ops.py": _is_smoke_torch_op_case, } @@ -474,6 +475,10 @@ def _is_smoke_fused_add_rms_norm_case(params): ) +def _is_smoke_merged_qk_rms_norm_case(params): + return params.get("seq_len") == 1 + + def _is_smoke_moe_sum_case(params): return ( params.get("topk") == 2 diff --git a/tests/test_merged_qk_rms_norm.py b/tests/test_merged_qk_rms_norm.py new file mode 100644 index 000000000..9b4f5aa8c --- /dev/null +++ b/tests/test_merged_qk_rms_norm.py @@ -0,0 +1,51 @@ +import infini.ops +import pytest +import torch + +from tests.utils import get_stream + + +@pytest.mark.parametrize("seq_len", (1, 17)) +def test_merged_qk_rms_norm(seq_len, device, implementation_index): + if device != "cuda" or implementation_index != 9: + pytest.skip("NineToothed NVIDIA packed Q/K fast path") + + torch.manual_seed(7) + dtype = torch.bfloat16 + query_heads = 16 + key_heads = 8 + head_dim = 128 + projection = torch.randn((seq_len, 32, head_dim), device=device, dtype=dtype) + input = projection[:, : query_heads + key_heads] + query_weight = torch.randn(head_dim, device=device, dtype=dtype) + key_weight = torch.randn(head_dim, device=device, dtype=dtype) + weight = torch.cat( + (query_weight.expand(query_heads, -1), key_weight.expand(key_heads, -1)) + ) + output = torch.empty( + (seq_len, query_heads + key_heads, head_dim), + device=device, + dtype=dtype, + ) + + infini.ops.merged_qk_rms_norm( + input, + weight, + 1e-6, + output, + implementation_index=implementation_index, + stream=get_stream(device), + ) + + expected = input.float() + expected[:seq_len, :query_heads] *= query_weight.float() + expected[:seq_len, query_heads:] *= key_weight.float() + expected *= torch.rsqrt( + input.float().square().sum(dim=-1, keepdim=True) / head_dim + 1e-6 + ) + torch.testing.assert_close( + output.float(), + expected.to(dtype).float(), + rtol=2e-2, + atol=1e-2, + )