[None][fix] Address code review: API safety, assertion, and comment accuracy

- Replace getFragmentPointerDevice()/getFragmentPointerHost() with a single
  getFragmentPointerSlot() that returns both rows atomically and advances workIdx
  once, making it impossible to call them in the wrong order (gripe 2).
- Add TLLM_CHECK that cacheBlockIds is non-empty on the cross-KV path before
  accessing cacheBlockIds[0] (gripe 5).
- Fix the Fix-3 description in encDecBeamSearchTest.cpp: this branch uses a
  repaired kernel (per-slot fragmentPointerDevice), not the direct-copy path
  described by the stale comment (gripe 3).

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

Author: achartier

cpp/include/tensorrt_llm/batch_manager/runtimeBuffers.h

@@ -224,17 +224,14 @@ class RuntimeBuffers
             workIdx = (workIdx + 1) % (fragmentPointerHost->getShape().d[0]);
         }
 
-        [[nodiscard]] TensorPtr getFragmentPointerHost()
+        //! Returns matching host and device pointer rows for the current workIdx, then advances
+        //! workIdx.  Always call this instead of the individual getters to avoid ordering bugs.
+        [[nodiscard]] std::pair<TensorPtr, TensorPtr> getFragmentPointerSlot()
         {
-            TensorPtr slice = runtime::ITensor::slice(fragmentPointerHost, workIdx, 1);
+            TensorPtr host = runtime::ITensor::slice(fragmentPointerHost, workIdx, 1);
+            TensorPtr device = runtime::ITensor::slice(fragmentPointerDevice, 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);
+            return {std::move(host), std::move(device)};
         };
     };
 

cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp

@@ -3948,6 +3948,8 @@ SizeType32 KVCacheManager::copyBlockOffsets(ITensor& output, SizeType32 outputSl
                 // For cross-KV (encoder features), all beams of a request share the same encoder output.
                 // Always use beam 0's block IDs so all beams attend to the correct encoder features.
                 auto const srcBeamIdx = isCrossKv() ? 0 : beamIdx;
+                TLLM_CHECK_WITH_INFO(!isCrossKv() || !cacheBlockIds.empty(),
+                    "Cross-KV sequence has no block IDs for request %lu", requestId);
                 auto const effectiveBlockCount = isCrossKv() ? cacheBlockIds[0].size() : cacheBlockIds[beamIdx].size();
                 auto const copyChunkSize = effectiveBlockCount * sizeof(tk::KVCacheIndex);
                 for (auto xIdx : {kIdx, vIdx})

cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp

@@ -103,14 +103,11 @@ void copyGenerationLogits(RuntimeBuffers::GenerationLogitsCache& generationLogit
 
     auto const fragmentSize = llmReq.getGenerationLogitsFragmentsSize();
 
-    // 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.
+    // Merge logits fragments on device.  getFragmentPointerSlot() returns the matching host and
+    // device rows for the current workIdx and advances the index atomically, so concurrent flushes
+    // for different requests in the same batch never clobber each other's pointer arrays.
     auto const& transposeBufferPtr = generationLogitsCache.transposedLogits;
-    auto const cachePointerDevice = generationLogitsCache.getFragmentPointerDevice();
-    auto const cachePointerHost = generationLogitsCache.getFragmentPointerHost();
+    auto [cachePointerHost, cachePointerDevice] = generationLogitsCache.getFragmentPointerSlot();
     tensorrt_llm::runtime::kernels::mergeLogitsFragments(bufferManager, *transposeBufferPtr,
         llmReq.getGenerationLogitsFragments(), *cachePointerDevice, *cachePointerHost, 0, 1, reqBeamWidth,
         bufferManager.getStream(), 0);

cpp/tests/unit_tests/batch_manager/encDecBeamSearchTest.cpp

@@ -25,12 +25,13 @@
 //   beam when isCrossKv() is true, since the encoder output is identical for
 //   all beams of a request.
 //
-// Fix 3 — copyGenerationLogits direct-copy places each (beam, step) logit at
-//   the correct host offset.
-//   Before: the mergeLogitsFragmentsKernel pointer-indirection caused
-//   intermittent GPU-state corruption when gather_generation_logits=True was
-//   combined with concurrent mixed beam-width requests.
-//   After:  each fragment is copied directly to the host without the kernel.
+// Fix 3 — fragmentPointerDevice is now per-batch-slot ([maxBatchSize, kCACHE_LENGTH]).
+//   Before: fragmentPointerDevice was a single shared row ([kCACHE_LENGTH]); sequential
+//   flushes from different requests in the same batch clobbered each other's GPU pointer
+//   arrays, causing the mergeLogitsFragmentsKernel to read stale fragment addresses and
+//   produce degenerate output with gather_generation_logits=True.
+//   After:  each request gets its own device-side pointer row via getFragmentPointerSlot(),
+//   eliminating cross-request interference.
 
 #include "tensorrt_llm/batch_manager/kvCacheManager.h"
 #include "tensorrt_llm/batch_manager/llmRequest.h"