[AutoDeploy] Flashinfer+trtllm IMA for (max_batch_size, max_seq_len)

When using trtllm runtime with flashinfer backend we get an illegal memory access when calling the model with input of shape (max_batch_size, max_seq_len), i.e., when we fully occupy the cache.

Interestingly, this doesn't happen with demollm. My current best guess is that there is an index offset by 1 in how trtllm runtime allocates cache positions/pages. The simple heuristic in demollm doesn't seem to run into that issue.

Error log below.

Related issue: https://github.com/NVIDIA/TensorRT-LLM/pull/3686

  File "/tmp/torchinductor_lliebenwein/xb/cxboxvid4wm7kweqmr6ksstvwsnk53r6tqu2vuataqchxfmfjksp.py", line 1527, in call
    buf23 = torch.ops.attention.flashinfer_mha_with_cache.default(buf16, buf17, reinterpret_tensor(buf22, (1, s1, 8, 128), (1024*s1, 1024, 128, 1), 0), arg8_1, arg9_1, arg11_1, arg12_1, arg14_1, arg15_1, arg30_1, arg29_1, arg25_1, 0.08838834764831845, 1.0, 1.0)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 756, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_library/custom_ops.py", line 335, in backend_impl
    result = self._backend_fns[device_type](*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_compile.py", line 51, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 838, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_library/custom_ops.py", line 367, in wrapped_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/dev/TensorRT-LLM/tensorrt_llm/_torch/auto_deploy/custom_ops/flashinfer_attention.py", line 284, in flashinfer_mha_with_cache
    flashinfer.page.append_paged_kv_cache(
  File "/usr/local/lib/python3.12/dist-packages/flashinfer/page.py", line 426, in append_paged_kv_cache
    _append_paged_kv_cache_kernel(
  File "/usr/local/lib/python3.12/dist-packages/flashinfer/page.py", line 150, in _append_paged_kv_cache_kernel
    get_module_attr("append_paged_kv_cache")(
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 756, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: AppendPagedKVCache failed with error: an illegal memory access was encountered
[05/20/2025-16:55:53] [TRT-LLM] [E] Encountered an error in forward function: AppendPagedKVCache failed with error: an illegal memory access was encountered
terminate called after throwing an instance of 'c10::Error'
  what():  CUDA error: an illegal memory access was encountered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Exception raised from c10_cuda_check_implementation at /opt/pytorch/pytorch/c10/cuda/CUDAException.cpp:43 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0x98 (0x7ffce67315e8 in /usr/local/lib/python3.12/dist-packages/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) + 0xe0 (0x7ffce66c64a2 in /usr/local/lib/python3.12/dist-packages/torch/lib/libc10.so)
frame #2: c10::cuda::c10_cuda_check_implementation(int, char const*, char const*, int, bool) + 0x3c2 (0x7ffd55faa2a2 in /usr/local/lib/python3.12/dist-packages/torch/lib/libc10_cuda.so)
frame #3: <unknown function> + 0xb7d265 (0x7ffce7323265 in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_cuda.so)
frame #4: <unknown function> + 0xb794eb (0x7ffce731f4eb in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_cuda.so)
frame #5: <unknown function> + 0xb80c04 (0x7ffce7326c04 in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_cuda.so)
frame #6: <unknown function> + 0x44c162 (0x7ffd4a657162 in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so)
frame #7: c10::TensorImpl::~TensorImpl() + 0x9 (0x7ffce670bf39 in /usr/local/lib/python3.12/dist-packages/torch/lib/libc10.so)
frame #8: <unknown function> + 0x703468 (0x7ffd4a90e468 in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so)
frame #9: <unknown function> + 0x703890 (0x7ffd4a90e890 in /usr/local/lib/python3.12/dist-packages/torch/lib/libtorch_python.so)
frame #10: /usr/bin/python3() [0x579cf2]
frame #11: /usr/bin/python3() [0x59f0b9]
frame #12: /usr/bin/python3() [0x579d32]
frame #13: /usr/bin/python3() [0x59f0b9]
frame #14: _PyEval_EvalFrameDefault + 0x8f04 (0x5df844 in /usr/bin/python3)
frame #15: /usr/bin/python3() [0x54cd32]
frame #16: _PyEval_EvalFrameDefault + 0x4c1b (0x5db55b in /usr/bin/python3)
frame #17: /usr/bin/python3() [0x54cd32]
frame #18: /usr/bin/python3() [0x6f826c]
frame #19: /usr/bin/python3() [0x6b917c]
frame #20: <unknown function> + 0x9caa4 (0x7ffff7d0eaa4 in /usr/lib/x86_64-linux-gnu/libc.so.6)
frame #21: <unknown function> + 0x129c3c (0x7ffff7d9bc3c in /usr/lib/x86_64-linux-gnu/libc.so.6)