Skip to content

fix(trtllm-mla): make spec-decode CUDA graph capture causal#616

Open
mesaleh wants to merge 1 commit into
lightseekorg:mainfrom
mesaleh:fix/trtllm-mla-cudagraph-capture-rebased
Open

fix(trtllm-mla): make spec-decode CUDA graph capture causal#616
mesaleh wants to merge 1 commit into
lightseekorg:mainfrom
mesaleh:fix/trtllm-mla-cudagraph-capture-rebased

Conversation

@mesaleh

@mesaleh mesaleh commented Jul 9, 2026

Copy link
Copy Markdown
Contributor

Replaces #260 after it was automatically closed for inactivity. GitHub is not allowing the closed PR to be reopened through the API, so this keeps the same fix alive on a fresh rebased branch.

Summary

  • Keeps spec-decode CUDA graph capture seq_lens, output cache slots, and dummy KV page internally consistent during warmup and capture.
  • Uses the reserved padding page for TRT-LLM MLA CUDA graph capture block tables.
  • Flattens multi-token TRT-LLM MLA decode for both target verification and draft catch-up so each token gets its own causal seq_len bound.

Validation

  • python3 -m py_compile python/tokenspeed/runtime/execution/cuda_graph_wrapper.py python/tokenspeed/runtime/layers/attention/backends/trtllm_mla.py
  • git diff --check

Signed-off-by: Moustafa Saleh <moustafa.saleh@omniva.com>
@mesaleh mesaleh force-pushed the fix/trtllm-mla-cudagraph-capture-rebased branch from 23fa23c to 3993b94 Compare July 9, 2026 01:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant