Skip to content

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

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

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

Conversation

@mesaleh

@mesaleh mesaleh commented May 26, 2026

Copy link
Copy Markdown
Contributor

Summary

  • initialize spec-decode CUDA graph capture with internally consistent synthetic sequence lengths and cache slots
  • point TRT-LLM MLA capture block tables at the same synthetic per-request pages and zero those pages before capture
  • flatten multi-token TRT-LLM MLA decode into per-token decode queries so target verification and draft catch-up use causal sequence bounds per token

Motivation

While bringing up Kimi K2.6 NVFP4 with EAGLE3 speculative decoding and --attention-backend trtllm_mla, CUDA graph startup could enter an impossible synthetic state: multi-token decode capture used grouped metadata that did not match the dummy KV/cache locations. This led to capture-time instability and downstream invalid draft ids.

This PR is intentionally independent of #217. It does not touch the EAGLE first-step reduce plumbing and merge-tests cleanly with pull/217/head.

Validation

  • python3 -m compileall -q python/tokenspeed/runtime/execution/cuda_graph_wrapper.py python/tokenspeed/runtime/layers/attention/backends/trtllm_mla.py
  • Local GB200 validation: Kimi K2.6 NVFP4 + EAGLE3 + trtllm_mla captured batch sizes [1, 2, 3, 4, 5, 6, 7, 8] and reached healthy readiness
  • OpenAI-compatible synthetic benchmark completed 10/10 requests with p50 TPOT 3.760 ms and mean decode 264.676 tok/s

Signed-off-by: Moustafa Saleh <moustafa.saleh@omniva.com>
@mesaleh
mesaleh force-pushed the omniva/fix-trtllm-mla-cudagraph-capture branch from 8d57363 to 476a102 Compare May 26, 2026 09:51
@mesaleh mesaleh closed this May 26, 2026
@mesaleh
mesaleh deleted the omniva/fix-trtllm-mla-cudagraph-capture branch May 26, 2026 10:10
@mesaleh
mesaleh restored the omniva/fix-trtllm-mla-cudagraph-capture branch May 26, 2026 10:13
@mesaleh mesaleh reopened this May 26, 2026
@mesaleh mesaleh closed this May 26, 2026
@mesaleh
mesaleh deleted the omniva/fix-trtllm-mla-cudagraph-capture branch May 26, 2026 10:13
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