@@ -94,6 +94,8 @@ if(${TRITON_ENABLE_NVTX})
add_definitions(-DTRITON_ENABLE_NVTX=1)
endif() # TRITON_ENABLE_NVTX
+find_package(XSched REQUIRED)
+
#
# Shared library implementing the Triton Backend API
#
@@ -197,6 +199,13 @@ target_link_libraries(
CUDA::cudart
)
+target_link_libraries(
+ triton-tensorrt-backend
+ PRIVATE
+ XSched::preempt
+ XSched::halcuda
+)
+
#
# Install
@@ -45,6 +45,8 @@
#include <set>
#include <thread>
#include <unordered_map>
+#include "xsched/cuda/hal.h"
+#include "xsched/xsched.h"
//
// TensorRT Backend that implements the TRITONBACKEND API.
@@ -2926,6 +2928,15 @@ ModelInstanceState::EvaluateTensorRTContext(
return nullptr;
}
+void SetStreamXQueuePriority(CUstream stream, int priority)
+{
+ HwQueueHandle hwqueue;
+ CudaQueueCreate(&hwqueue, (CUstream)stream);
+ XQueueHandle xqueue;
+ XQueueCreate(&xqueue, hwqueue, kPreemptLevelDeactivate, kQueueCreateFlagNone);
+ XHintPriority(xqueue, -priority); // In XSched, lower number means lower priority
+}
+
TRITONSERVER_Error*
ModelInstanceState::InitStreamsAndEvents()
{
@@ -2961,15 +2972,19 @@ ModelInstanceState::InitStreamsAndEvents()
CreateCudaStream(DeviceId(), cuda_stream_priority_, &stream_));
}
}
+ SetStreamXQueuePriority(stream_, cuda_stream_priority_);
#ifdef TRITON_ENABLE_STATS
RETURN_IF_ERROR(
CreateCudaStream(DeviceId(), cuda_stream_priority_, &signal_stream_));
+ SetStreamXQueuePriority(signal_stream_, cuda_stream_priority_);
#endif // TRITON_ENABLE_STATS
RETURN_IF_ERROR(
CreateCudaStream(DeviceId(), cuda_stream_priority_, &input_copy_stream_));
+ SetStreamXQueuePriority(input_copy_stream_, cuda_stream_priority_);
if (model_state_->SeparateOutputStream()) {
RETURN_IF_ERROR(CreateCudaStream(
DeviceId(), cuda_stream_priority_, &output_copy_stream_));
+ SetStreamXQueuePriority(output_copy_stream_, cuda_stream_priority_);
}
// Create CUDA events associated with the execution states
RETURN_IF_ERROR(InitEventSet(model_state_->BusyWaitEvents()));