Integrate int4kv to benchmark_attn_decode.py #1029

scxiao · 2024-04-19T13:46:22Z

What does this PR do?

Integrate the int4kv standalone test to the benchmark benchmark_attn_decode.py

Before submitting

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

facebook-github-bot · 2024-04-19T13:46:28Z

Hi @scxiao!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [email protected]. Thanks!

codecov-commenter · 2024-04-19T14:43:26Z

Codecov Report

Attention: Patch coverage is 0% with 1 lines in your changes are missing coverage. Please review.

Project coverage is 59.93%. Comparing base (67f38b9) to head (50d747c).

Files	Patch %	Lines
xformers/ops/fmha/triton_splitk.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1029   +/-   ##
=======================================
  Coverage   59.93%   59.93%           
=======================================
  Files         113      113           
  Lines       10007    10007           
=======================================
  Hits         5998     5998           
  Misses       4009     4009

Flag	Coverage Δ
Python	`59.93% <0.00%> (ø)`

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

facebook-github-bot · 2024-04-19T22:06:46Z

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

scxiao · 2024-04-24T19:23:52Z

Hi @jianyuh, could you please help to review this PR to include int4kv in the FA decode benchmark? It seems like I cannot add reviewer here. Thanks

jianyuh · 2024-04-25T06:13:22Z

Thanks @scxiao ! Could you fix the lint issue?

jianyuh · 2024-04-25T06:14:40Z

xformers/benchmarks/benchmark_attn_decoding.py

 class AttentionDecodingPyTorchRepeat(AttentionDecodingBase):
    def fw(self) -> None:
        B, Mq, Mkv, Hq, Hkv, K = self.shapes
        scale = 1 / K**0.5
        q = self.q.reshape([B, Mq, -1, K]).permute(0, 2, 1, 3)
        k = self.k.reshape([B, Mkv, -1, K]).permute(0, 2, 1, 3)
        v = self.v.reshape([B, Mkv, -1, K]).permute(0, 2, 1, 3)
-        attn = (q @ k.transpose(-1, -2)).softmax(-1) * scale
+        attn = (q @ k.transpose(-1, -2) * scale).softmax(-1)


Thanks for catching this! Looks like we didn't sync with https://github.com/ROCm/phantom_amd_llm_operators/pull/18/files ..

bottler · 2024-04-25T11:36:14Z

xformers/ops/fmha/triton_splitk.py

@@ -637,7 +637,7 @@ def dequantize(
            x_[:, :, None, :] >> offsets
        )  # (BLOCK_N, D // PACKED_PER_VAL, PACKED_PER_VAL)

-        quant_offset = tl.reshape(
+        quant_offset = tl.view(


I think tl.view, which we used to use here and which used to do a non obvious thing, still officially has unspecified behaviour in triton. tl.reshape actually claims to return the right thing, which is nice 😊. Are you changing this here because you observe a speedup, or because you expect one, or is this something else e.g. a fix?

I saw an error complaining no tl.reshape, and suggesting use tl.view. Let me double check.

scxiao · 2024-04-25T22:01:17Z

Thanks @scxiao ! Could you fix the lint issue?

It seems like I did not change that file, but I fixed the format anyway. Is that OK for you? Thanks

jianyuh · 2024-04-26T01:52:30Z

xformers/benchmarks/benchmark_attn_decoding.py

@@ -141,14 +178,93 @@ class AttentionDecodingCKSplitKV(AttentionDecodingBase):
    OP = xops.fmha.ck_splitk.FwOp


+class AttentionDecodingSplitInt4KV(AttentionDecodingBase):
+    OP = xops.fmha.triton_splitk.FwOp
+    def __init__(self, B: int, Mq: int, Mkv: int, Hq: int, Hkv: int, K: int, bw: bool,


We have lint issues.

./xformers/benchmarks/benchmark_attn_decoding.py:183:5: E301 expected 1 blank line, found 0 def __init__(self, B: int, Mq: int, Mkv: int, Hq: int, Hkv: int, K: int, bw: bool, ^ ./xformers/benchmarks/benchmark_attn_decoding.py:1[8](https://github.com/facebookresearch/xformers/actions/runs/8839890145/job/24274714295?pr=1029#step:10:9)5:5: E124 closing bracket does not match visual indentation ) -> None: ^ ./xformers/benchmarks/benchmark_attn_decoding.py:258:1: W2[9](https://github.com/facebookresearch/xformers/actions/runs/8839890145/job/24274714295?pr=1029#step:10:10)3 blank line contains whitespace ^

scxiao · 2024-04-29T18:52:17Z

Hi @sgrigory, could you please help review this PR when you get a chance? Thanks

bottler · 2024-05-02T10:25:51Z

xformers/ops/fmha/__init__.py

@@ -569,9 +569,9 @@ def merge_attentions(
    concat_path = attn_is_concat and lse_is_concat
    if not concat_path:
        if attn_is_concat:
-            attn_split = cast(torch.Tensor, attn_split).unbind(0)


Was this change fixing a failure, or was it to make the type checker happy, or something else?

this change seems wrong to me

The typing is not quite right at the moment - because unbind returns a tuple and the code which uses it needs a list I think, although it just needs a Sequence.

This change will /work/ because a Tensor is usable as a Sequence of tensors, which is what the non-concat path needs, so the unbind was unnecessary. And this change probably also fixes the type error because mypy doesn't get finer than knowing attn_split is a Union. But this change does make the code slightly less understandable.

I think OP is just trying to fix the existing lint errors in this PR: here and in the predicated...h file (which is clearly the correct fix). I'm fine if they leave them broken or just this one broken.

Thanks @bottler. Yes, the changes here are to fix lint error.

scxiao · 2024-05-06T13:54:01Z

Hi @sgrigory, when you get a chance, could you please help review this PR, so we can get it merge soon? Thanks.
CC: @jianyuh.

sgrigory · 2024-05-06T14:10:09Z

hi @scxiao <https://github.com/scxiao>, I'm on vacation and travelling till May 11 - I can review it after that, or @bottler <https://github.com/bottler> will have a look in the meantime

…

On Mon 6. May 2024 at 22:54, Shucai Xiao ***@***.***> wrote: Hi @sgrigory <https://github.com/sgrigory>, when you get a chance, could you please help review this PR, so we can get it merge soon? Thanks. CC: @jianyuh <https://github.com/jianyuh>. — Reply to this email directly, view it on GitHub <#1029 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AFRAGLQZW5THRIHFBC722ELZA6DRBAVCNFSM6AAAAABGPHEJR2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDAOJWGA3TIOJQGA> . You are receiving this because you were mentioned.Message ID: ***@***.***>

bottler · 2024-05-07T15:41:33Z

LGTM except the reshape->view change which I'm still uncertain about. Can we leave unchanged?

scxiao · 2024-05-08T16:48:47Z

LGTM except the reshape->view change which I'm still uncertain about. Can we leave unchanged?

I just tried again, it seems like the error related to reshape is from our rocm fork. I reverted this change.

scxiao marked this pull request as draft April 19, 2024 13:46

scxiao marked this pull request as ready for review April 19, 2024 19:32

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 Apr 19, 2024

jianyuh requested review from sgrigory and bottler April 25, 2024 06:12

jianyuh reviewed Apr 25, 2024

View reviewed changes

bottler reviewed Apr 25, 2024

View reviewed changes

jianyuh reviewed Apr 26, 2024

View reviewed changes

scxiao force-pushed the int4kv_fa_decode branch from 1f30ac9 to dfd0f39 Compare April 26, 2024 19:57

jianyuh approved these changes Apr 28, 2024

View reviewed changes

bottler reviewed May 2, 2024

View reviewed changes

scxiao added 5 commits May 3, 2024 17:20

add int4kv to FA docode

4b8cb8b

backup changes

18fa47f

code cleanup

aac4824

int4kv test passed

2ea2b49

fix format

50d747c

scxiao force-pushed the int4kv_fa_decode branch from f965a72 to 50d747c Compare May 3, 2024 22:28

revert a change

65f404c

bottler merged commit 60d5f11 into facebookresearch:main May 8, 2024
2 of 4 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Integrate int4kv to benchmark_attn_decode.py #1029

Integrate int4kv to benchmark_attn_decode.py #1029

scxiao commented Apr 19, 2024

facebook-github-bot commented Apr 19, 2024

codecov-commenter commented Apr 19, 2024 •

edited

Loading

facebook-github-bot commented Apr 19, 2024

scxiao commented Apr 24, 2024

jianyuh commented Apr 25, 2024

jianyuh Apr 25, 2024

bottler Apr 25, 2024

scxiao Apr 25, 2024

scxiao commented Apr 25, 2024

jianyuh Apr 26, 2024

scxiao commented Apr 29, 2024

bottler May 2, 2024

danthe3rd May 2, 2024

bottler May 2, 2024

scxiao May 2, 2024

scxiao commented May 6, 2024

sgrigory commented May 6, 2024 via email

bottler commented May 7, 2024

scxiao commented May 8, 2024

Integrate int4kv to benchmark_attn_decode.py #1029

Integrate int4kv to benchmark_attn_decode.py #1029

Conversation

scxiao commented Apr 19, 2024

What does this PR do?

Before submitting

PR review

facebook-github-bot commented Apr 19, 2024

Action Required

Process

codecov-commenter commented Apr 19, 2024 • edited Loading

Codecov Report

facebook-github-bot commented Apr 19, 2024

scxiao commented Apr 24, 2024

jianyuh commented Apr 25, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

scxiao commented Apr 25, 2024

Choose a reason for hiding this comment

scxiao commented Apr 29, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

scxiao commented May 6, 2024

sgrigory commented May 6, 2024 via email

bottler commented May 7, 2024

scxiao commented May 8, 2024

codecov-commenter commented Apr 19, 2024 •

edited

Loading