Building a Custom Triton Inference Server C++ Backend for Sequence-to-Sequence Models — code snippets

snippet-1.txt

// triton_seq2seq_backend.h
#pragma once

#include <memory>
#include <string>
#include <vector>
#include <cuda_runtime.h>
#include "triton/backend/backend_common.h"
#include "triton/backend/backend_model.h"
#include "triton/backend/backend_model_instance.h"

namespace triton { namespace backend { namespace seq2seq {

// ModelState manages configuration and weight resources shared across instances.
class ModelState : public ModelStateBase {
 public:
  static TRITONSERVER_Error* Create(
      TRITONBACKEND_Model* triton_model, ModelState** state);
  virtual ~ModelState() = default;

  const std::string& EnginePath() const { return engine_path_; }

 private:
  ModelState(TRITONBACKEND_Model* triton_model);
  std::string engine_path_;
};

// ModelInstanceState handles the active execution loop on a specific GPU.
class ModelInstanceState : public ModelInstance {
 public:
  static TRITONSERVER_Error* Create(
      ModelState* model_state, TRITONBACKEND_ModelInstance* triton_model_instance,
      ModelInstanceState** state);
  virtual ~ModelInstanceState();

  // The primary entrypoint executed by Triton's scheduler threads
  TRITONSERVER_Error* Execute(
      std::vector<TRITONBACKEND_Request*>& requests, const uint32_t request_count);

 private:
  ModelInstanceState(
      ModelState* model_state, TRITONBACKEND_ModelInstance* triton_model_instance);

  TRITONSERVER_Error* InitCudaStream();

  ModelState* model_state_;
  cudaStream_t stream_;
  void* device_workspace_;
  size_t workspace_size_bytes_;
};

}}} // namespace triton::backend::seq2seq

snippet-2.txt

// triton_seq2seq_backend.cc
#include "triton_seq2seq_backend.h"

namespace triton { namespace backend { namespace seq2seq {

TRITONSERVER_Error* ModelInstanceState::Execute(
    std::vector<TRITONBACKEND_Request*>& requests, const uint32_t request_count) {
  
  std::vector<const int32_t*> input_ids_buffers(request_count, nullptr);
  std::vector<int64_t> input_ids_lengths(request_count, 0);
  std::vector<int32_t> max_gen_lengths(request_count, 128); // Standard production default

  for (size_t i = 0; i < request_count; ++i) {
    TRITONBACKEND_Request* request = requests[i];

    // Fetch the primary token sequence tensor
    TRITONBACKEND_Input* input_ids_tensor = nullptr;
    RETURN_IF_TRITONSERVER_ERROR(TRITONBACKEND_RequestInput(request, "INPUT_IDS", &input_ids_tensor));

    const void* buffer = nullptr;
    uint64_t byte_size = 0;
    TRITONSERVER_MemoryType memory_type;
    int64_t memory_type_id;
    RETURN_IF_TRITONSERVER_ERROR(TRITONBACKEND_InputBuffer(
        input_ids_tensor, 0 /* index */, &buffer, &byte_size, &memory_type, &memory_type_id));

    // For tokenizer alignment, we expect the input token IDs on host memory
    if (memory_type == TRITONSERVER_MEMORY_GPU) {
      return TRITONSERVER_ErrorNew(
          TRITONSERVER_ERROR_INVALID_ARG, "INPUT_IDS tensor must reside in host memory");
    }

    input_ids_buffers[i] = reinterpret_cast<const int32_t*>(buffer);
    input_ids_lengths[i] = byte_size / sizeof(int32_t);

    // Extract optional generation metadata
    TRITONBACKEND_Input* max_len_tensor = nullptr;
    TRITONSERVER_Error* err = TRITONBACKEND_RequestInput(request, "MAX_GEN_LEN", &max_len_tensor);
    if (err == nullptr && max_len_tensor != nullptr) {
      const void* len_buffer = nullptr;
      uint64_t len_byte_size = 0;
      TRITONSERVER_MemoryType len_mem_type;
      int64_t len_mem_id;
      RETURN_IF_TRITONSERVER_ERROR(TRITONBACKEND_InputBuffer(
          max_len_tensor, 0, &len_buffer, &len_byte_size, &len_mem_type, &len_mem_id));
      max_gen_lengths[i] = *reinterpret_cast<const int32_t*>(len_buffer);
    } else if (err != nullptr) {
      // Clear error struct if the parameter was simply omitted
      TRITONSERVER_ErrorDelete(err);
    }
  }

  // batch execution execution path follows...
  return nullptr;
}

}}} // namespace triton::backend::seq2seq

snippet-3.txt

// memory_utils.cc
#include <triton/backend/backend_common.h>
#include <cuda_runtime.h>
#include <vector>

namespace triton { namespace backend { namespace seq2seq {

TRITONSERVER_Error* AllocateOutputOnDevice(
    TRITONBACKEND_Response* response, const char* tensor_name,
    TRITONSERVER_DataType data_type, const std::vector<int64_t>& shape,
    size_t expected_bytes, cudaStream_t stream, void** device_ptr) {
  
  TRITONBACKEND_Output* output_tensor = nullptr;
  RETURN_IF_TRITONSERVER_ERROR(TRITONBACKEND_ResponseOutput(
      response, &output_tensor, tensor_name, data_type, shape.data(), shape.size()));

  void* allocated_buffer = nullptr;
  TRITONSERVER_MemoryType memory_type = TRITONSERVER_MEMORY_GPU;
  int64_t memory_type_id = 0; // Target device GPU 0

  // Triton maps the buffer. If configured, this yields direct GPU pointers
  RETURN_IF_TRITONSERVER_ERROR(TRITONBACKEND_OutputBuffer(
      output_tensor, &allocated_buffer, expected_bytes, &memory_type, &memory_type_id));

  if (memory_type == TRITONSERVER_MEMORY_GPU) {
    // Zero-copy path: directly reference the device memory block
    *device_ptr = allocated_buffer;
  } else {
    // Fail immediately if Triton allocates Host memory, as fallback copies degrade p99s
    return TRITONSERVER_ErrorNew(
        TRITONSERVER_ERROR_INTERNAL, 
        "Direct GPU output allocation failed; Triton returned host memory instead.");
  }

  return nullptr;
}

}}} // namespace triton::backend::seq2seq

snippet-4.txt

// generator.cc
#include <triton/backend/backend_common.h>
#include <cuda_runtime.h>
#include <vector>

namespace triton { namespace backend { namespace seq2seq {

// Interface defining interactions with our underlying inference engine (e.g. TensorRT C++)
class ModelExecutor {
 public:
  virtual TRITONSERVER_Error* RunEncoder(
      const int32_t* input_ids, size_t length, cudaStream_t stream, void* enc_hidden_states) = 0;
  virtual TRITONSERVER_Error* RunDecoderStep(
      int32_t current_token, void* enc_hidden_states, cudaStream_t stream, int32_t* next_token) = 0;
};

TRITONSERVER_Error* ProcessStreamingRequest(
    TRITONBACKEND_Request* request, ModelExecutor* executor,
    const int32_t* input_ids, size_t length, int32_t max_gen_len,
    int32_t eos_token_id, cudaStream_t stream) {

  // Allocate dynamic memory on GPU for encoder sequence representations
  void* dev_enc_hidden_states = nullptr;
  cudaError_t cuda_err = cudaMallocAsync(&dev_enc_hidden_states, 2048 * sizeof(float), stream);
  if (cuda_err != cudaSuccess) {
    return TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL, cudaGetErrorString(cuda_err));
  }

  // Execute Encoder once
  RETURN_IF_TRITONSERVER_ERROR(executor->RunEncoder(input_ids, length, stream, dev_enc_hidden_states));

  // Initialize auto-regressive generation targeting the final input token
  int32_t current_token = input_ids[length - 1]; 
  std::vector<int64_t> output_shape = {1};

  for (int32_t step = 0; step < max_gen_len; ++step) {
    int32_t next_token = 0;
    
    // Execute a single decoder token step
    RETURN_IF_TRITONSERVER_ERROR(executor->RunDecoderStep(
        current_token, dev_enc_hidden_states, stream, &next_token));
    
    // Create an independent streaming response container
    TRITONBACKEND_Response* token_response = nullptr;
    RETURN_IF_TRITONSERVER_ERROR(TRITONBACKEND_ResponseNew(&token_response, request));

    TRITONBACKEND_Output* output_tensor = nullptr;
    RETURN_IF_TRITONSERVER_ERROR(TRITONBACKEND_ResponseOutput(
        token_response, &output_tensor, "OUTPUT_IDS", TRITONSERVER_TYPE_INT32,
        output_shape.data(), output_shape.size()));

    void* out_buffer = nullptr;
    TRITONSERVER_MemoryType out_mem_type = TRITONSERVER_MEMORY_CPU;
    int64_t out_mem_id = 0;
    
    // Bind token storage to CPU-pinned host memory to guarantee fast asynchronous transfer
    RETURN_IF_TRITONSERVER_ERROR(TRITONBACKEND_OutputBuffer(
        output_tensor, &out_buffer, sizeof(int32_t), &out_mem_type, &out_mem_id));

    *reinterpret_cast<int32_t*>(out_buffer) = next_token;

    // Check generation stop limits
    bool is_last = (next_token == eos_token_id) || (step == max_gen_len - 1);
    uint32_t send_flags = is_last ? TRITONSERVER_RESPONSE_COMPLETE_FINAL : TRITONSERVER_RESPONSE_COMPLETE_NONE;

    // Asynchronously dispatch token to Triton client transport layer
    RETURN_IF_TRITONSERVER_ERROR(TRITONBACKEND_ResponseSend(token_response, send_flags, nullptr));

    current_token = next_token;
    if (is_last) {
      break;
    }
  }

  cudaFreeAsync(dev_enc_hidden_states, stream);
  return nullptr;
}

}}} // namespace triton::backend::seq2seq

snippet-5.txt

# config.pbtxt for Seq2Seq Custom C++ Backend
name: "seq2seq_cpp_backend"
backend: "seq2seq_cpp"
max_batch_size: 64

model_transaction_policy {
  decoupled: true
}

input [
  {
    name: "INPUT_IDS"
    data_type: TYPE_INT32
    dims: [ -1 ]
  },
  {
    name: "MAX_GEN_LEN"
    data_type: TYPE_INT32
    dims: [ 1 ]
  }
]

output [
  {
    name: "OUTPUT_IDS"
    data_type: TYPE_INT32
    dims: [ 1 ]
  }
]

instance_group [
  {
    count: 2
    kind: KIND_GPU
    gpus: [ 0 ]
  }
]

parameters: {
  key: "engine_path"
  value: {
    string_value: "/models/seq2seq_cpp_backend/1/model.bin"
  }
}