feat: add NineToothed fused QK norm RoPE

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)

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 <cassert>
+#include <cstddef>
+
+#include "operator.h"
+#include "tensor.h"
+
+namespace infini::ops {
+
+class MergedQkRmsNorm : public Operator<MergedQkRmsNorm> {
+ 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_

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,
+    )

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 <cassert>
+#include <cstdint>
+#include <vector>
+
+#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<MergedQkRmsNorm, Device::Type::kNvidia, 9>
+    : 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<std::uint64_t> weight_sizes{
+        static_cast<std::uint64_t>(input.size(0)),
+        static_cast<std::uint64_t>(input.size(1)),
+        static_cast<std::uint64_t>(input.size(2))};
+    std::vector<std::int64_t> weight_strides{
+        0, static_cast<std::int64_t>(weight.stride(0)),
+        static_cast<std::int64_t>(weight.stride(1))};
+    std::uint64_t scalar_shape[1] = {};
+    std::int64_t scalar_strides[1] = {};
+    double epsilon_value = static_cast<double>(epsilon);
+    std::int64_t normalized_width = static_cast<std::int64_t>(input.size(2));
+
+    ninetoothed::Tensor input_tensor(input);
+    ninetoothed::Tensor weight_tensor(const_cast<void*>(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<NineToothedStream>(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_

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

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,
+    )