fix: fix guided decoding state corruption in turbomind when tp>1

Author: windreamer

src/turbomind/layers/BaseDynamicDecodeLayer.h

@@ -31,6 +31,7 @@ class BaseDynamicDecodeLayer {
         int                   vocab_size_padded;
         cudaStream_t          stream;
         const cudaDeviceProp* device_prop;
+        int                   tp_rank;
     };
 
     virtual ~BaseDynamicDecodeLayer() = default;
@@ -42,6 +43,7 @@ class BaseDynamicDecodeLayer {
         vocab_size_padded_ = param.vocab_size_padded;
         stream_            = param.stream;
         device_prop_       = param.device_prop;
+        tp_rank_           = param.tp_rank;
     };
 
     virtual void Setup(const std::vector<const Request*>& rs, const TensorMap& args) = 0;
@@ -54,6 +56,7 @@ class BaseDynamicDecodeLayer {
     int                   vocab_size_padded_;
     cudaStream_t          stream_;
     const cudaDeviceProp* device_prop_;
+    int                   tp_rank_;
 };
 
 }  // namespace turbomind

src/turbomind/layers/DynamicDecodeLayer.cc

@@ -31,15 +31,22 @@ DynamicDecodeLayer::DynamicDecodeLayer(DataType              dtype,
                                        int                   vocab_size,
                                        int                   vocab_size_padded,
                                        cudaStream_t          stream,
-                                       const cudaDeviceProp* device_prop)
+                                       const cudaDeviceProp* device_prop,
+                                       int                   tp_rank):
+    tp_rank_(tp_rank)
 {
     TM_LOG_DEBUG(__PRETTY_FUNCTION__);
     TM_CHECK(dtype == kFloat32);
-    BaseDynamicDecodeLayer::BaseParam param{max_batch_size, vocab_size, vocab_size_padded, stream, device_prop};
+    BaseDynamicDecodeLayer::BaseParam param{
+        max_batch_size, vocab_size, vocab_size_padded, stream, device_prop, tp_rank};
     layers_.emplace_back(new LogitsProcessorLayer<float>{param});
-    layers_.emplace_back(new GuidedDecodeMaskLayer<float>{param});
+    if (tp_rank == 0) {
+        layers_.emplace_back(new GuidedDecodeMaskLayer<float>{param});
+    }
     layers_.emplace_back(new SamplingLayer<float>{param});
-    layers_.emplace_back(new GuidedDecodeUpdateLayer<float>{param});
+    if (tp_rank == 0) {
+        layers_.emplace_back(new GuidedDecodeUpdateLayer<float>{param});
+    }
     layers_.emplace_back(new StopCriteriaLayer<float>{param});
 }
 
@@ -48,6 +55,7 @@ DynamicDecodeLayer::~DynamicDecodeLayer() {}
 void DynamicDecodeLayer::Setup(const std::vector<const Request*>& rs, const TensorMap& args)
 {
     TM_LOG_DEBUG(__PRETTY_FUNCTION__);
+    TM_LOG_INFO("[Dynamic Decode] Setup layer for %d", tp_rank_);
     for (const auto& layer : layers_) {
         layer->Setup(rs, args);
     }
@@ -82,6 +90,7 @@ void DynamicDecodeLayer::Forward(TensorMap& args)
      *   \param  sampled_nums [batch_size, 1], optional
      */
 
+    TM_LOG_INFO("[Dynamic Decode] Forward for %d", tp_rank_);
     for (const auto& layer : layers_) {
         layer->Forward(args);
     }

src/turbomind/layers/DynamicDecodeLayer.h

@@ -33,7 +33,8 @@ class DynamicDecodeLayer {
                        int                   vocab_size,
                        int                   vocab_size_padded,
                        cudaStream_t          stream,
-                       const cudaDeviceProp* device_prop);
+                       const cudaDeviceProp* device_prop,
+                       int                   tp_rank);
 
     ~DynamicDecodeLayer();
 
@@ -42,6 +43,7 @@ class DynamicDecodeLayer {
     void Forward(TensorMap& args);
 
 private:
+    int                                                  tp_rank_;
     std::vector<std::unique_ptr<BaseDynamicDecodeLayer>> layers_;
 };
 

src/turbomind/models/llama/LlamaV2.cc

@@ -90,7 +90,7 @@ LlamaV2::LlamaV2(DataType                     dtype,
 
     // using float to avoid data overflow
     dynamic_decode_ = std::make_unique<DynamicDecodeLayer>(
-        kFloat32, max_batch_size, vocab_size_, vocab_size_padded_, stream_, &ctx.device_prop);
+        kFloat32, max_batch_size, vocab_size_, vocab_size_padded_, stream_, &ctx.device_prop, tp_rank_);
 }
 
 void LlamaV2::updateEmbedding(char*            decoder_input,