[Bug] Cuda Graph Issues with TRTLLM-GEN Backend with BatchDecodeWithPagedKVCacheWrapper

[NihalPotdar]

Hi! I was running into an issue with the BatchDecodeWithPagedKVCacheWrapper and trtllm-gen backend on B200 where it returns different results when played with cuda graphs then without.

I have attached a snippet of the cuda graph code as reference:

During init:

self._cuda_graph_kv_indptr = torch.zeros(
    max_concurrent_requests + 1,
    dtype=torch.int32,
    device=self.device,
)

self._cuda_graph_kv_indices = torch.zeros(
    max_concurrent_requests * max_num_pages,
    dtype=torch.int32,
    device=self.device,
)

self._cuda_graph_kv_last_page_len = torch.ones(
    max_concurrent_requests,
    dtype=torch.int32,
    device=self.device,
)

During capture:

kv_indices=self._cuda_graph_kv_indices
kv_indptr=self._cuda_graph_kv_indptr[:batch_size + 1]
kv_last_page_len=self._cuda_graph_kv_last_page_len[:batch_size]

Before cuda graph replay:

self._cuda_graph_kv_indices[:total_pages].copy_(attention_data.kv_indices[:total_pages])
self._cuda_graph_kv_indptr[:batch_size + 1].copy_(attention_data.kv_indptr[:batch_size + 1])
self._cuda_graph_kv_last_page_len[:batch_size].copy_(attention_data.kv_last_page_len[:batch_size])

Questions:

• I noticed that in sglang which directly calls these kernels and does not use the wrapper, the block table is passed in during the capture phase. Is that a requirement for this backend?

Environment

Kernel: BatchDecodeWithPagedKVCacheWrapper
Batch size: 8
GPU: SM100 (B200)

[saltyminty]

In the BatchDecodeWithPagedKVCacheWrapper API, block_tables is either passed in or automatically constructed during the plan() phase and cached for future replays. It seems you are updating kv_indices/kv_indptr/kv_last_page_len in-place, leading to the block table being out of sync (as it isn't updated without being manually updated or a plan() call.

Is my understanding correct?

[bkryu]

Every time plan() is called without an explicit block_tables argument, it allocates a brand-new tensor. Then run() passes that tensor to the kernel.

One situation @NihalPotdar you might be describing is, when we use CUDA graphs,

1. During graph capture: plan() is called → creates self._block_tables at memory address A. Then run() is captured into the graph, which bakes in address A.
2. During graph replay: plan() is called again (outside the graph) → if we give self._block_tables = None (since block_tables param defaults to None) → then the function allocates a new tensor at memory address B, fills it with the updated page table.
3. But the graph replays run() using the old address A — which has stale data. So the kernel reads the wrong block table.

I noticed that in sglang which directly calls these kernels and does not use the wrapper, the block table is passed in during the capture phase. Is that a requirement for this backend?

I believe the answer is yes. For CUDAGraph to work with the trtllm-gen backend, block_tables needs to be a fixed-address, pre-allocated buffer that gets updated in-place, just like indptr/indices/last_page_len.