본문 바로가기
Project/RISC-V 기반 LLMs Accelerator Design

[1] RISC-V 기반 LLMs Accelerator Design (feat. cursor)

by 한PU 2025. 4. 15.
반응형

요즘 Cursor가 진짜 야무지다.

코딩도 잘하고, 정리도 잘해주고...

cursor와 함께 RISC-V 기반 LLMs Accelerator를 설계해보려고 한다.

Cursor 프롬프트

 

기능 구조

# RISC-V 기반 LLM 가속기 기능 구조

## 개요
RISC-V CPU 코어와 LLM(Large Language Model) 연산 가속을 위한 하드웨어 블록을 통합한 시스템 온 칩 (SoC) 형태를 가정합니다. 호스트 CPU(RISC-V)는 전체 시스템 제어 및 가속기 오프로딩을 담당하고, 가속기는 LLM의 핵심 연산을 병렬 처리하여 성능을 극대화합니다.

## 주요 구성 요소

1.  **RISC-V 코어 (RISC-V Core):**
    *   **역할:**
        *   운영체제 및 애플리케이션 실행
        *   LLM 모델 로딩 및 전처리/후처리
        *   가속기 제어 및 작업 스케줄링
        *   가속기와 데이터 통신
    *   **특징:**
        *   표준 RISC-V 명령어 세트 아키텍처 (ISA) 지원
        *   (선택 사항) LLM 가속기 제어를 위한 커스텀 명령어 확장

2.  **LLM 가속 유닛 (LLM Accelerator Unit):**
    *   **처리 요소 배열 (Processing Element Array - PE Array):**
        *   **구조:** 다수의 PE가 2D 또는 1D 형태로 배열됨.
        *   **기능:** 행렬 곱셈 (GEMM), 벡터 연산 등 LLM의 핵심 연산(주로 Transformer의 Feed-Forward Network, Attention 계산)을 병렬로 수행합니다. 각 PE는 MAC(Multiply-Accumulate) 유닛을 포함하는 것이 일반적입니다.
        *   **데이터 흐름:** Systolic array, Output Stationary, Weight Stationary 등 다양한 데이터 흐름 아키텍처 적용 가능.
    *   **온칩 메모리 (On-chip Memory / Scratchpad):**
        *   **종류:** SRAM 기반의 고속 버퍼 메모리.
        *   **용도:** 가중치(Weights), 활성화(Activations), 중간 결과 등을 저장하여 데이터 재사용성을 높이고 외부 메모리 접근 지연 시간을 줄입니다. 계층적 구조(L1, L2 캐시와 유사)를 가질 수 있습니다.
    *   **데이터 로더/저장 유닛 (Data Loader/Storer):**
        *   **기능:** 외부 메모리(DRAM)와 온칩 메모리 간의 데이터 전송을 담당합니다. DMA(Direct Memory Access) 엔진을 사용하여 CPU 개입 없이 대량 데이터 전송을 효율적으로 처리합니다.
        *   **최적화:** 데이터 압축/해제, 희소(Sparsity) 행렬 처리 기능 포함 가능.
    *   **제어 로직 (Control Logic):**
        *   **기능:** PE 배열, 온칩 메모리, 데이터 로더/저장 유닛의 동작을 제어하고 동기화합니다. RISC-V 코어로부터 명령을 받아 연산 흐름을 관리합니다.
        *   **구현:** 마이크로 코드 시퀀서 또는 하드웨어 상태 머신(FSM)으로 구현될 수 있습니다.

3.  **메모리 인터페이스 (Memory Interface):**
    *   **기능:** SoC와 외부 주 메모리(예: DDR4/5, LPDDR4/5, HBM) 간의 데이터 통신을 위한 인터페이스를 제공합니다.
    *   **고려 사항:** 높은 대역폭(Bandwidth)과 낮은 지연 시간(Latency)이 중요합니다.

4.  **상호 연결망 (Interconnect Fabric):**
    *   **기능:** RISC-V 코어, LLM 가속 유닛, 메모리 인터페이스 및 기타 주변 장치들 간의 데이터 및 제어 신호를 전송하는 버스 또는 네트워크입니다.
    *   **종류:** AXI (Advanced eXtensible Interface), NoC (Network-on-Chip) 등 다양한 프로토콜 및 구조 사용 가능.

## 동작 흐름 (예시: 행렬 곱셈 연산)

1.  **RISC-V 코어:** LLM 애플리케이션 실행 중 행렬 곱셈 연산 필요시, 가속기에 해당 연산을 오프로딩.
    *   입력 행렬(활성화) 및 가중치 행렬의 메모리 주소 전달.
    *   가속기 제어 레지스터 설정 (연산 종류, 크기 등).
    *   가속기 시작 신호 전달 (커스텀 명령어 또는 메모리 매핑된 I/O 사용).
2.  **제어 로직:** RISC-V 코어로부터 명령 수신 및 해석.
3.  **데이터 로더:** 외부 메모리에서 필요한 가중치 및 활성화 데이터를 온칩 메모리로 로딩.
4.  **제어 로직:** 로딩된 데이터를 PE 배열에 적절히 분배하고 연산 시작 제어.
5.  **PE 배열:** 분배된 데이터를 사용하여 병렬로 행렬 곱셈 (MAC 연산) 수행.
6.  **제어 로직:** 연산 완료 확인.
7.  **데이터 저장 유닛:** 계산된 결과(출력 행렬)를 온칩 메모리에서 외부 메모리로 저장.
8.  **제어 로직:** 연산 완료 및 결과 저장 완료를 RISC-V 코어에 알림 (인터럽트 또는 상태 레지스터 사용).
9.  **RISC-V 코어:** 가속 완료 확인 후, 결과를 후속 연산에 사용.

 

레퍼런스 코드

- python 사용

- cursor가 짜줌 ㅋㅋ

 

import numpy as np
import time

# --- 시뮬레이션 구성 요소 ---

def ExternalMemory(size_in_mb):
    """외부 주 메모리를 시뮬레이션합니다."""
    print(f"[Memory] {size_in_mb}MB 외부 메모리 초기화.")
    # 실제 데이터 저장을 위한 딕셔너리 사용 (주소 -> 데이터)
    return {'data': {}, 'size_mb': size_in_mb}

def OnChipMemory(size_in_kb):
    """온칩 스크래치패드 메모리를 시뮬레이션합니다."""
    print(f"[Memory] {size_in_kb}KB 온칩 메모리 초기화.")
    return {'data': {}, 'size_kb': size_in_kb}

def DataLoaderStorer(external_mem, on_chip_mem):
    """데이터 로더/저장 유닛을 시뮬레이션합니다."""
    def load(ext_addr, on_chip_addr, size):
        """외부 메모리 -> 온칩 메모리 데이터 로딩."""
        # print(f"[DMA] Loading {size} bytes from External:{ext_addr} to OnChip:{on_chip_addr}")
        # 실제 구현에서는 데이터 크기, 주소 유효성 검사 등이 필요
        if ext_addr in external_mem['data']:
            on_chip_mem['data'][on_chip_addr] = external_mem['data'][ext_addr][:size] # 실제 데이터 복사 시뮬레이션 (간략화)
        else:
            print(f"[DMA Error] External memory address {ext_addr} not found.")
        # time.sleep(0.01) # 데이터 전송 지연 시뮬레이션

    def store(on_chip_addr, ext_addr, size):
        """온칩 메모리 -> 외부 메모리 데이터 저장."""
        # print(f"[DMA] Storing {size} bytes from OnChip:{on_chip_addr} to External:{ext_addr}")
        if on_chip_addr in on_chip_mem['data']:
            external_mem['data'][ext_addr] = on_chip_mem['data'][on_chip_addr][:size] # 실제 데이터 복사 시뮬레이션 (간략화)
            # 온칩 메모리 데이터는 유지될 수도, 삭제될 수도 있음 (정책에 따라 다름)
            # del on_chip_mem['data'][on_chip_addr] # 예: 저장 후 삭제
        else:
            print(f"[DMA Error] On-chip memory address {on_chip_addr} not found.")
        # time.sleep(0.01) # 데이터 전송 지연 시뮬레이션

    return load, store

def PEArray(rows, cols):
    """PE 배열을 시뮬레이션합니다. (행렬 곱셈 기능)"""
    print(f"[PE Array] {rows}x{cols} PE 배열 초기화.")
    pe_count = rows * cols

    def execute_gemm(matrix_a, matrix_b):
        """행렬 곱셈(GEMM) 연산을 수행합니다."""
        # print(f"[PE Array] Executing GEMM: {matrix_a.shape} x {matrix_b.shape}")
        if not isinstance(matrix_a, np.ndarray) or not isinstance(matrix_b, np.ndarray):
             raise TypeError("Inputs must be NumPy arrays for PE Array simulation")
        if matrix_a.shape[1] != matrix_b.shape[0]:
            raise ValueError("Incompatible matrix dimensions for multiplication.")

        # 실제 PE 배열 연산은 병렬 처리되지만, 여기서는 numpy로 기능만 시뮬레이션
        result = np.dot(matrix_a, matrix_b)
        # time.sleep(0.05) # 연산 지연 시간 시뮬레이션
        # print(f"[PE Array] GEMM complete. Result shape: {result.shape}")
        return result

    return execute_gemm

def ControlLogic(pe_array_executor, data_loader, data_storer, on_chip_mem_ref):
    """제어 로직을 시뮬레이션합니다."""
    print("[Control] 제어 로직 초기화.")
    tasks = [] # (operation_type, params)

    def receive_instruction(instruction):
        """호스트로부터 명령을 수신합니다."""
        print(f"[Control] 명령 수신: {instruction['type']}")
        tasks.append(instruction)

    def run_task():
        """큐에 있는 작업을 처리합니다."""
        if not tasks:
            # print("[Control] 처리할 작업 없음.")
            return None # 처리 완료 또는 할 일 없음

        instruction = tasks.pop(0)
        op_type = instruction['type']
        params = instruction['params']

        if op_type == 'gemm':
            print("[Control] GEMM 작업 시작.")
            # 1. 데이터 로딩
            print("[Control] 데이터 로딩 시작...")
            data_loader(params['a_ext_addr'], params['a_onchip_addr'], params['a_size'])
            data_loader(params['b_ext_addr'], params['b_onchip_addr'], params['b_size'])
            print("[Control] 데이터 로딩 완료.")

            # 온칩 메모리에서 데이터 가져오기 (시뮬레이션 편의상)
            matrix_a = on_chip_mem_ref['data'].get(params['a_onchip_addr'])
            matrix_b = on_chip_mem_ref['data'].get(params['b_onchip_addr'])

            if matrix_a is None or matrix_b is None:
                print("[Control Error] 온칩 메모리에서 데이터 로드 실패.")
                return 'error'

            # 2. PE 배열 연산 실행
            print("[Control] PE 배열 연산 실행...")
            result_matrix = pe_array_executor(matrix_a, matrix_b)
            print("[Control] PE 배열 연산 완료.")

            # 임시로 결과를 온칩 메모리에 저장 (시뮬레이션)
            on_chip_mem_ref['data'][params['res_onchip_addr']] = result_matrix

            # 3. 결과 저장
            print("[Control] 결과 저장 시작...")
            data_storer(params['res_onchip_addr'], params['res_ext_addr'], params['res_size'])
            # 결과 저장 후 온칩 메모리 정리 (선택적)
            if params['res_onchip_addr'] in on_chip_mem_ref['data']:
                del on_chip_mem_ref['data'][params['res_onchip_addr']]
            print("[Control] 결과 저장 완료.")

            print("[Control] GEMM 작업 완료.")
            return 'completed' # 작업 완료 상태 반환
        else:
            print(f"[Control] 알 수 없는 연산 타입: {op_type}")
            return 'unknown_op'

    return receive_instruction, run_task

def HostCPU(accelerator_ctrl_receiver):
    """호스트 CPU(RISC-V)의 역할을 시뮬레이션합니다."""
    print("[Host] 호스트 CPU 시작.")

    def prepare_data(external_mem):
        """연산에 필요한 데이터를 외부 메모리에 준비합니다."""
        print("[Host] 데이터 준비 중...")
        # 예시 데이터 생성
        matrix_a = np.random.rand(64, 128).astype(np.float32)
        matrix_b = np.random.rand(128, 256).astype(np.float32)

        # 외부 메모리 주소 할당 (간단한 시뮬레이션)
        addr_a = 0x1000
        addr_b = 0x2000
        addr_res = 0x3000

        external_mem['data'][addr_a] = matrix_a
        external_mem['data'][addr_b] = matrix_b
        print(f"[Host] 데이터 준비 완료. A: {addr_a}, B: {addr_b}, Result: {addr_res}")
        return addr_a, addr_b, addr_res, matrix_a.nbytes, matrix_b.nbytes, (matrix_a.shape[0] * matrix_b.shape[1] * np.dtype(np.float32).itemsize) # 결과 크기

    def offload_gemm_task(ext_a, ext_b, ext_res, size_a, size_b, size_res):
        """GEMM 연산을 가속기에 오프로딩합니다."""
        print("[Host] GEMM 작업을 가속기에 오프로딩...")
        # 온칩 메모리 주소 계획 (간단한 시뮬레이션)
        onchip_a = 0x100
        onchip_b = 0x200
        onchip_res = 0x300

        instruction = {
            'type': 'gemm',
            'params': {
                'a_ext_addr': ext_a, 'a_onchip_addr': onchip_a, 'a_size': size_a,
                'b_ext_addr': ext_b, 'b_onchip_addr': onchip_b, 'b_size': size_b,
                'res_ext_addr': ext_res, 'res_onchip_addr': onchip_res, 'res_size': size_res
            }
        }
        accelerator_ctrl_receiver(instruction) # 제어 로직에 명령 전달
        print("[Host] 작업 전달 완료.")

    def check_completion(external_mem, result_addr):
         """가속기 작업 완료 및 결과 확인."""
         # 실제 시스템에서는 인터럽트나 상태 레지스터 폴링 사용
         print("[Host] 가속기 완료 확인 중...")
         # 여기서는 외부 메모리에 결과가 쓰여졌는지 확인하는 방식으로 간략화
         if result_addr in external_mem['data']:
             print("[Host] 작업 완료 확인. 결과 데이터 접근 가능.")
             return external_mem['data'][result_addr]
         else:
             # print("[Host] 아직 작업 미완료...")
             return None

    return prepare_data, offload_gemm_task, check_completion

# --- 시뮬레이션 실행 ---

if __name__ == "__main__":
    print("--- RISC-V LLM 가속기 시뮬레이션 시작 ---")

    # 1. 하드웨어 구성 요소 초기화
    external_memory = ExternalMemory(size_in_mb=1024) # 1GB DRAM
    on_chip_memory = OnChipMemory(size_in_kb=512)    # 512KB SRAM
    data_loader, data_storer = DataLoaderStorer(external_memory, on_chip_memory)
    pe_array_executor = PEArray(rows=8, cols=8) # 8x8 PE 배열
    ctrl_receive, ctrl_run = ControlLogic(pe_array_executor, data_loader, data_storer, on_chip_memory)
    host_prepare, host_offload, host_check = HostCPU(ctrl_receive)

    print("\n--- 시뮬레이션 흐름 ---")

    # 2. 호스트: 데이터 준비
    addr_a, addr_b, addr_res, size_a, size_b, size_res = host_prepare(external_memory)
    original_a = external_memory['data'][addr_a]
    original_b = external_memory['data'][addr_b]

    # 3. 호스트: 작업 오프로딩
    host_offload(addr_a, addr_b, addr_res, size_a, size_b, size_res)

    # 4. 가속기: 작업 실행 (제어 로직이 주도)
    status = None
    max_wait_cycles = 10
    cycles = 0
    while status != 'completed' and cycles < max_wait_cycles:
        print(f"\n--- Cycle {cycles + 1} ---")
        status = ctrl_run() # 제어 로직의 작업 처리 함수 호출
        if status == 'error' or status == 'unknown_op':
            print(f"[Simulation Error] 가속기 오류 발생: {status}")
            break
        time.sleep(0.1) # 시뮬레이션 시간 진행
        cycles += 1

    if status == 'completed':
      print("\n--- 작업 완료 ---")
      # 5. 호스트: 결과 확인
      result_matrix = host_check(external_memory, addr_res)

      if result_matrix is not None:
          print(f"[Host] 결과 매트릭스 확인 완료. Shape: {result_matrix.shape}")

          # (선택적) 결과 검증
          print("[Host] Numpy를 사용한 결과 검증 중...")
          expected_result = np.dot(original_a, original_b)
          if np.allclose(result_matrix, expected_result):
              print("[Host] 검증 성공: 가속기 결과가 예상과 일치합니다.")
          else:
              print("[Host] 검증 실패: 가속기 결과가 예상과 다릅니다.")
      else:
          print("[Host] 결과 데이터를 외부 메모리에서 찾을 수 없습니다.")

    elif cycles >= max_wait_cycles:
        print("\n[Simulation Timeout] 가속기 작업이 시간 내에 완료되지 않았습니다.")
    else:
        print("\n[Simulation End] 시뮬레이션이 오류 상태로 종료되었습니다.")


    print("\n--- 시뮬레이션 종료 ---")

 

레퍼런스 결과

레퍼런스 코드 결과

성능 비교가 추가되면 좋을듯

반응형