fp8_gemm def update and workload

[YiyanZhai]

No description provided.

[YiyanZhai]

WIP: reference impl

[gemini-code-assist]

Summary of Changes

Hello @YiyanZhai, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refines the definition and implementation of FP8 General Matrix Multiply (GEMM) operations by introducing block-wise scaling factors for inputs A and B. This change is reflected in the updated JSON definitions, which now specify K-block-A and K-block-B constants and adjust the dimensionality of the A_scale and B_scale inputs accordingly. Furthermore, the PR transitions several gemm_fp8 workloads from using randomly generated input data to leveraging pre-defined tensors stored in newly added .safetensors files, ensuring consistent and reproducible testing.

Highlights

• Updated GEMM Definitions: Modified several gemm_fp8 definition files to introduce K-block-A and K-block-B constants, which define block sizes for scaling factors used in FP8 operations.
• Dimensionality Change for Scaling Factors: The A_scale and B_scale inputs in the GEMM definitions were updated from 1D to 2D arrays, incorporating the newly defined K-block-A and K-block-B dimensions for more granular scaling.
• New Safetensors Workload Data: A large number of .safetensors files were added across various gemm_fp8 directories, providing concrete tensor data for workloads.
• Workload Configuration Update: Existing gemm_fp8 workload configurations were updated to reference these new .safetensors files for their A, B, A_scale, and B_scale inputs, transitioning from randomly generated data to pre-defined, reproducible tensor data.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request updates several fp8 GEMM definitions to use block-wise scaling and adds corresponding workloads with pre-generated data. While the move to use .safetensors for workloads is a good improvement for deterministic benchmarking, I've found critical issues in the definition files. The Python reference implementation for dequantization is incorrect in all updated definitions, as it hasn't been updated to support the new block-wise scaling scheme. Additionally, one of the definitions (gemm_fp8_n7168_k2304.json) contains an invalid value for K-block-B that will lead to errors. These issues need to be addressed to ensure the correctness of the definitions and allow for successful verification.

[gemini-code-assist]

The reference implementation on line 73 is now incorrect due to the changes for block-wise scaling. The dequantization logic A_scale.view(-1, 1) is for per-tensor or per-row scaling and will fail with the new 2D A_scale and B_scale tensors. The reference implementation needs to be updated to correctly handle block-wise dequantization.

A correct implementation would be:

"reference": "import torch\n\ndef run(A, B, A_scale, B_scale):\n    A_scale_rep = A_scale.repeat_interleave(A.shape[1] // A_scale.shape[1], dim=1)\n    A_deq = A.to(torch.float32) * A_scale_rep\n    B_scale_rep = B_scale.repeat_interleave(B.shape[1] // B_scale.shape[1], dim=1)\n    B_deq = B.to(torch.float32) * B_scale_rep\n    C = torch.matmul(A_deq, B_deq.T)\n    return C.to(torch.bfloat16)"

Since I cannot comment on line 73 directly, I'm adding this comment here. Please update the reference field.

[gemini-code-assist]

The reference implementation on line 73 is now incorrect due to the changes for block-wise scaling. The dequantization logic A_scale.view(-1, 1) is for per-tensor or per-row scaling and will fail with the new 2D A_scale and B_scale tensors. The reference implementation needs to be updated to correctly handle block-wise dequantization.

A correct implementation would be:

"reference": "import torch\n\ndef run(A, B, A_scale, B_scale):\n    A_scale_rep = A_scale.repeat_interleave(A.shape[1] // A_scale.shape[1], dim=1)\n    A_deq = A.to(torch.float32) * A_scale_rep\n    B_scale_rep = B_scale.repeat_interleave(B.shape[1] // B_scale.shape[1], dim=1)\n    B_deq = B.to(torch.float32) * B_scale_rep\n    C = torch.matmul(A_deq, B_deq.T)\n    return C.to(torch.bfloat16)"

Since I cannot comment on line 73 directly, I'm adding this comment here. Please update the reference field.

[gemini-code-assist]

The reference implementation on line 73 is now incorrect due to the changes for block-wise scaling. The dequantization logic A_scale.view(-1, 1) is for per-tensor or per-row scaling and will fail with the new 2D A_scale and B_scale tensors. The reference implementation needs to be updated to correctly handle block-wise dequantization.

A correct implementation would be:

"reference": "import torch\n\ndef run(A, B, A_scale, B_scale):\n    A_scale_rep = A_scale.repeat_interleave(A.shape[1] // A_scale.shape[1], dim=1)\n    A_deq = A.to(torch.float32) * A_scale_rep\n    B_scale_rep = B_scale.repeat_interleave(B.shape[1] // B_scale.shape[1], dim=1)\n    B_deq = B.to(torch.float32) * B_scale_rep\n    C = torch.matmul(A_deq, B_deq.T)\n    return C.to(torch.bfloat16)"

Since I cannot comment on line 73 directly, I'm adding this comment here. Please update the reference field.

[gemini-code-assist]

The reference implementation on line 73 is now incorrect due to the changes for block-wise scaling. The dequantization logic A_scale.view(-1, 1) is for per-tensor or per-row scaling and will fail with the new 2D A_scale and B_scale tensors. The reference implementation needs to be updated to correctly handle block-wise dequantization.

A correct implementation would be:

"reference": "import torch\n\ndef run(A, B, A_scale, B_scale):\n    A_scale_rep = A_scale.repeat_interleave(A.shape[1] // A_scale.shape[1], dim=1)\n    A_deq = A.to(torch.float32) * A_scale_rep\n    B_scale_rep = B_scale.repeat_interleave(B.shape[1] // B_scale.shape[1], dim=1)\n    B_deq = B.to(torch.float32) * B_scale_rep\n    C = torch.matmul(A_deq, B_deq.T)\n    return C.to(torch.bfloat16)"

Since I cannot comment on line 73 directly, I'm adding this comment here. Please update the reference field.

[gemini-code-assist]

The reference implementation on line 73 is now incorrect due to the changes for block-wise scaling. The dequantization logic A_scale.view(-1, 1) is for per-tensor or per-row scaling and will fail with the new 2D A_scale and B_scale tensors. The reference implementation needs to be updated to correctly handle block-wise dequantization.

A correct implementation would be:

"reference": "import torch\n\ndef run(A, B, A_scale, B_scale):\n    A_scale_rep = A_scale.repeat_interleave(A.shape[1] // A_scale.shape[1], dim=1)\n    A_deq = A.to(torch.float32) * A_scale_rep\n    B_scale_rep = B_scale.repeat_interleave(B.shape[1] // B_scale.shape[1], dim=1)\n    B_deq = B.to(torch.float32) * B_scale_rep\n    C = torch.matmul(A_deq, B_deq.T)\n    return C.to(torch.bfloat16)"

Since I cannot comment on line 73 directly, I'm adding this comment here. Please update the reference field.

[gemini-code-assist]

The value of K-block-B is 5, but the K dimension is 2304. For block-wise scaling to work, the K dimension must be divisible by the number of blocks. Since 2304 % 5 != 0, this configuration is invalid and will cause errors during dequantization. Please verify and provide a correct value for K-block-B that is a divisor of 2304. For example, K-block-A is 18, which divides 2304. Perhaps K-block-B should also be 18?

[gemini-code-assist]

The reference implementation on line 73 is now incorrect due to the changes for block-wise scaling. The dequantization logic A_scale.view(-1, 1) is for per-tensor or per-row scaling and will fail with the new 2D A_scale and B_scale tensors. The reference implementation needs to be updated to correctly handle block-wise dequantization.

A correct implementation would be:

"reference": "import torch\n\ndef run(A, B, A_scale, B_scale):\n    A_scale_rep = A_scale.repeat_interleave(A.shape[1] // A_scale.shape[1], dim=1)\n    A_deq = A.to(torch.float32) * A_scale_rep\n    B_scale_rep = B_scale.repeat_interleave(B.shape[1] // B_scale.shape[1], dim=1)\n    B_deq = B.to(torch.float32) * B_scale_rep\n    C = torch.matmul(A_deq, B_deq.T)\n    return C.to(torch.bfloat16)"

Since I cannot comment on line 73 directly, I'm adding this comment here. Please update the reference field.