phind-justin
commented
6 months ago I was looking into the XQA prepare and dispatch functions here https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.h#L185
it seems like the mPrepareCalled could be subject to a race condition, however I'm sure you have more context on the synchonization conditions than me. I can't tell what other mechanisms are in place to make mPrepareCalled thread-safe
void prepare(const XQAParams& xqa_params)
{
if (!mPrepareCalled)
{
this->prepareForRun(xqa_params);
mPrepareCalled = true;
}
}
template <typename KVCacheBuffer>
void dispatch(const XQAParams& xqa_params, KVCacheBuffer& kv_cache_buffer, const cudaStream_t& stream)
{
if (!mPrepareCalled)
{
TLLM_THROW("DecoderXQARunner::prepare() hasn't been called before DecoderXQARunner::dispatch().");
}
sync_check_cuda_error();
this->run(xqa_params, kv_cache_buffer, stream);
}
private:
bool shouldUseImpl(const XQAParams& xqa_params);
void prepareForRun(const XQAParams& xqa_params);
template <typename KVCacheBuffer>
void run(const XQAParams& xqa_params, KVCacheBuffer& kv_cache_buffer, const cudaStream_t& stream);
static constexpr int kMaxBeamWidth = 4;
// Cache the grid_size and block_size that gives the highest occupancy for
// invokeApplyBiasRopeUpdateKVCache.
int2 mLaunchGridBlockCache = make_int2(0, 0);
bool mPrepareCalled;
michaelroyzen
ming-wei
juney-nvidia
System Info
x86 8x h100 80g v0.8.0
Who can help?
@Tracin
Information
Tasks
examplesfolder (such as GLUE/SQuAD, ...)Reproduction
I then use triton to serve the model and send a few inferences
it seems to work for a few inferences, especially when there's only 1-2 at a time. But after a few inferences, or if we send 4 at a time, it doesn't work.
Expected behavior
it runs successfully and does not crash for batch sizes up to 32!
actual behavior
It crashes with an XQA error after a few inferences
codellama_70b_xqa.log
additional notes
.