feat: Refactor LLM model zoo and add KV cache support #3527

peri044 · 2025-05-20T23:30:28Z

Description

This PR redesigns our LLM model compilation, unifies it, fixes output mismatch and performance issues. This PR also implements KV caching using native TensorRT.

Fixes # (issue)

Type of change

New feature (non-breaking change which adds functionality)
This change requires a documentation update

Checklist:

My code follows the style guidelines of this project (You can use the linters)
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas and hacks
I have made corresponding changes to the documentation
I have added tests to verify my fix or my feature
New and existing unit tests pass locally with my changes
I have added the relevant labels to my PR in so that relevant reviewers are notified

…s landed

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

narendasan · 2025-06-13T21:01:54Z

tools/llm/cache_utils.py

+from torch.utils._pytree import _LEAF_SPEC
+
+
+@torch_tensorrt.dynamo.conversion.dynamo_tensorrt_converter(


Do we want to upstream this eventually?

narendasan · 2025-06-13T21:03:57Z

tools/llm/sdpa_converter.py

@@ -62,25 +62,15 @@ def scaled_dot_product_attention(
 ) -> TRTTensor:


I would put the torch_trt extensions in a subdirectory all together. like //tools/llm/torch_trt_ext/

narendasan · 2025-06-13T21:04:44Z

tools/llm/test_qwen2.5_components.py

@@ -0,0 +1,193 @@
+import torch


Should we create a tests/py/tools/llm?

narendasan · 2025-06-13T21:10:20Z

tools/llm/cache_utils.py

+) -> Node:
+    """Add a graph input to the given GraphModule and return the newly created node.
+
+    NOTE: function does NOT do any graph canonicalization. This is left to the user!


If you dont want users to hit this directly consider prefixing the method name with _ like _add_graph_input

zewenli98

Mostly LGTM. just some questions

zewenli98 · 2025-06-14T00:45:57Z

py/torch_tensorrt/dynamo/_compiler.py

@@ -837,6 +859,7 @@ def contains_metadata(gm: torch.fx.GraphModule) -> bool:
            continue
        submodule_node_dict[node.name] = node

+    preserve_module_specs(original_in_spec, original_out_spec, partitioned_module)


Should set spec for partitioned_module's children?

zewenli98 · 2025-06-14T05:24:03Z

examples/dynamo/weight_streaming_example.py

-from utils import export_llm
+
+
+def export_llm(model, inputs, min_seq_len=1, max_seq_len=16):


why not directly call tools/llm/utils.py?

peri044 and others added 9 commits April 25, 2025 21:03

chore: kv cache prototyping

e5723db

chore: add sdpa converter/lowering

3e0b46a

feat: implement static/dynamic kv cache in Torch-TRT

e30fa42

chore: updates

a3a202f

chore: updates

3642614

chore: updates

3688630

chore: refactor updates

c9f5f27

chore: refactor updates

0cb0dcc

chore: updates

6cbb1bd

peri044 added the WIP Work is in progress, pull request should not be merged yet label May 20, 2025

facebook-github-bot added the cla signed label May 20, 2025

peri044 marked this pull request as draft May 20, 2025 23:30

github-actions bot requested a review from gs-olive May 20, 2025 23:30

peri044 changed the title ~~feat : caching attempts~~ feat: caching attempts May 20, 2025

peri044 removed the request for review from gs-olive May 21, 2025 00:40

Chengzhe Xu and others added 3 commits May 27, 2025 19:13

chore: updates

f539b55

chore: updates

0dc3a7e

chore: rebase with main

095b5cf

This comment was marked as spam.

Sign in to view

This comment was marked as outdated.

Sign in to view

peri044 added 2 commits June 4, 2025 09:15

feat: Refactor LLM runner and implemented support for Qwen family

600e363

chore: updates

9309725

peri044 force-pushed the kv_cache branch from 19bc80b to 9309725 Compare June 4, 2025 09:17

This comment was marked as outdated.

Sign in to view

chore: updates

a50e7ac

peri044 added 2 commits June 11, 2025 08:40

chore: remove conditional branching for causal attention

535c6a8

chore: remove conditional branching for causal attention

7b7ac04

This comment was marked as spam.

Sign in to view

chore: remove is_causal input now that causal attention enhancement i…

c1f0053

…s landed

This comment was marked as spam.

Sign in to view

chore: refactor

8301ee6

This comment was marked as spam.

Sign in to view

peri044 removed WIP Work is in progress, pull request should not be merged yet labels Jun 12, 2025

peri044 marked this pull request as ready for review June 12, 2025 22:36

This comment was marked as spam.

Sign in to view

chore: move code to tools/llm

ecf88d1

This comment was marked as spam.

Sign in to view

peri044 added 4 commits June 12, 2025 22:55

chore: remove files not considered for release

723d1b2

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

chore: rebase with main

b835ae9

chore: updates

5eabf65

chore: Add README.md

806616d

peri044 changed the title ~~feat: caching attempts~~ feat: Refactor LLM optimization and add KV cache support Jun 13, 2025

peri044 changed the title ~~feat: Refactor LLM optimization and add KV cache support~~ feat: Refactor LLM model zoo and add KV cache support Jun 13, 2025

chore: rebase with main

53a2c56

github-actions bot removed the component: lowering Issues re: The lowering / preprocessing passes label Jun 13, 2025

chore: add docs

c2cbd5a

github-actions bot added the documentation Improvements or additions to documentation label Jun 13, 2025

peri044 added 2 commits June 13, 2025 13:31

chore: Add a tutorial

81114f8

chore: fix model name

57da513

narendasan reviewed Jun 13, 2025

View reviewed changes

zewenli98 reviewed Jun 14, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: Refactor LLM model zoo and add KV cache support #3527

feat: Refactor LLM model zoo and add KV cache support #3527

Uh oh!

peri044 commented May 20, 2025 •

edited

Loading

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as spam.

Uh oh!

narendasan Jun 13, 2025

Uh oh!

narendasan Jun 13, 2025

Uh oh!

narendasan Jun 13, 2025 •

edited

Loading

Uh oh!

narendasan Jun 13, 2025

Uh oh!

zewenli98 left a comment

Uh oh!

zewenli98 Jun 14, 2025

Uh oh!

zewenli98 Jun 14, 2025

Uh oh!

Uh oh!

		from torch.utils._pytree import _LEAF_SPEC


		@torch_tensorrt.dynamo.conversion.dynamo_tensorrt_converter(

		@@ -62,25 +62,15 @@ def scaled_dot_product_attention(
		) -> TRTTensor:

		from utils import export_llm


		def export_llm(model, inputs, min_seq_len=1, max_seq_len=16):

feat: Refactor LLM model zoo and add KV cache support #3527

Are you sure you want to change the base?

feat: Refactor LLM model zoo and add KV cache support #3527

Uh oh!

Conversation

peri044 commented May 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Checklist:

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as spam.

Uh oh!

This comment was marked as spam.

Uh oh!

narendasan Jun 13, 2025

Choose a reason for hiding this comment

Uh oh!

narendasan Jun 13, 2025

Choose a reason for hiding this comment

Uh oh!

narendasan Jun 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

narendasan Jun 13, 2025

Choose a reason for hiding this comment

Uh oh!

zewenli98 left a comment

Choose a reason for hiding this comment

Uh oh!

zewenli98 Jun 14, 2025

Choose a reason for hiding this comment

Uh oh!

zewenli98 Jun 14, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

peri044 commented May 20, 2025 •

edited

Loading

narendasan Jun 13, 2025 •

edited

Loading