Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
25 commits
Select commit Hold shift + click to select a range
603c8a0
BCG V1
GuanhuaWang2001 Jul 20, 2026
fe93728
bcg v2, support more model, fix adp bug and dsv4 bug
GuanhuaWang2001 Jul 21, 2026
3838856
[07/23/14:03] share BCG pool across segments
GuanhuaWang2001 Jul 23, 2026
73d4da9
[07/23/16:37] benchmark DSV4 1P1D AAAgent with BCG
GuanhuaWang2001 Jul 23, 2026
f7b3e9f
[07/24/20:09] refactor DSv4 BCG epilogue fusion
GuanhuaWang2001 Jul 24, 2026
fbb0147
fix bcg and epilogue
GuanhuaWang2001 Jul 27, 2026
980525e
rename some var name to make mla more clear
GuanhuaWang2001 Jul 28, 2026
2f46386
[None][test] add focused BCG accuracy coverage
GuanhuaWang2001 Jul 28, 2026
73fb1ed
[None][chore] apply pre-commit fixes to BCG changes
GuanhuaWang2001 Jul 29, 2026
8897e87
[None][fix] preserve PCG behavior with breakable CUDA graphs
GuanhuaWang2001 Jul 29, 2026
06f1313
[None][test] remove DeepSeek V4 BCG accuracy tests
GuanhuaWang2001 Jul 30, 2026
6dacfc5
[None][fix] repair BCG tests after rebase
GuanhuaWang2001 Aug 3, 2026
e4c21f1
[None][fix] avoid implicit legacy PCG bucket conflicts
GuanhuaWang2001 Aug 4, 2026
557406f
[None][fix] restrict eager graph captured values
GuanhuaWang2001 Aug 4, 2026
12a879e
reorganize tests
GuanhuaWang2001 Aug 4, 2026
5aff7d3
[None][test] provide valid prepared graph inputs
GuanhuaWang2001 Aug 5, 2026
89fe680
[None][fix] keep prefill warmup ranks aligned
GuanhuaWang2001 Aug 5, 2026
b401fd5
[None][fix] simplify breakable CUDA graph helpers
GuanhuaWang2001 Aug 5, 2026
9a3ff89
chore: apply post-rebase formatting
GuanhuaWang2001 Aug 8, 2026
ab3d2b7
[None][fix] address breakable CUDA graph review feedback
GuanhuaWang2001 Aug 10, 2026
f328f84
[None][fix] update prefill graph flag after rebase
GuanhuaWang2001 Aug 10, 2026
5ad889c
[None][fix] repair breakable graph CI regressions
GuanhuaWang2001 Aug 10, 2026
e4d0e14
[None][fix] retain eager bridge output storage
GuanhuaWang2001 Aug 11, 2026
6410451
[None][fix] clean up failed BCG capture
GuanhuaWang2001 Aug 12, 2026
e016428
[None][docs] clarify DeepSeek-V4 output contract
GuanhuaWang2001 Aug 12, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 34 additions & 5 deletions docs/source/features/torch_compile_and_piecewise_cuda_graph.md
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# Torch Compile & Piecewise CUDA Graph
# Torch Compile & Prefill CUDA Graph

In this guide, we show how to enable torch.compile and Piecewise CUDA Graph in TensorRT LLM. TensorRT LLM uses torch.compile for lightweight vertical fusion and Piecewise CUDA Graph.

Expand Down Expand Up @@ -41,12 +41,40 @@ To enable torch.compile and Piecewise CUDA Graph, add the following configuratio

```yaml
... # Other extra config
prefill_cuda_graph_backend: piecewise
prefill_capture_num_tokens: '${capture_num_tokens}' # e.g. [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, ..., 3072]
torch_compile_config:
capture_num_tokens: '${capture_num_tokens}' # List of num tokens to capture. e.g., [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, ..., 3072]
enable_userbuffers: false
enable_piecewise_cuda_graph: true
```

`TorchCompileConfig.enable_piecewise_cuda_graph` and
`TorchCompileConfig.capture_num_tokens` are deprecated aliases for these
prefill-specific options.

The experimental breakable implementation can capture the model body without
torch.compile:

```yaml
prefill_cuda_graph_backend: breakable
Comment thread
GuanhuaWang2001 marked this conversation as resolved.
prefill_capture_num_tokens: [128, 256, 512]
```

The breakable backend is experimental. The integration coverage in this change
includes BF16 Qwen3.5 on one GPU and NVFP4 DeepSeek models on multiple GPUs,
with both context-only and mixed context/decode batches using the KV cache.

The following restrictions are enforced:

- `torch_compile_config`, LoRA, and multimodal models are rejected during
engine initialization.
- Speculative decoding is supported.
- Context-logit requests run eagerly instead of replaying a breakable CUDA
graph.

Other model families, quantization modes, and parallel configurations are not
yet covered by this experimental backend's integration tests and should be
validated before use.

## Tips for Piecewise CUDA Graph

### Piecewise CUDA Graph & Generation Only CUDA Graph
Expand All @@ -59,9 +87,10 @@ cuda_graph_config:
max_batch_size: 1024 # Specify max capture batch size for generation only cuda graph. By default, TensorRT LLM will generate a capture list based on it.

torch_compile_config:
capture_num_tokens: '${capture_num_tokens}' # Specify capture_num_tokens for piecewise cuda graph
enable_userbuffers: false
enable_piecewise_cuda_graph: true

prefill_cuda_graph_backend: piecewise
prefill_capture_num_tokens: '${capture_num_tokens}'
```

### Piecewise CUDA Graph Padding
Expand Down
Loading
Loading