Skip to content

[Issue #4403][refactor] Move pattern matching transforms to new InferenceOptimizer - #98

Closed
h-guo18 wants to merge 14 commits into
feat/ad-2025-07-07from
haoguo/move-transforms
Closed

[Issue #4403][refactor] Move pattern matching transforms to new InferenceOptimizer#98
h-guo18 wants to merge 14 commits into
feat/ad-2025-07-07from
haoguo/move-transforms

Conversation

@h-guo18

@h-guo18 h-guo18 commented Jul 19, 2025

Copy link
Copy Markdown

GH Issue #4403 [refactor] Move pattern matching transforms to new InferenceOptimizer

Description

  • Moved the following transformation into the new configurable inference optimizer:
    • quantize
    • moe
    • KVCache
    • ROPE
  • Updated unit test of the corresponding transforms to use the new inference optimizer.

Test Coverage

Unite tests. See changed files.

h-guo18 added 3 commits July 18, 2025 19:16
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
@h-guo18 h-guo18 self-assigned this Jul 19, 2025
@h-guo18
h-guo18 requested a review from lucaslie July 19, 2025 00:02
@h-guo18 h-guo18 changed the title Haoguo/move transforms [GH Issue #4403](https://github.com/NVIDIA/TensorRT-LLM/issues/4403) [refactor] Move KVCache, Quantization to new InferenceOptimizer Jul 19, 2025
@h-guo18 h-guo18 changed the title [GH Issue #4403](https://github.com/NVIDIA/TensorRT-LLM/issues/4403) [refactor] Move KVCache, Quantization to new InferenceOptimizer [GH Issue #4403][refactor] Move KVCache, Quantization to new InferenceOptimizer Jul 19, 2025
@h-guo18 h-guo18 changed the title [GH Issue #4403][refactor] Move KVCache, Quantization to new InferenceOptimizer [Issue #4403][refactor] Move KVCache, Quantization to new InferenceOptimizer Jul 19, 2025
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
@nv-auto-deploy nv-auto-deploy deleted a comment from github-actions Bot Jul 19, 2025
@nv-auto-deploy nv-auto-deploy deleted a comment from github-actions Bot Jul 19, 2025
@h-guo18
h-guo18 marked this pull request as ready for review July 19, 2025 00:22
h-guo18 added 2 commits July 20, 2025 05:01
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
@h-guo18
h-guo18 marked this pull request as draft July 20, 2025 21:29
h-guo18 added 2 commits July 20, 2025 23:06
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
@h-guo18 h-guo18 changed the title [Issue #4403][refactor] Move KVCache, Quantization to new InferenceOptimizer [Issue #4403][refactor] Move pattern matching transforms to new InferenceOptimizer Jul 21, 2025
h-guo18 added 2 commits July 21, 2025 02:21
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>

@lucaslie lucaslie left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For any transform that we move we should:

  1. remove the corresponding transform from the old InferenceOptimizer
  2. configure the default settings in auto_deploy/config/default.yaml

Once we have finalized the reviews, would you be comfortable submitting one PR per transform? (or at least one PR per a couple of transforms). This well help with tracking potential regressions in case we face any

Comment thread tensorrt_llm/_torch/auto_deploy/transform/library/fuse_moe.py Outdated

# TODO:(hg) confirm this
info = TransformInfo(
skipped=False, num_matches=num_moe_patterns, is_clean=False, has_valid_shapes=True

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
skipped=False, num_matches=num_moe_patterns, is_clean=False, has_valid_shapes=True
skipped=False, num_matches=num_moe_patterns, is_clean=False, has_valid_shapes=False

This is safer unless we know that the transform correctly assigns and updates shapes.

@Fridah-nv can also comment on that

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree on this. I think all the transformations should be able to preserve valid shapes except for those using torch._inductor pattern matcher.
Should we require the other transformations to preserve it to avoid running shape propagation multiple times? cc: @lucaslie

Comment thread tensorrt_llm/_torch/auto_deploy/transform/library/fuse_moe.py Outdated
Comment thread tensorrt_llm/_torch/auto_deploy/transform/library/fuse_moe.py Outdated

# TODO:(hg) confirm this
info = TransformInfo(
skipped=False, num_matches=fused_key_counter, is_clean=False, has_valid_shapes=True

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
skipped=False, num_matches=fused_key_counter, is_clean=False, has_valid_shapes=True
skipped=False, num_matches=fused_key_counter, is_clean=False, has_valid_shapes=False

@Fridah-nv can also comment on this

Comment thread tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py
Comment thread tensorrt_llm/_torch/auto_deploy/transform/library/quantize_moe.py
Comment thread tensorrt_llm/_torch/auto_deploy/transform/library/quantization.py
Comment thread tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py

# TODO:(hg) confirm this
info = TransformInfo(
skipped=False, num_matches=num_matches, is_clean=False, has_valid_shapes=True

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Fridah-nv please help confirm @h-guo18 as well

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

has_valid_shapes should be set to False since the pattern matcher utility won't preserve shape information correctly

h-guo18 added 4 commits July 23, 2025 03:25
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
Signed-off-by: haoguo <67671475+h-guo18@users.noreply.github.com>
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.

3 participants