feat: Support cos_sin_cache in all cases. by yuxianq · Pull Request #3020 · NVIDIA/TensorRT-LLM

yuxianq · 2025-03-24T09:44:21Z

This MR contains the following updates:

Handle fuse_pos_embd=True/False and create RotaryEmbedding inside attention module, so that the users don't need to handle it in the modeling files.
Cache cos_sin for unfused rope implementation. If flashinfer is available, use apply_rope_with_cos_sin_cache_inplace instead of apply_rope_inplace. Otherwise, we fallback to pure pytorch implementation, which can support any rope now.
We use create_rope_const_params to create and cache cos_sin_cache for all rope types, including Deepseek yarn rope.

yuxianq · 2025-03-24T09:45:06Z

/bot run --add-multi-gpu-test

niukuo · 2025-03-24T09:50:37Z

PR_Github #283 [ run ] triggered by Bot

niukuo · 2025-03-24T11:15:30Z

PR_Github #283 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #272 completed with status: 'FAILURE'

yuxianq · 2025-03-25T06:39:08Z

/bot run --add-multi-gpu-test

niukuo · 2025-03-25T06:46:37Z

PR_Github #387 [ run ] triggered by Bot

niukuo · 2025-03-25T07:09:26Z

PR_Github #387 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #345 completed with status: 'FAILURE'

yuxianq · 2025-03-25T12:01:38Z

/bot run --add-multi-gpu-test

niukuo · 2025-03-25T12:10:32Z

PR_Github #430 [ run ] triggered by Bot

niukuo · 2025-03-25T14:03:13Z

PR_Github #430 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #369 completed with status: 'FAILURE'

yuxianq · 2025-03-26T04:01:04Z

/bot run --add-multi-gpu-test

niukuo · 2025-03-26T04:07:45Z

PR_Github #510 [ run ] triggered by Bot

BestJuly · 2025-03-26T05:32:21Z

I think I am pinged by mistake, is the review request actually pointed to @litaotju ?

niukuo · 2025-03-26T05:33:38Z

PR_Github #510 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #437 completed with status: 'FAILURE'

QiJune · 2025-03-26T06:15:04Z

@yuxianq Can we split this PR to several small PRs? For example, the first item can be a single PR.

Handle fuse_pos_embd=True/False and create RotaryEmbedding inside attention module, so that the users don't need to handle it in the modeling files.

yuxianq · 2025-03-26T07:07:03Z

Can we split this PR to several small PRs? For example, the first item can be a single PR.

@QiJune I will have a try. Let me pass the CI first to validate that these features work correctly.

yuxianq · 2025-03-26T08:47:15Z

/bot run --add-multi-gpu-test

niukuo · 2025-03-26T08:56:31Z

PR_Github #550 [ run ] triggered by Bot

niukuo · 2025-03-26T11:38:21Z

PR_Github #550 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #469 completed with status: 'FAILURE'

yuxianq · 2025-03-26T12:13:38Z

/bot run --disable-fail-fast --add-multi-gpu-test

niukuo · 2025-03-26T12:21:57Z

PR_Github #584 [ run ] triggered by Bot

niukuo · 2025-03-26T15:24:50Z

PR_Github #584 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #497 completed with status: 'FAILURE'

litaotju · 2025-04-01T04:06:33Z

Just curious, are someone importing this file?

It should be used a cli command only.

No. I refactor it when I debug thread/process leak for multiprocessing module. It should be fine for a cli command.

litaotju · 2025-04-01T04:10:58Z

This looks somewaht ad-hoc to me. Is it possible to not relying on the model type hard code string here in a general interface?

It is a workaround for DeepSeek V3 Lite since its rope_scaling is null in config.json. DeepSeek V3/R1 should be fine without this code. I have added a comment to explain it.

litaotju · 2025-04-01T04:12:36Z

Maybe Nit.

Can we say create*positions_for_yarn or something similar w/o to be too specific to DeepSeek easily?

I have renamed it to create_sinusoidal_positions_yarn. DeepSeek use a extended version of yarn with additional mscale and mscale_all_dim arguments. Since it is compatible with the original version when using default value for mscale and mscale_all_dim, we can merge the Deepseek variant with the original version.

litaotju · 2025-04-01T04:25:11Z

you should be the custom op will do rope fusion right?

Maybe renaming to something like? self.rope_fused_in_custom_op = True

I have renamed it to enable_rope_fusion.

litaotju · 2025-04-01T04:26:57Z

"support" is a vague word.
support means which one?

configuable, both fused qkv and unfused qkv can run. or

requires fused qkv?

I use "support" because a backend can support both fused qkv and unfused qkv, so they are not mutually exclusive. Currently each backend only supports one case.

litaotju · 2025-04-01T04:30:15Z

when will this be true?
Can we always attach the extra_attrs in _load_model such that this won't be needed?

Some tests use stub models and override _load_model (see https://github.com/NVIDIA/TensorRT-LLM/blob/7441714d66e179192ef0642a9a150f5486e5b06b/tests/unittest/_torch/test_pytorch_model_engine.py#L50). So we have to add this code to ensure the extra_attrs is attached. I have added comment to explain it.

litaotju · 2025-04-01T04:31:27Z

shall we wrap the whole forward function be inside this context manager?

Good point, have wrapped the whole forward function.

litaotju · 2025-04-01T04:32:50Z

@Superjomn do you remember if we have some reason to make this "wait=False"?

It can avoid thread leak, which can further avoid memory leak.

litaotju · 2025-04-01T04:33:54Z

will not this llm object be automatically destoryed when the function return? its just a local var.

Yes, it is unnecessary. Just for debugging OOM issues. I have removed all unnecessary del llm.

litaotju · 2025-04-01T04:34:52Z

same here, is there some reason that the object not deleted by python when function returns?

I have removed all unnecessary del llm.

yuxianq · 2025-04-02T08:27:01Z

/bot run --disable-fail-fast --stage-list "A30-7"

yuxianq · 2025-04-02T09:48:57Z

/bot run --disable-fail-fast --stage-list "A30-7"

tensorrt-cicd · 2025-04-02T09:55:44Z

PR_Github #1005 [ run ] triggered by Bot

tensorrt-cicd · 2025-04-02T13:32:29Z

PR_Github #1005 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #776 (Partly Tested) completed with status: 'FAILURE'

yuxianq · 2025-04-02T15:47:13Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2025-04-02T15:52:35Z

PR_Github #1030 [ run ] triggered by Bot

tensorrt-cicd · 2025-04-02T18:14:20Z

PR_Github #1030 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #794 completed with status: 'FAILURE'

yuxianq · 2025-04-03T07:28:39Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2025-04-03T07:34:14Z

PR_Github #1088 [ run ] triggered by Bot

tensorrt-cicd · 2025-04-03T09:55:46Z

PR_Github #1088 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #829 completed with status: 'FAILURE'

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

yuxianq · 2025-04-03T10:47:17Z

/bot run --disable-fail-fast --add-multi-gpu-test

tensorrt-cicd · 2025-04-03T10:52:44Z

PR_Github #1110 [ run ] triggered by Bot

tensorrt-cicd · 2025-04-03T13:19:12Z

PR_Github #1110 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #843 completed with status: 'FAILURE'

yuxianq · 2025-04-14T04:20:48Z

This PR has been split into several PRs. Move the main part of this PR to #3517 and close this outdated PR.

yuxianq requested review from BestJuly, QiJune, hlu1 and kaiyux March 24, 2025 09:44

yuxianq force-pushed the user/yuxianq/cos-sin-cache branch 2 times, most recently from 95840f8 to 54d797b Compare March 26, 2025 04:00

BestJuly removed their request for review March 26, 2025 05:20

yuxianq requested a review from litaotju March 26, 2025 07:01

yuxianq changed the title ~~Support cos_sin_cache in all cases.~~ feat: Support cos_sin_cache in all cases. Mar 26, 2025

litaotju reviewed Apr 1, 2025

View reviewed changes

yuxianq force-pushed the user/yuxianq/cos-sin-cache branch from 77b5b14 to 7441714 Compare April 2, 2025 09:48

yuxianq added 6 commits April 3, 2025 10:40

Add thread leak check and fix thread/memory leak issues.

33ca134

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Fix thread leak issues.

c7e685d

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Refine attention backend interface.

3b9951c

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Cache sin cos in model instead of global LRU cache.

885b9bb

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Handle rope fusion inside attention module instead of modeling files.

8f59736

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

Support unfused rope in MLA.

e774b7c

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>

yuxianq force-pushed the user/yuxianq/cos-sin-cache branch from 269a813 to e774b7c Compare April 3, 2025 10:46

yuxianq closed this Apr 14, 2025

chenfeiz0326 mentioned this pull request Oct 16, 2026

[https://nvbugs/6843484][fix] Keep multi-node aggregated perf-sanity jobs in one NVLink domain #20418

Merged

1 task done

Conversation

yuxianq commented Mar 24, 2025

Uh oh!

yuxianq commented Mar 24, 2025

Uh oh!

niukuo commented Mar 24, 2025

Uh oh!

niukuo commented Mar 24, 2025

Uh oh!

yuxianq commented Mar 25, 2025

Uh oh!

niukuo commented Mar 25, 2025

Uh oh!

niukuo commented Mar 25, 2025

Uh oh!

yuxianq commented Mar 25, 2025

Uh oh!

niukuo commented Mar 25, 2025

Uh oh!

niukuo commented Mar 25, 2025

Uh oh!

yuxianq commented Mar 26, 2025

Uh oh!

niukuo commented Mar 26, 2025

Uh oh!

BestJuly commented Mar 26, 2025

Uh oh!

niukuo commented Mar 26, 2025

Uh oh!

QiJune commented Mar 26, 2025

Uh oh!

yuxianq commented Mar 26, 2025

Uh oh!

yuxianq commented Mar 26, 2025

Uh oh!

niukuo commented Mar 26, 2025

Uh oh!

niukuo commented Mar 26, 2025

Uh oh!

yuxianq commented Mar 26, 2025

Uh oh!

niukuo commented Mar 26, 2025

Uh oh!

niukuo commented Mar 26, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment