[Bug][AutoDeploy]: Investigate _scratch_block_offsets in trtllm_attention.py
@govind-ramnarayan is already working on this.
Since Apr 16, 2026.
- Dominant language
- Python
- Stars
- 14.7k
- Forks
- 2.8k
- Avg merge
- 2d 23h
- Merged PRs (30d)
- 489
Description
Problem
We observed crashes without a defensive copy of block_offsets into a scratch buffer before each thop.attention call. This occurred with Llama + TRTLLM Attention + overlap scheduler ON + torch-cudagraph. Adding the workaround described below fixed it - we should investigate why this was.
The hypothesis is that metadata planning for an upcoming batch can write to shared buffers (like block_offsets) while a concurrent thop.attention call is still reading from them.
Current workaround
A _scratch_block_offsets buffer is allocated in _GlobalTrtllmPlanner.init_spec_decoding() and used in trtllm_mha_with_cache (gated on is_spec_decoding_enabled). See trtllm_attention.py.
TODO
- Understand the root cause of the race more deeply
- Determine if non-spec-dec paths (vanilla Llama + overlap scheduler) also need the scratch copy
- If so, move the copy out of the spec-dec gate — this may fix the known torch-simple + overlap scheduler crash. If not, diagnose the root cause and fix it instead of this scratch block offsets.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.