辅导 program、讲解 Java/Python程序语言
Assignment 4
220 pts (+ 140 bonus pts)
Due date: 11:59 PM PST, Dec 6 (Fri), 2024
Read the entire document carefully and provide appropriate answers based on the
context.
Figure 1. An illustration of the batched matrix multiplication (BMM) operator [ref].
Figure 1 illustrates the batched matrix multiplication (BMM) operator. BMM operators
are commonly used in self-attention of Transformer-based models such as GPT. BMM
operation can be represented in the loop nest presented in Figure 1.
Figure 2. Self-attention module
Figure 2 illustrates a self-attention module discussed in L04_Transformer slides. The
self-attention module consists of two BMMs. Generating Q, K, and V tensors requires
one BMM for each. See L04_Transformer slides for details.
inputA inputB output
for b in range(B):
for i in range(M):
for j in range(N):
for k in range(K):
output[b][i][j] += inputA[b][i][k] * inputB[b][k][j]Problem 1. Accelerator Design for an LLM
(170 pts + 90 bonus pts)
You are a DL accelerator architect in a major tech company. Your company has
developed a new large language model (LLM), and the model team is developing a
product like Chat-GPT based on the new LLM. To win the competition, your company
set the low latency service as their strategy. Your team is trying to develop a systolic
array-based DL accelerator to provide low latency to the new LLM.
Assumptions
Designing an accelerator for LLMs is a realistic and complicated problem. Therefore, we
will make the following assumptions to simplify the problem for this assignment.
- Like current state-of-the-art LLMs like GPT, we assume that the LLM is a variant
of the Transformer: We will assume that the model consists of 100 single-head
self-attention modules only (i.e., 100 of the blocks illustrated in Figure 2).
o We will only consider BMM computations in slides 20-26 in
L04_Transformer slides. (Ignore latencies for scale, mask, and softmax.)
o We will ignore all the other parts (e.g., embedding generation and the
linear layer after a self-attention block) for simplicity
- One “inference run” in this problem means the entire forward pass computation
for one set of input tensors (Q, K, and V).
- We will consider prefill stage only for this problem for simpilcity (i.e., no KV
cache)
A short summary of assumptions: The LLM is a Transformer model variant. In our
problems, we will ignore everything except BMMs in self-attention modules (+ Q, K, and
V projection). The LLM model consists of 100 self-attention modules.
* Note: The problems are organized to be aligned with the design process of an
accelerator. Note that this is not a complete list; the design process in this assignment
covers some essential parts of the design process, not everything.
Step 1. Workload analysis (70 pts)
Your team analyzed the new LLM and listed up useful information regarding the model
as follows:
- Number of attention blocks: 100
- Embedding dimension (d): 2048
- Query/Key/Value dimensions (Dk and Dv): 4096
- Sequence length (L): 1024 (fixed)
- Batch size (B): 1 – 128 (varied)
* Review L04_Transformer slides for the conventions Problem 1-1. Assuming the batch size is 1, what is the total number of MAC operations
to complete one inference? In addition to the answer, explain how you computed it and
why you took that approach. (30 pts)
Problem 1-2. Assuming the batch size is 1, what is the minimum number of PEs to
complete the LLM inference within one second (i.e., assuming 100% utilization of PEs
without any delay in the memory and network-on-chip, what is the required number of
PEs to meet the goal)? Assume the clock frequency of your accelerator is 1.5 GHz.
Assume each PE can compute one MAC in each cycle. In addition to the answer,
explain how you computed it and why you took that approach. (40 pts)
Step 2) Hardware Design Parameter Exploration
(50 pts + 60 bonus pts)
Based on the workload analysis, your team decided to design a systolic array with 4096
PEs (2
12
PEs).
Problem 2-1. There are many ways to organize a 2D PE array for the systolic array:
1x4096, 2x2048, … 64x64. Assuming the batch size is 1, which PE array organization
do you suggest? Qualitatively and quantitatively support your suggestion. (50 pts)
* Hint: Consider the PE array utilization
Problem 2-2. [Bonus problem] Assuming the batch size is 1, what is the minimum
memory size that minimizes the DRAM traffic? Minimum DRAM traffic is achieved when
you access DRAM only for fetching model parameters (weights) and the first input tensor
and sending the final output back to the DRAM. Extra DRAM traffic (which you need to
avoid) is incurred if you do not have sufficient memory to keep all the necesasry
intermediate tensors with model weights. In addition to the answer, explain how you
computed it and why you took that approach. Assume that all data are in FP8 format. (60
pts)
Step 3) Dataflow Exploration (50 pts + 30 bonus pts)
After many technical meetings, your team decided to design a 128 x 32 systolic array.
Problem 3-1. If the batch size is 1, which dimensions do you suggest to parallelize over
the rows and columns of the PE array? Also, based on your selection, what is the
processing order you suggest (i.e., the order of B, M, N, and K dimension movement)?
For simplicity, we will only consider the BMMs for Q, K, and V projections. Since
computations for Q, K, and V projections are the same, you can consider only one of them
in your answer.
Qualitatively and quantitatively support your answers. (50 pts)
Problem 3-2. [Bonus problem] If the batch size is 32, does that change your suggestion?
For simplicity, only consider the BMMs for Q, K, and V projections. Qualitatively discuss
the implication of larger batch size. (30 pts)
Problem 2. Systolic array and Dataflow Description
(50 pts + bonus 50 pts)
You are trying to compute a matrix multiplication where the matrix dimensions are as
follows: [M: 16, N: 4, K: 16].
Figure 3. Matrix multiplication
Assume the following for this problem:
• MAC computation on each PE requires one clock cycle
• Data forwarding between adjacent PEs requires one clock cycle
o Like examples in our slides, a PE can receive incoming data from a
neighbor and compute MAC at the same cycle. However, the output of the
MAC forwarded to the next PE is available in the next cycle for the next
PE.
o Similarly, operand a PE forward to the next PE is available in the next
cycle.
o Review L06_DL Accelerators slide deck carefully to fully understand how
a systolic array works. Assume the same style when you solve problems.
• All the necessary data are pre-loaded to a sufficiently large on-chip global buffer
for matrices A and B.
• You don't have to count the cycles to send output from the global shared buffer
(L2) to DRAM; only consider cycles consumed within the PE array
(Hint: The assumptions above create the same environment as the systolic array
examples in lecture slides)
Problem 4-1. Assuming a systolic array of 16x16 PEs (row: parallelize the K dimension,
column: parallelize the M dimension), estimate the latency of the systolic array for
computing the matrix multiplication in clock cycle counts. In addition to the answer,
explain how you computed it and why you took that approach. (50 pts)
Problem 4-2. [Bonus problem] Describe the dataflow given in Problem 4-1. You can
use a loop nest or the data-centric directives (50 pts)