use index_put only in kv cache update to reduce number of operators #3786

cccclai · 2024-05-30T22:42:02Z

Summary:
The decomposition from

class IndexPut(torch.nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, x, input_pos, value):
        x[:, :, input_pos] = value
        return x

is

opcode         name             target                      args                                             kwargs
-------------  ---------------  --------------------------  -----------------------------------------------  --------
placeholder    x                x                           ()                                               {}
placeholder    input_pos        input_pos                   ()                                               {}
placeholder    value            value                       ()                                               {}
call_function  slice_1          aten.slice.Tensor           (x, 0, 0, 9223372036854775807)                   {}
call_function  slice_2          aten.slice.Tensor           (slice_1, 1, 0, 9223372036854775807)             {}
call_function  index_put        aten.index_put.default      (slice_2, [None, None, input_pos], value)        {}
call_function  slice_3          aten.slice.Tensor           (x, 0, 0, 9223372036854775807)                   {}
call_function  slice_scatter    aten.slice_scatter.default  (slice_3, index_put, 1, 0, 9223372036854775807)  {}
call_function  slice_scatter_1  aten.slice_scatter.default  (x, slice_scatter, 0, 0, 9223372036854775807)    {}
output         output           output                      ((slice_scatter_1, slice_scatter_1),)            {}

however x[:, :, input_pos] = value really is just updating the content inside x with value, essentially just index_put

By replacing x[:, :, input_pos] = value with torch.ops.aten.index_put_(x, [None, None, input_pos], value), we reduce the number of operators from 6 to 1.

class IndexPut(torch.nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, x, indices, values):
        torch.ops.aten.index_put_(x, [None, None, input_pos], value)
        return x

decomposition is

opcode         name       target                  args                                 kwargs
-------------  ---------  ----------------------  -----------------------------------  --------
placeholder    x          x                       ()                                   {}
placeholder    input_pos  input_pos               ()                                   {}
placeholder    value      value                   ()                                   {}
call_function  index_put  aten.index_put.default  (x, [None, None, input_pos], value)  {}
output         output     output                  ((index_put, index_put),)            {}

A more proper way to address this in long term is via pattern matching to replace the patterns with the simplified pattern

Differential Revision: D57949659

pytorch-bot · 2024-05-30T22:42:05Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/3786

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit 6574788 with merge base 0333390 ():

NEW FAILURES - The following jobs have failed:

trunk / test-llama-runner-mac (fp32, cmake, coreml) / macos-job (gh)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
trunk / test-llama-runner-mac (fp32, cmake, mps) / macos-job (gh)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

facebook-github-bot · 2024-05-30T22:42:11Z

This pull request was exported from Phabricator. Differential Revision: D57949659

…ytorch#3786) Summary: Pull Request resolved: pytorch#3786 The decomposition from ``` class IndexPut(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, input_pos, value): x[:, :, input_pos] = value return x ``` is ``` opcode name target args kwargs ------------- --------------- -------------------------- ----------------------------------------------- -------- placeholder x x () {} placeholder input_pos input_pos () {} placeholder value value () {} call_function slice_1 aten.slice.Tensor (x, 0, 0, 9223372036854775807) {} call_function slice_2 aten.slice.Tensor (slice_1, 1, 0, 9223372036854775807) {} call_function index_put aten.index_put.default (slice_2, [None, None, input_pos], value) {} call_function slice_3 aten.slice.Tensor (x, 0, 0, 9223372036854775807) {} call_function slice_scatter aten.slice_scatter.default (slice_3, index_put, 1, 0, 9223372036854775807) {} call_function slice_scatter_1 aten.slice_scatter.default (x, slice_scatter, 0, 0, 9223372036854775807) {} output output output ((slice_scatter_1, slice_scatter_1),) {} ``` however `x[:, :, input_pos] = value` really is just updating the content inside `x` with value, essentially just `index_put` By replacing `x[:, :, input_pos] = value` with `torch.ops.aten.index_put_(x, [None, None, input_pos], value)`, we reduce the number of operators from 6 to 1. ``` class IndexPut(torch.nn.Module): def __init__(self): super().__init__() def forward(self, x, indices, values): torch.ops.aten.index_put_(x, [None, None, input_pos], value) return x ``` decomposition is ``` opcode name target args kwargs ------------- --------- ---------------------- ----------------------------------- -------- placeholder x x () {} placeholder input_pos input_pos () {} placeholder value value () {} call_function index_put aten.index_put.default (x, [None, None, input_pos], value) {} output output output ((index_put, index_put),) {} ``` A more proper way to address this in long term is via pattern matching to replace the patterns with the simplified pattern Perf: For stories, before the diff ``` I 00:00:03.437290 executorch:runner.cpp:419] Prompt Tokens: 9 Generated Tokens: 118 I 00:00:03.437295 executorch:runner.cpp:425] Model Load Time: 0.763000 (seconds) I 00:00:03.437301 executorch:runner.cpp:435] Total inference time: 2.661000 (seconds) Rate: 44.344231 (tokens/second) I 00:00:03.437305 executorch:runner.cpp:443] Prompt evaluation: 0.185000 (seconds) Rate: 48.648649 (tokens/second) I 00:00:03.437309 executorch:runner.cpp:454] Generated 118 tokens: 2.476000 (seconds) Rate: 47.657512 (tokens/second) I 00:00:03.437313 executorch:runner.cpp:462] Time to first generated token: 0.206000 (seconds) I 00:00:03.437315 executorch:runner.cpp:469] Sampling time over 127 tokens: 0.042000 (seconds) ``` After the diff ``` I 00:00:03.195257 executorch:runner.cpp:419] Prompt Tokens: 9 Generated Tokens: 118 I 00:00:03.195295 executorch:runner.cpp:425] Model Load Time: 0.683000 (seconds) I 00:00:03.195314 executorch:runner.cpp:435] Total inference time: 2.502000 (seconds) Rate: 47.162270 (tokens/second) I 00:00:03.195319 executorch:runner.cpp:443] Prompt evaluation: 0.175000 (seconds) Rate: 51.428571 (tokens/second) I 00:00:03.195323 executorch:runner.cpp:454] Generated 118 tokens: 2.327000 (seconds) Rate: 50.709067 (tokens/second) I 00:00:03.195327 executorch:runner.cpp:462] Time to first generated token: 0.195000 (seconds) I 00:00:03.195330 executorch:runner.cpp:469] Sampling time over 127 tokens: 0.049000 (seconds) ``` Differential Revision: D57949659

facebook-github-bot · 2024-07-18T22:08:40Z

This pull request was exported from Phabricator. Differential Revision: D57949659

YifanShenSZ · 2024-07-18T22:10:20Z

examples/models/llama2/llama_transformer.py

+            )
+            v_out = torch.ops.aten.index_put_(
+                self.v_cache, [None, None, input_pos], v_val
+            )


Is the 2nd v_out redundant?

larryliu0820

Can you address the comment then land?

iseeyuan · 2025-02-11T01:10:46Z

@cccclai Do you have any performance numbers for before and after comparison?

cccclai · 2025-02-11T20:00:18Z

@cccclai Do you have any performance numbers for before and after comparison?

about 5% IIRC.

cccclai · 2025-06-06T00:55:42Z

Stale PR

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 30, 2024

facebook-github-bot added the fb-exported label May 30, 2024

cccclai mentioned this pull request Jun 5, 2024

add 16a4w_hqq quant mode #3752

Closed

cccclai added the ciflow/trunk label Jul 18, 2024

cccclai force-pushed the export-D57949659 branch from d9aed5a to 6574788 Compare July 18, 2024 22:08

YifanShenSZ reviewed Jul 18, 2024

View reviewed changes

larryliu0820 approved these changes Jul 18, 2024

View reviewed changes

cccclai mentioned this pull request Jul 30, 2024

Use Core ML Quantizer in Llama Export #4458

Closed

cccclai closed this Jun 6, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

use index_put only in kv cache update to reduce number of operators #3786

use index_put only in kv cache update to reduce number of operators #3786

Uh oh!

cccclai commented May 30, 2024

Uh oh!

pytorch-bot bot commented May 30, 2024 •

edited

Loading

Uh oh!

facebook-github-bot commented May 30, 2024

Uh oh!

facebook-github-bot commented Jul 18, 2024

Uh oh!

YifanShenSZ Jul 18, 2024

Uh oh!

larryliu0820 left a comment

Uh oh!

iseeyuan commented Feb 11, 2025

Uh oh!

cccclai commented Feb 11, 2025

Uh oh!

cccclai commented Jun 6, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

use index_put only in kv cache update to reduce number of operators #3786

use index_put only in kv cache update to reduce number of operators #3786

Uh oh!

Conversation

cccclai commented May 30, 2024

Uh oh!

pytorch-bot bot commented May 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/3786

❌ 2 New Failures

Uh oh!

facebook-github-bot commented May 30, 2024

Uh oh!

facebook-github-bot commented Jul 18, 2024

Uh oh!

YifanShenSZ Jul 18, 2024

Choose a reason for hiding this comment

Uh oh!

larryliu0820 left a comment

Choose a reason for hiding this comment

Uh oh!

iseeyuan commented Feb 11, 2025

Uh oh!

cccclai commented Feb 11, 2025

Uh oh!

cccclai commented Jun 6, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

pytorch-bot bot commented May 30, 2024 •

edited

Loading