Accept left offsets when applying position encodings

CMakeLists.txt

@@ -155,6 +155,8 @@ set(SOURCES
  src/ops/mul.cc
  src/ops/multinomial.cc
  src/ops/multinomial_cpu.cc
+ src/ops/position_encodings_add.cc
+ src/ops/position_encodings_add_cpu.cc
  src/ops/quantize.cc
  src/ops/quantize_cpu.cc
  src/ops/relu.cc
@@ -511,6 +513,7 @@ if (WITH_CUDA)
  src/ops/layer_norm_gpu.cu
  src/ops/mean_gpu.cu
  src/ops/multinomial_gpu.cu
+ src/ops/position_encodings_add_gpu.cu
  src/ops/rms_norm_gpu.cu
  src/ops/rotary_gpu.cu
  src/ops/softmax_gpu.cu

include/ctranslate2/layers/common.h

@@ -89,10 +89,14 @@ namespace ctranslate2 {
  // Base class for position encoders.
  class PositionEncoder : public Layer {
  public:
- void operator()(StorageView& input, dim_t index = 0);
- void operator()(const StorageView& input, StorageView& output, dim_t index = 0);
+ void operator()(const StorageView& input,
+ StorageView& output,
+ dim_t step = 0,
+ const StorageView* offsets = nullptr);
  protected:
  virtual const StorageView& get_position_encoding(dim_t max_time) = 0;
+ private:
+ ops::PositionEncodingsAdd _add_op;
  };
 
  // Concrete position encoder loading encoding vectors from the model.

include/ctranslate2/ops/ops.h

@@ -39,3 +39,4 @@
 #include "median_filter.h"
 #include "rotary.h"
 #include "alibi_add.h"
+#include "position_encodings_add.h"

include/ctranslate2/ops/position_encodings_add.h

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "op.h"
+
+namespace ctranslate2 {
+ namespace ops {
+
+ class PositionEncodingsAdd : public Op {
+ public:
+ void operator()(const StorageView& input,
+ const StorageView& encodings,
+ StorageView& output,
+ const StorageView* offsets = nullptr,
+ const dim_t step = 0) const;
+
+ private:
+ template <Device D, typename T>
+ void compute(const dim_t step,
+ const StorageView* offsets,
+ const StorageView& input,
+ const StorageView& encodings,
+ StorageView& output) const;
+ };
+
+ }
+}

src/layers/common.cc

@@ -4,7 +4,6 @@
 
 #include "ctranslate2/ops/activation.h"
 #include "cpu/backend.h"
-#include "dispatch.h"
 
 namespace ctranslate2 {
  namespace layers {
@@ -148,34 +147,13 @@ namespace ctranslate2 {
  }
 
 
- void PositionEncoder::operator()(StorageView& input, dim_t index) {
+ void PositionEncoder::operator()(const StorageView& input,
+ StorageView& output,
+ dim_t step,
+ const StorageView* offsets) {
  const dim_t time = input.dim(1);
- const dim_t depth = input.dim(-1);
- const dim_t max_time = std::max(time, index + 1);
- const StorageView& encodings = get_position_encoding(max_time);
- const dim_t num_encodings = encodings.dim(0);
-
- if (max_time > num_encodings)
- throw std::runtime_error("No position encodings are defined for positions >= "
- + std::to_string(num_encodings)
- + ", but got position "
- + std::to_string(max_time - 1));
- if (depth != encodings.dim(1))
- throw std::invalid_argument("Shape mismatch: position encodings have depth "
- + std::to_string(encodings.dim(1))
- + ", but the input has depth "
- + std::to_string(depth));
-
- DEVICE_AND_TYPE_DISPATCH(input.device(), input.dtype(),
- primitives<D>::add_batch_broadcast(encodings.data<T>() + index * depth,
- input.data<T>(),
- time * depth,
- input.size()));
- }
-
- void PositionEncoder::operator()(const StorageView& input, StorageView& output, dim_t index) {
- output = input;
- operator()(output, index);
+ const StorageView& encodings = get_position_encoding(step + time);
+ _add_op(input, encodings, output, offsets, step);
  }
 
 

src/layers/transformer.cc

@@ -259,7 +259,7 @@ namespace ctranslate2 {
  if (_embeddings_scale)
  ops::Mul()(input, *_embeddings_scale, input);
  if (_position_encoder)
- (*_position_encoder)(input);
+ (*_position_encoder)(input, input);
  if (_layernorm_embedding)
  (*_layernorm_embedding)(input, input);
 
@@ -460,7 +460,7 @@ namespace ctranslate2 {
  if (layer_in.rank() == 2)
  layer_in.expand_dims(1);
  if (_position_encoder)
- (*_position_encoder)(layer_in, std::max(step, dim_t(0)));
+ (*_position_encoder)(layer_in, layer_in, std::max(step, dim_t(0)));
  if (_layernorm_embedding)
  (*_layernorm_embedding)(layer_in, layer_in);
 

src/layers/whisper.cc

@@ -52,7 +52,7 @@ namespace ctranslate2 {
  _gelu(output, output);
 
  _transpose(output, input);
- _position_embedding(input);
+ _position_embedding(input, input);
 
  for (const auto& layer : _layers) {
  (*layer)(input, nullptr, output);

src/ops/position_encodings_add.cc

@@ -0,0 +1,50 @@
+#include "ctranslate2/ops/position_encodings_add.h"
+
+#include "dispatch.h"
+
+namespace ctranslate2 {
+ namespace ops {
+
+ void PositionEncodingsAdd::operator()(const StorageView& input,
+ const StorageView& encodings,
+ StorageView& output,
+ const StorageView* offsets,
+ const dim_t step) const {
+ PROFILE("PositionEncodingsAdd");
+
+ const dim_t time = input.dim(1);
+ const dim_t depth = input.dim(2);
+ const dim_t max_time = time + step;
+
+ if (max_time > encodings.dim(0))
+ throw std::runtime_error("No position encodings are defined for positions >= "
+ + std::to_string(encodings.dim(0))
+ + ", but got position "
+ + std::to_string(max_time - 1));
+
+ if (depth != encodings.dim(1))
+ throw std::invalid_argument("Shape mismatch: position encodings have depth "
+ + std::to_string(encodings.dim(1))
+ + ", but the input has depth "
+ + std::to_string(depth));
+
+ output.resize_as(input);
+
+ if (offsets) {
+ DEVICE_AND_FLOAT_DISPATCH(
+ "PositionEncodingsAdd", input.device(), input.dtype(),
+ (compute<D, T>(step, offsets, input, encodings, output)));
+
+ } else {
+ DEVICE_AND_FLOAT_DISPATCH(
+ "PositionEncodingsAdd", input.device(), input.dtype(),
+ (primitives<D>::add_batch_broadcast(encodings.data<T>() + step * depth,
+ input.data<T>(),
+ output.data<T>(),
+ time * depth,
+ input.size())));
+ }
+ }
+
+ }
+}

src/ops/position_encodings_add_cpu.cc

@@ -0,0 +1,48 @@
+#include "ctranslate2/ops/position_encodings_add.h"
+
+#include "cpu/parallel.h"
+
+namespace ctranslate2 {
+ namespace ops {
+
+ template <Device D, typename T>
+ void PositionEncodingsAdd::compute(const dim_t step,
+ const StorageView* offsets,
+ const StorageView& input,
+ const StorageView& encodings,
+ StorageView& output) const {
+ const dim_t batch_size = input.dim(0);
+ const dim_t time = input.dim(1);
+ const dim_t depth = input.dim(2);
+
+ cpu::parallel_for(0, batch_size * time, 1, [&](const dim_t begin, const dim_t end) {
+ for (dim_t i = begin; i < end; ++i) {
+ const dim_t b = i / time;
+ const dim_t t = i % time;
+
+ const dim_t offset = offsets ? offsets->at<int32_t>(b) : 0;
+ const dim_t encoding_offset = t - offset + step;
+
+ if (encoding_offset < 0)
+ continue;
+
+ primitives<Device::CPU>::add(encodings.index<float>({encoding_offset, 0}),
+ input.index<float>({b, t, 0}),
+ output.index<float>({b, t, 0}),
+ depth);
+ }
+ });
+ }
+
+#define DECLARE_IMPL(T) \
+ template void \
+ PositionEncodingsAdd::compute<Device::CPU, T>(const dim_t, \
+ const StorageView*, \
+ const StorageView&, \
+ const StorageView&, \
+ StorageView&) const;
+
+ DECLARE_IMPL(float)
+
+ }
+}

src/ops/position_encodings_add_gpu.cu

@@ -0,0 +1,73 @@
+#include "ctranslate2/ops/position_encodings_add.h"
+
+#include "type_dispatch.h"
+#include "cuda/helpers.h"
+
+namespace ctranslate2 {
+ namespace ops {
+
+ template <typename T, typename AddFunc>
+ __global__ void position_encodings_add_kernel(const T* input,
+ const T* encodings,
+ T* output,
+ const int32_t* offsets,
+ cuda::index_t step,
+ cuda::index_t max_time,
+ cuda::index_t depth,
+ const AddFunc& add_func) {
+ const cuda::index_t batch = blockIdx.x / max_time;
+ const cuda::index_t time = blockIdx.x % max_time;
+
+ const int32_t offset = offsets ? offsets[batch] : 0;
+ const int32_t encoding_offset = time - offset + step;
+
+ if (encoding_offset < 0)
+ return;
+
+ input += blockIdx.x * depth;
+ output += blockIdx.x * depth;
+ encodings += encoding_offset * depth;
+
+ for (cuda::index_t i = threadIdx.x; i < depth; i += blockDim.x) {
+ output[i] = add_func(input[i], encodings[i]);
+ }
+ }
+
+ template <Device D, typename T>
+ void PositionEncodingsAdd::compute(const dim_t step,
+ const StorageView* offsets,
+ const StorageView& input,
+ const StorageView& encodings,
+ StorageView& output) const {
+ const dim_t batch_size = input.dim(0);
+ const dim_t time = input.dim(1);
+ const dim_t depth = input.dim(2);
+
+ const dim_t blocks = std::min(batch_size * time, cuda::max_blocks);
+ const dim_t threads = std::min(depth, cuda::max_threads);
+
+ position_encodings_add_kernel<<<blocks, threads, 0, cuda::get_cuda_stream()>>>(
+ cuda::device_cast(input.data<T>()),
+ cuda::device_cast(encodings.data<T>()),
+ cuda::device_cast(output.data<T>()),
+ offsets ? offsets->data<int32_t>() : nullptr,
+ step,
+ time,
+ depth,
+ cuda::plus<cuda::device_type<T>>());
+ }
+
+#define DECLARE_IMPL(T) \
+ template void \
+ PositionEncodingsAdd::compute<Device::CUDA, T>(const dim_t, \
+ const StorageView*, \
+ const StorageView&, \
+ const StorageView&, \
+ StorageView&) const;
+
+ DECLARE_IMPL(float)
+ DECLARE_IMPL(float16_t)
+ DECLARE_IMPL(bfloat16_t)
+
+ }
+}

tests/layers_test.cc

@@ -223,6 +223,36 @@ TEST(LayerTest, PadderIgnore) {
  expect_storage_eq(x, original);
 }
 
+TEST_P(LayerDeviceFPTest, PositionEncoderOffset) {
+ const Device device = GetParam().device;
+ const DataType dtype = GetParam().dtype;
+ const float error = GetParam().error;
+
+ layers::SinusoidalPositionEncoder position_encoder(4, dtype, device);
+
+ StorageView offsets({2}, std::vector<int32_t>{3, 1}, device);
+ dim_t step = 5;
+
+ StorageView expected_encodings(dtype, device);
+
+ {
+ StorageView zero({2, 5, 4}, 0.f, device);
+ StorageView encodings(dtype, device);
+ position_encoder(zero.to(dtype), encodings);
+
+ StorageView position_ids({2, 1}, std::vector<int32_t>{2, 4}, device);
+ ops::Gather(/*axis=*/1, /*batch_dims=*/1)(encodings, position_ids, expected_encodings);
+ }
+
+ {
+ StorageView zero({2, 1, 4}, 0.f, device);
+ StorageView encodings(dtype, device);
+ position_encoder(zero.to(dtype), encodings, step, &offsets);
+
+ expect_storage_eq(encodings.to_float32(), expected_encodings.to_float32(), error);
+ }
+}
+
 TEST(LayerTest, PositionEncoderNoSharedState) {
  // Test case for issue: http://forum.opennmt.net/t/ctranslate2-c-api-returns-strange-results-when-initializing-2-models/3208
  layers::SinusoidalPositionEncoder position_encoder_1(4);
@@ -233,7 +263,7 @@ TEST(LayerTest, PositionEncoderNoSharedState) {
  {1, 1, 4}, std::vector<float>{0.1, -2.3, 0.5, 1.2});
  StorageView expected(
  {1, 1, 4}, std::vector<float>{0.941471, -2.2999, 1.0403, 2.2});
- position_encoder_1(input);
+ position_encoder_1(input, input);
  expect_storage_eq(input, expected, 1e-5);
  }
 
@@ -242,7 +272,7 @@ TEST(LayerTest, PositionEncoderNoSharedState) {
  {1, 1, 6}, std::vector<float>{-0.2, -1.3, 0.1, -0.6, 2.0, 1.1});
  StorageView expected(
  {1, 1, 6}, std::vector<float>{0.641471, -1.29, 0.1001, -0.0596977, 2.99995, 2.1});
- position_encoder_2(input);
+ position_encoder_2(input, input);
  expect_storage_eq(input, expected, 1e-5);
  }
 }