Method forward

src/fastertransformer/models/decoder/Decoder.cc:160–321

template<typename T>
void Decoder<T>::forward(std::vector<Tensor>*                      output_tensors,
                         const std::vector<Tensor>*                input_tensors,
                         const std::vector<DecoderLayerWeight<T>>* decoder_layer_weight)
{
    // input tensors:
    //      decoder_input [batch_size, hidden_dimension],
    //      encoder_output [batch_size, mem_max_seq_len, memory_hidden_dimension],
    //      encoder_sequence_length [batch_size],
    //      finished [batch_size],
    //      step [1] on cpu
    //      sequence_lengths [batch_size]
    //      cache_indirection [local_batch_size / beam_width, beam_width, max_seq_len]
    //              Here, local_batch_size contains the beam_width, so local_batch_size / beam_width
    //              is real local_batch_size.

    // output tensors:
    //      decoder_output [batch_size, hidden_dimension],
    //      key_cache [num_layer, batch, head_num, size_per_head // x, max_seq_len, x]
    //      value_cache [num_layer, batch, head_num, max_seq_len, size_per_head]
    //      key_mem_cache [num_layer, batch_size, mem_max_seq_len, hidden_dimension],
    //      value_mem_cache [num_layer, batch_size, mem_max_seq_len, hidden_dimension]

    FT_CHECK(input_tensors->size() == 7);
    FT_CHECK(output_tensors->size() == 5);
    isValidBatchSize(input_tensors->at(0).shape[0]);
    allocateBuffer(input_tensors->at(0).shape[0]);

    const size_t   batch_size      = (size_t)input_tensors->at(0).shape[0];
    const size_t   mem_max_seq_len = (size_t)input_tensors->at(1).shape[1];
    const DataType data_type       = getTensorType<T>();

    TensorMap cross_attention_input_tensors{
        {"input_query", Tensor{MEMORY_GPU, data_type, {batch_size, hidden_units_}, normed_self_attn_output_}},
        {"encoder_output", input_tensors->at(1)},
        {"encoder_sequence_length", input_tensors->at(2)},
        {"finished", input_tensors->at(3)},
        {"step", input_tensors->at(4)}};
    TensorMap self_attention_input_tensors{
        {"input_query", Tensor{MEMORY_GPU, data_type, {batch_size, hidden_units_}, decoder_normed_input_}},
        {"finished", input_tensors->at(3)},
        {"sequence_lengths", input_tensors->at(5)},
        {"step", input_tensors->at(4)}};

    TensorMap ffn_input_tensors(
        {{"ffn_input", Tensor{MEMORY_GPU, data_type, {batch_size, hidden_units_}, normed_cross_attn_output_}}});

    for (uint l = 0; l < num_layer_; l++) {
        const T* decoder_input = (const T*)((l == 0) ? input_tensors->at(0).getPtr<const T>() : decoder_layer_output_);
        T* decoder_output = (T*)((l == num_layer_ - 1) ? output_tensors->at(0).getPtr<T>() : decoder_layer_output_);

        size_t self_key_cache_offset = l;
        for (auto t = output_tensors->at(1).shape.begin() + 1; t != output_tensors->at(1).shape.end(); ++t) {
            self_key_cache_offset *= (*t);
        }
        size_t self_value_cache_offset = l;
        for (auto t = output_tensors->at(2).shape.begin() + 1; t != output_tensors->at(2).shape.end(); ++t) {
            self_value_cache_offset *= (*t);
        }