[None][fix] Fix gather_generation_logits corruption: per-slot fragmentPointerDevice

The mergeLogitsFragmentsKernel reads fragment GPU addresses from fragmentPointerDevice
([kCACHE_LENGTH]).  When multiple requests in the same batch flush their logits fragments
sequentially, each overwrites that shared device buffer before the previous kernel has
committed its result on the stream — a race that caused intermittent degenerate output
("happ happ happ") with gather_generation_logits=True.

Fix: resize fragmentPointerDevice from [kCACHE_LENGTH] to [maxBatchSize, kCACHE_LENGTH],
matching the layout of fragmentPointerHost.  Add getFragmentPointerDevice() to return the
current-workIdx row without advancing it, and call it before getFragmentPointerHost() so
both use the same per-slot row.  Each request in the batch now has its own device-side
pointer array, eliminating cross-request clobbering.

Signed-off-by: Aurelien Chartier <2567591+achartier@users.noreply.github.com>

Author: achartier

cpp/include/tensorrt_llm/batch_manager/runtimeBuffers.h

@@ -208,7 +208,9 @@ class RuntimeBuffers
         //! Temporarily store the transposed results of multiple fragment logits, [maxBeamWidth, kCACHE_LENGTH]
         TensorPtr transposedLogits;
 
-        //! Temporarily store logits buffer address during the transposing, [kCACHE_LENGTH]
+        //! Temporarily store logits buffer address during the transposing, [maxBatchSize, kCACHE_LENGTH]
+        //! One row per batch slot (same layout as fragmentPointerHost) so concurrent flushes for
+        //! different requests in the same batch never clobber each other's pointer arrays.
         TensorPtr fragmentPointerDevice;
 
         //! Temporarily store logits buffer address during the transposing, [maxBatchSize, kCACHE_LENGTH]
@@ -227,6 +229,12 @@ class RuntimeBuffers
             TensorPtr slice = runtime::ITensor::slice(fragmentPointerHost, workIdx, 1);
             cycleWorkIdx();
             return slice;
+        }
+
+        //! Returns the device pointer row for the current workIdx (called before cycleWorkIdx).
+        [[nodiscard]] TensorPtr getFragmentPointerDevice()
+        {
+            return runtime::ITensor::slice(fragmentPointerDevice, workIdx, 1);
         };
     };
 

cpp/tensorrt_llm/batch_manager/runtimeBuffers.cpp

@@ -152,8 +152,8 @@ void RuntimeBuffers::create(SizeType32 maxBatchSize, SizeType32 maxBeamWidth,
             ITensor::makeShape({GenerationLogitsCache::kCACHE_LENGTH, maxBatchSize * maxBeamWidth, vocabSizePadded}),
             logitsType);
 
-        generationLogitsCache.fragmentPointerDevice
-            = manager.gpu(ITensor::makeShape({GenerationLogitsCache::kCACHE_LENGTH}), nvinfer1::DataType::kINT64);
+        generationLogitsCache.fragmentPointerDevice = manager.gpu(
+            ITensor::makeShape({maxBatchSize, GenerationLogitsCache::kCACHE_LENGTH}), nvinfer1::DataType::kINT64);
         generationLogitsCache.fragmentPointerHost = tensorrt_llm::runtime::BufferManager::pinnedPool(
             ITensor::makeShape({maxBatchSize, GenerationLogitsCache::kCACHE_LENGTH}), nvinfer1::DataType::kINT64);
     }

cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp

@@ -103,10 +103,14 @@ void copyGenerationLogits(RuntimeBuffers::GenerationLogitsCache& generationLogit
 
     auto const fragmentSize = llmReq.getGenerationLogitsFragmentsSize();
 
-    // Merge logits fragments on device
+    // Merge logits fragments on device.
+    // getFragmentPointerDevice() reads the current workIdx row; getFragmentPointerHost() reads the
+    // same row then advances workIdx.  Calling device-slot first preserves the pairing so both
+    // host and device pointer arrays use the same per-batch-slot row, preventing concurrent
+    // flushes for different requests from clobbering each other's pointer arrays.
     auto const& transposeBufferPtr = generationLogitsCache.transposedLogits;
-    auto const& cachePointerDevice = generationLogitsCache.fragmentPointerDevice;
-    auto const& cachePointerHost = generationLogitsCache.getFragmentPointerHost();
+    auto const cachePointerDevice = generationLogitsCache.getFragmentPointerDevice();
+    auto const cachePointerHost = generationLogitsCache.getFragmentPointerHost();
     tensorrt_llm::runtime::kernels::mergeLogitsFragments(bufferManager, *transposeBufferPtr,
         llmReq.getGenerationLogitsFragments(), *cachePointerDevice, *cachePointerHost, 0, 1, reqBeamWidth,
         bufferManager.getStream(), 0);