[sharktank] restore custom matmul kernel #896

dan-garvey · 2025-02-02T00:10:53Z

fix iree-org/iree#19859 (comment)

python3 -m sharktank.examples.export_paged_llm_v1 --irpa-file=/sharedfile/llama3_8b_fp8.irpa \
--output-mlir=fp8_dan1.mlir \
--output-config=config1.json \
--bs=1 --attention-kernel torch \
--attention-dtype=float8_e4m3fnuz --activation-dtype=bfloat16

AmosLewis · 2025-02-03T17:04:59Z

cannot reproduce the numeric result with this patch,here is the log danfix.log

This argument can be skipped. Then the result dtype is inferred from the arguments. This requires iree-org/iree-turbine#451

sogartar · 2025-02-04T02:52:43Z

I added some modifications that require merging iree-org/iree-turbine#451 first.

I made the batch_matmul_transpose_b kernel accept an optional accumulation dtype.

sogartar · 2025-02-05T16:18:54Z

This change needs an IREE Turbine bump #916.

dan-garvey · 2025-02-07T20:08:06Z

are you still working on this @sogartar

sogartar · 2025-02-08T02:31:01Z

@dan-garvey the kernel is ready. I also cleaned up linear and qlinear.
Should this kernel be also used in the sharktank.ops.matmul? I have not added it there. The only places that this seems useful is when we want to control the accumulation dtype. I assume that b transposition and then matmul would be fused by the compiler.

sogartar · 2025-02-08T02:36:54Z

I did not notice this failing unit test. I will see what is the problem.
I added one xfail batch_matmul_transpose_b kernel test that can't be compiled for llvm-cpu. It does compile for amdgpu.

sogartar · 2025-02-08T16:19:42Z

The kernel does not work with unsigned types. It needs to be adapted.
The problem is that we reinterpret cast to signless integer types and lose this info when we get to linalg.batch_matmul_transpose_b. Maybe linalg.batch_matmul_transpose_b when promoting from i8 to i32 assumes a
signed type even though i8 is signless.

I made an exception in qlinear if there are unsigned integer types to first promote to the accumulation dtype.

sogartar · 2025-02-08T16:23:55Z

@dan-garvey let me know if it looks good. I will approve the PR as you can't probably approve your own PR.

AmosLewis

@dan-garvey Before merged, we need this issue fix. I got error when I tried to export input bin for f8 model, the error tracked to batch_matmul_transpose_b.py
paged_llm_v_export_f8_input_bin_bug.log @aviator19941 need your review as well to ensure the paged_llm_v1.py work as well.

sogartar · 2025-02-10T23:02:55Z

@AmosLewis we get this error when compiling for llvm-cpu. It compiles for amdgpu. I guess we should put an eager implementation for the kernel.

sogartar · 2025-02-10T23:10:16Z

I will implement to eager path to not use IREE.

AmosLewis · 2025-02-11T01:05:38Z

@AmosLewis we get this error when compiling for llvm-cpu. It compiles for amdgpu. I guess we should put an eager implementation for the kernel.

Get a similar error in paged_llama_attention_block.py with rebase eager commit https://gist.github.com/AmosLewis/3860a0371236b528b24c86c48e1e31c2?permalink_comment_id=5433466#gistcomment-5433466

sogartar · 2025-02-11T18:38:53Z

I will run the model locally and debug this problem.

archana-ramalingam · 2025-02-14T07:42:50Z

@sogartar git commit 5bf4636 fixing .index_copy_ breaks iree-compile for fp8 8b.
Error: https://gist.github.com/archana-ramalingam/7347234893c5eba773d826283d5a635d

sogartar · 2025-02-14T16:55:31Z

May bad. I will fix it.

sogartar · 2025-02-14T23:45:41Z

Here is an IREE issue for the compilation failure iree-org/iree#20002.

IanWood1 · 2025-02-19T00:07:20Z

sharktank/sharktank/layers/kv_cache.py

+            page_table.index_put_(
+                indices=indices, values=ops.to(cache_partition, dtype=page_table.dtype)
+            )


I'm seeing element-wise conversion from f8 to int8 before writing to the cache. I think you need view(dtype) instead of to()

Thank you, I noticed this problem. I am preparing a fix, but there are other numerical issues that also need to be resolved.

I'm seeing this compilation error now with those changes:

fp8_dan_attn_q_bf16_f32_v4.mlir:488:12: error: 'flow.tensor.bitcast' op value set has 1 dynamic dimensions but only 0 dimension values are attached %178 = torch.aten.view.dtype %168, %int1_151 : !torch.vtensor<[?,32,8,128],f8E4M3FNUZ>, !torch.int -> !torch.vtensor<[?,32,8,128],si8> ^ fp8_dan_attn_q_bf16_f32_v4.mlir:488:12: note: see current operation: %244 = "flow.tensor.bitcast"(%243) <{operandSegmentSizes = array<i32: 1, 0, 0>}> : (tensor<?x32x8x128xf8E4M3FNUZ>) -> tensor<?x32x8x128xi8>

This was recently fixed in iree. You need an iree build with iree-org/iree@d38b22e (merged yesterday)

ah ok cool thanks!

…le_custom_fp8_matmul

archana-ramalingam · 2025-02-20T21:06:04Z

sharktank/sharktank/utils/hf_datasets.py

-            extra_filenames=["tokenizer.json"],
-        ),
-    ),
-).alias_to("llama3_8B_fp16")


This model cannot be removed as it is part of llama_serving.md release docs, that was tested in sharktank/shortfin. However, agree that we need only one of the 2 llama 8b fp16 models listed here.

The docs could be updated to use standard huggingface tooling instead of sharktank.utils.hf_datasets. As written today, the hf_datasets file should only be a utility for project development and testing, not something user-facing.

I fixed the doc.

llama3_8B_fp16 is the right instruct version that was tested to be numerically right for the release, the other model is non-instruct version and I remember it was generating repetitive tokens.
This can be a separate PR where we consult with shortfin folks if this model switch on the release docs can be made.

… error

sogartar · 2025-02-21T00:31:37Z

@archana-ramalingam, although I technically reviewed the PR, could you review it as I am now the one making changes.

Make the quantized linear op use the batch_matmul_transpose_b kernel. Make the batch_matmul_transpose_b kernel accept accum_dtype argument that specifies in what precision the operation is performed. Make LLM models accept a kv-cache-dtype argument that specifies only the dtype of the KV cache that can be different from the attention dtype. This config option is used in Llama f8. For Llama f8 the IREE change iree-org/iree#20005 is required. Fixes iree-org/iree#19859 (comment) ``` python3 -m sharktank.examples.export_paged_llm_v1 --irpa-file=/sharedfile/llama3_8b_fp8.irpa \ --output-mlir=fp8_dan1.mlir \ --output-config=config1.json \ --bs=1 \ --attention-kernel torch \ --attention-dtype=bfloat16 \ --activation-dtype=bfloat16 \ --kv-cache-dtype=float8_e4m3fnuz ``` --------- Signed-off-by: Boian Petkantchin<boian.petkantchin@amd.com> Co-authored-by: aviator19941 <avinash.sharma@amd.com> Co-authored-by: archana-ramalingam <archana.ramalingam@amd.com>

dan-garvey requested a review from nithinsubbiah February 2, 2025 00:10

restore custom matmul kernel

8300bc8

dan-garvey mentioned this pull request Feb 2, 2025

Numeric issue for llama_8b_fp8 model on hip iree-org/iree#19859

Closed

not mergeable as-is

a98a332

Make batch_matmul_transpose_b accept accumulation dtype

80fee98

This argument can be skipped. Then the result dtype is inferred from the arguments. This requires iree-org/iree-turbine#451

Merge batch_matmul_transpose_b export tests into 1

9ff020e

sogartar requested review from rsuderman, AmosLewis and IanNod February 5, 2025 03:10

sogartar mentioned this pull request Feb 5, 2025

Bump pinned iree-turbine version to 3.2.0rc20250205 #916

Merged

Merge branch 'main' into users/dan-garvey/enable_custom_fp8_matmul

7b93a6a

sogartar added 2 commits February 8, 2025 16:08

Add exception to qlinear to not use the kernel when unsigned ints

781c8e8

Small fix

9f1c3d4

sogartar approved these changes Feb 8, 2025

View reviewed changes

AmosLewis reviewed Feb 10, 2025

View reviewed changes

AmosLewis requested a review from aviator19941 February 10, 2025 04:32

sogartar added 2 commits February 10, 2025 23:19

Add eager execution to circamvent failure to compile for llvm-cpu

82b032a

Merge branch 'main' into users/dan-garvey/enable_custom_fp8_matmul

aa5c7b0

sogartar added 2 commits February 14, 2025 05:06

Fix KV cache index_copy_ f8 workaround

462ddc4

In linear for (Tensor, QuantizedTensor) raise if accum_dtype is given

4dc2ac2

sogartar mentioned this pull request Feb 14, 2025

torch.aten.view.dtype from f8E4M3FNUZ to si8 with dynamic dimension fails to lower iree-org/iree#20002

Closed

IanWood1 reviewed Feb 19, 2025

View reviewed changes

sogartar added 9 commits February 20, 2025 16:49

Fix KV cache f8

13bfc68

Remove unused HF dataset

8b20445

Add KV cache dtype different from attention dtype

9f160f1

Add more KV cache tests for various dtypes

77a8443

Remove some unwanted corner casehandlings in linear layer

1ea608a

Add more linear layer tests

6f0c98b

Refactor quark parity test to use tmp dir

664a847

Fix KV cache dtype CLI arg parsing

9816c35

Merge remote-tracking branch 'origin/main' into users/dan-garvey/enab…

b8ff8cc

…le_custom_fp8_matmul

archana-ramalingam reviewed Feb 20, 2025

View reviewed changes

sogartar added 3 commits February 20, 2025 21:55

Change doc example to not use the removed Llama dataset

c17629e

Add KV cache dtype to benchmark

9b7dfdf

Change testBenchmark8B_fp8_Non_Decomposed xfail reason to compilation…

55c8701

… error

archana-ramalingam and others added 2 commits February 20, 2025 16:37

Merge branch 'main' into users/dan-garvey/enable_custom_fp8_matmul

fea5204

Put back in the llama3_8B_fp16 HF dataset

740bb80

archana-ramalingam approved these changes Feb 21, 2025

View reviewed changes

sogartar added 3 commits February 21, 2025 14:20

Remove left behind comment

02215d5

Make quark parity test use f8 KV cache

40f993a

Add more bf16 qlinear tests and make ref dtype be f64

85053ef

sogartar merged commit 6397ead into main Feb 21, 2025
36 checks passed

sogartar deleted the users/dan-garvey/enable_custom_fp8_matmul branch February 21, 2025 21:05

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[sharktank] restore custom matmul kernel #896

[sharktank] restore custom matmul kernel #896

dan-garvey commented Feb 2, 2025 •

edited by AmosLewis

Loading

AmosLewis commented Feb 3, 2025

sogartar commented Feb 4, 2025 •

edited

Loading

sogartar commented Feb 5, 2025

dan-garvey commented Feb 7, 2025

sogartar commented Feb 8, 2025

sogartar commented Feb 8, 2025

sogartar commented Feb 8, 2025

sogartar commented Feb 8, 2025

AmosLewis left a comment •

edited

Loading

sogartar commented Feb 10, 2025

sogartar commented Feb 10, 2025

AmosLewis commented Feb 11, 2025 •

edited

Loading

sogartar commented Feb 11, 2025

archana-ramalingam commented Feb 14, 2025

sogartar commented Feb 14, 2025

sogartar commented Feb 14, 2025

IanWood1 Feb 19, 2025 •

edited

Loading

sogartar Feb 19, 2025

aviator19941 Feb 19, 2025

IanWood1 Feb 19, 2025

aviator19941 Feb 19, 2025

archana-ramalingam Feb 20, 2025

ScottTodd Feb 20, 2025

sogartar Feb 20, 2025

archana-ramalingam Feb 20, 2025 •

edited

Loading

sogartar commented Feb 21, 2025

[sharktank] restore custom matmul kernel #896

[sharktank] restore custom matmul kernel #896

Conversation

dan-garvey commented Feb 2, 2025 • edited by AmosLewis Loading

AmosLewis commented Feb 3, 2025

sogartar commented Feb 4, 2025 • edited Loading

sogartar commented Feb 5, 2025

dan-garvey commented Feb 7, 2025

sogartar commented Feb 8, 2025

sogartar commented Feb 8, 2025

sogartar commented Feb 8, 2025

sogartar commented Feb 8, 2025

AmosLewis left a comment • edited Loading

Choose a reason for hiding this comment

sogartar commented Feb 10, 2025

sogartar commented Feb 10, 2025

AmosLewis commented Feb 11, 2025 • edited Loading

sogartar commented Feb 11, 2025

archana-ramalingam commented Feb 14, 2025

sogartar commented Feb 14, 2025

sogartar commented Feb 14, 2025

IanWood1 Feb 19, 2025 • edited Loading

Choose a reason for hiding this comment

sogartar Feb 19, 2025

Choose a reason for hiding this comment

aviator19941 Feb 19, 2025

Choose a reason for hiding this comment

IanWood1 Feb 19, 2025

Choose a reason for hiding this comment

aviator19941 Feb 19, 2025

Choose a reason for hiding this comment

archana-ramalingam Feb 20, 2025

Choose a reason for hiding this comment

ScottTodd Feb 20, 2025

Choose a reason for hiding this comment

sogartar Feb 20, 2025

Choose a reason for hiding this comment

archana-ramalingam Feb 20, 2025 • edited Loading

Choose a reason for hiding this comment

sogartar commented Feb 21, 2025

dan-garvey commented Feb 2, 2025 •

edited by AmosLewis

Loading

sogartar commented Feb 4, 2025 •

edited

Loading

AmosLewis left a comment •

edited

Loading

AmosLewis commented Feb 11, 2025 •

edited

Loading

IanWood1 Feb 19, 2025 •

edited

Loading

archana-ramalingam Feb 20, 2025 •

edited

Loading