커널 설계

첫 번째 벡터 커널

상수 덧셈

목표

벡터의 모든 요소에 대해 \(\text{out}[i] = \text{in}[i] + 1\)을 계산합니다. 이 패턴은 잔차 덧셈과 바이어스 덧셈을 다룹니다.

데이터 이동

벡터를 HBM에서 DM으로 옮기고, 각 DM Slice를 Fetch, Collect, Vector Engine, Commit을 거쳐 스트리밍한 뒤, 결과를 DM에서 HBM으로 다시 옮깁니다.

flowchart TB
    HOST[Host] <-->|PCIe DMA| HBM[(HBM)]
    HBM <-->|Tensor DMA| DM[(DM)]

    subgraph TU[Tensor Unit]
        direction TB
        FE[Fetch] --> CO[Collect] --> VE["Vector AddFxp +1"] --> CM[Commit]
    end

    DM -->|stream| FE
    CM -->|stream| DM

디바이스 소스

이 커널은 Chip 하나와 Cluster 두 개 중 하나를 사용하며, 해당 Cluster의 Slice 256개를 모두 사용합니다. 각 Slice는 i32 값 여덟 개를 받습니다. TagMode::Zero는 매 사이클마다 Vector Engine을 실행합니다. to_dm은 벡터를 여러 Slice에 분배하고, begin → fetch → collect → vector_init → vector_intra_slice_tag → vector_fxp → vector_final → commit 체인은 각 Slice를 한 번의 패스로 처리합니다.

디바이스 소스: src/kernel/constant_add_kernel.rs.

use furiosa_opt_std::prelude::*;

axes![A = 2048];

pub type Chip = m![1];
pub type Cluster = m![1 # 2];
pub type Slice = m![A / 8 # 256];

#[device(chip = 1)]
pub fn constant_add_kernel(device: &mut Device, input: &HbmTensor<i32, Chip, m![A]>) -> HbmTensor<i32, Chip, m![A]> {
    // HBM → DM: split 2048 elements across 256 slices (8 elements per slice)
    let dm = input.to_dm::<Cluster, Slice, m![A % 8]>(&mut device.tdma);

    let result = device
        .main
        .begin(dm.view())
        // Fetch: stream 8-element packets from DM into the pipeline
        .fetch::<m![1], m![A % 8]>()
        // Collect: normalize the stream into 32-byte flits (8 × i32)
        .collect::<m![1], m![A % 8]>()
        // Vector Engine: enter pipeline and arm unconditionally
        .vector_init()
        .vector_intra_slice_tag(TagMode::Zero)
        // Add the scalar constant 1 to every element
        .vector_fxp(FxpBinaryOp::AddFxp, 1)
        // Exit VE and commit: trim the packet to the commit width, then write
        // results back to DM
        .vector_final()
        .commit_trim::<m![A % 8]>()
        .commit::<m![A % 8]>();

    // DM → HBM
    result.to_hbm(&mut device.tdma)
}

호스트 프로그램과 오라클

호스트 프로그램은 입력을 생성하여 HBM에 업로드하고, 디바이스 함수를 실행한 뒤, 테스트에서 결과를 호스트 수식과 비교합니다.

기본 템플릿 호스트 프로그램과 오라클: src/constant_add.rs.

use furiosa_opt_std::prelude::*;
use {{ crate_name }}::kernel::constant_add_kernel::{A, constant_add_kernel};
use rand::SeedableRng;
use rand::rngs::SmallRng;

#[tokio::main]
async fn main() -> Result<(), Error> {
    let mut device = Device::new(constant_add_kernel.topology())?;
    let mut rng = SmallRng::seed_from_u64(42);
    let input = HostTensor::<i32, m![A]>::rand(&mut rng);
    let in_hbm = input.to_hbm(&mut device.pdma).await?;
    let _out_hbm = launch(constant_add_kernel, (&mut device, &in_hbm)).await?;
    println!("Constant Add: kernel ran");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn matches_reference() {
        let mut device = Device::new(constant_add_kernel.topology()).unwrap();

        let mut rng = SmallRng::seed_from_u64(42);
        let input = HostTensor::<i32, m![A]>::rand(&mut rng);
        let in_hbm = input.to_hbm(&mut device.pdma).await.unwrap();

        // Reference: out[i] = in[i] + 1.
        let in_buf: Vec<i32> = input.into_vec();
        let expected: Vec<i32> = in_buf.iter().map(|&x| x.wrapping_add(1)).collect();

        let out_hbm = launch(constant_add_kernel, (&mut device, &in_hbm)).await.unwrap();

        let actual: Vec<i32> = out_hbm.to_host::<m![A]>(&mut device.pdma).await.unwrap().into_vec();
        for (i, (&e, &a)) in expected.iter().zip(&actual).enumerate() {
            assert_eq!(e, a, "constant_add mismatch at i={i}: expected {e}, actual {a}");
        }
    }
}

패턴별 매핑과 메모리

TCP에는 네 단계로 중첩된 하드웨어 계층이 있습니다.

텐서 타입은 요소 타입과 각 논리 축이 이 계층 전반에 어떻게 분포되는지를 인코딩합니다. 예를 들어 DmTensor<bf16, m![1], m![1 # 2], m![A / 8 # 256], m![A % 8]>는 축 A가 256개의 Slice에 나뉘고 Slice마다 요소가 8개씩 있는 bf16 벡터를 나타냅니다. A의 각 요소는 하나의 Slice 안에서 명확하게 정의된 위치를 하나씩 가집니다.

매핑 연산자는 축이 어디로 배치되는지를 기술합니다.

• /는 축을 스트라이드 단위로 분할합니다.
• %는 단위 내 개수를 나타냅니다.
• #는 매핑을 하드웨어 단위 개수에 맞게 패딩합니다.
• =는 뷰가 기존 항을 포괄하는 범위를 줄이며, 타일은 이 방식으로 텐서의 한 조각을 선택합니다.

%와 =는 모두 축을 좁히므로 둘의 차이를 짚어 둘 필요가 있습니다. %는 새로운 내부 항을 도입합니다. m![B / 2048, B % 2048]에서 첫 번째 항은 B의 2048개 요소 블록을 세고, 두 번째 항은 한 블록 안의 위치를 셉니다. =는 항을 그대로 두고 뷰가 읽는 범위를 줄입니다. 따라서 m![B / 2048 = 1 # 2]는 두 블록 중 하나를 선택하며, 이때 # 2는 선택된 범위가 아닌 전체 축의 스트라이드를 유지합니다. 타일링된 커널은 m![B / 2048 = 1 # 2, B % 2048]처럼 둘을 함께 사용하여 한 블록을 고르고 그 안의 모든 위치를 포괄합니다. 일반 규칙은 타일링을 참조하십시오.

Tensor Unit 텐서는 파이프라인 반복에 Time을, 한 반복 내의 요소에 Packet을 추가로 사용합니다. Tensor Unit 파이프라인은 Fetch → Switch → Collect → Contraction → Vector → Cast → Transpose → Commit 순서입니다. Switch Engine은 Slice 사이를 연결하며, 대부분의 단계는 Slice 안에서 독립적으로 동작합니다.

이러한 선택은 상수 덧셈 패턴을 특수화합니다. 전체 매핑 모델은 텐서 매핑을, 메모리 이동은 텐서 이동을, 파이프라인 API는 텐서 계산을 참조하십시오.

고정 벡터 피연산자

요소별 곱셈

목표

\(\text{out}[i] = \text{lhs}[i] \times \text{rhs}[i]\)를 계산합니다. 이 패턴은 게이트 스케일링, GLU, 어텐션 스케일링에 적용됩니다.

매핑과 데이터 이동

sub 컨텍스트는 rhs를 VRF에 미리 적재하고, main 컨텍스트는 lhs를 스트리밍하면서 Vector Engine 사이클마다 VRF를 읽습니다. 두 DM 영역은 겹치면 안 됩니다.

flowchart TB
    LHS_HBM[(lhs HBM)] -->|Tensor DMA| LHS_DM[(lhs DM)]
    RHS_HBM[(rhs HBM)] -->|Tensor DMA| RHS_DM[(rhs DM)]

    subgraph sub[sub context]
        direction LR
        sFE[Fetch] --> sCO[Collect] --> VRF[(VRF)]
    end

    subgraph main[main context]
        direction LR
        mFE[Fetch] --> mCO[Collect] --> VE["Vector MulInt"] --> CM[Commit]
    end

    RHS_DM --> sFE
    LHS_DM --> mFE
    VRF --> VE
    CM --> OUT_DM[(result DM)]
    OUT_DM -->|Tensor DMA| OUT_HBM[(HBM)]

모든 디바이스 커널은 별도의 하드웨어 자원 위에서 동작하는 device.main 실행 컨텍스트와 device.sub 실행 컨텍스트를 가집니다. main은 주요 계산을 수행하고, sub는 보통 피연산자를 프리페치합니다. main은 sub가 아직 생성하지 않은 데이터가 필요하면 대기하며, 두 컨텍스트는 평면 구조의 Chip 내장 SRAM을 공유합니다.

자동 배치를 사용하려면 DM 주소를 생략하십시오. 알고리즘에 서로 겹치지 않는 명시적 주소가 필요한 경우에만 _at API를 사용하십시오. Tu const 제네릭은 텐서가 { Tu::Main }과 { Tu::Sub } 중 어느 쪽을 통해 흐르는지 식별합니다. sub 컨텍스트는 Fetch, Collect, .to_vrf()를 거쳐 rhs_dm을 적재하고, main은 그 고정 VRF 피연산자와 함께 lhs_dm을 MulInt로 스트리밍합니다. 자원 의존성이 허용하면 두 컨텍스트는 동시에 실행됩니다.

디바이스 소스

디바이스 소스: src/kernel/elementwise_mul_kernel.rs.

use furiosa_opt_std::prelude::*;

axes![A = 2048];

pub type Chip = m![1];
pub type Cluster = m![1 # 2];
pub type Slice = m![A / 8 # 256];

#[device(chip = 1)]
pub fn elementwise_mul_kernel(
    device: &mut Device,
    lhs: &HbmTensor<i32, Chip, m![A]>,
    rhs: &HbmTensor<i32, Chip, m![A]>,
) -> HbmTensor<i32, Chip, m![A]> {
    // Move both operands from HBM to DM (DM placement is assigned automatically).
    let lhs_dm = lhs.to_dm::<Cluster, Slice, m![A % 8]>(&mut device.tdma);
    let rhs_dm = rhs.to_dm::<Cluster, Slice, m![A % 8]>(&mut device.tdma);

    // Sub context: load rhs into VRF (runs concurrently with the main context below).
    // VRF holds a per-slice operand that the Vector Engine reads every cycle.
    let rhs_vrf: VrfTensor<i32, Chip, Cluster, Slice, m![A % 8]> = device
        .sub
        .begin(rhs_dm.view())
        .fetch::<m![1], m![A % 8]>()
        .collect::<m![A % 8 / 8], m![A % 8 % 8]>()
        .to_vrf();

    // Main context: multiply every lhs element by its rhs counterpart from VRF
    let result = device
        .main
        .begin(lhs_dm.view())
        .fetch::<m![1], m![A % 8]>()
        .collect::<m![1], m![A % 8]>()
        .vector_init()
        .vector_intra_slice_tag(TagMode::Zero)
        // Each slice multiplies its 8 lhs elements by the matching 8 rhs elements in VRF
        .vector_fxp(FxpBinaryOp::MulInt, &rhs_vrf)
        .vector_final()
        .commit_trim::<m![A % 8]>()
        .commit::<m![A % 8]>();

    result.to_hbm(&mut device.tdma)
}

호스트 프로그램과 오라클

기본 템플릿 호스트 프로그램과 오라클: src/elementwise_mul.rs.

use furiosa_opt_std::prelude::*;
use {{ crate_name }}::kernel::elementwise_mul_kernel::{A, elementwise_mul_kernel};
use rand::SeedableRng;
use rand::rngs::SmallRng;

#[tokio::main]
async fn main() -> Result<(), Error> {
    let mut device = Device::new(elementwise_mul_kernel.topology())?;
    let mut rng = SmallRng::seed_from_u64(42);
    let lhs = HostTensor::<i32, m![A]>::rand(&mut rng);
    let rhs = HostTensor::<i32, m![A]>::rand(&mut rng);
    let lhs_hbm = lhs.to_hbm(&mut device.pdma).await?;
    let rhs_hbm = rhs.to_hbm(&mut device.pdma).await?;
    let _out_hbm = launch(elementwise_mul_kernel, (&mut device, &lhs_hbm, &rhs_hbm)).await?;
    println!("Elementwise Mul: kernel ran");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn matches_reference() {
        let mut device = Device::new(elementwise_mul_kernel.topology()).unwrap();

        let mut rng = SmallRng::seed_from_u64(42);
        let lhs = HostTensor::<i32, m![A]>::rand(&mut rng);
        let rhs = HostTensor::<i32, m![A]>::rand(&mut rng);

        let lhs_hbm = lhs.to_hbm(&mut device.pdma).await.unwrap();
        let rhs_hbm = rhs.to_hbm(&mut device.pdma).await.unwrap();

        // Reference: out[i] = lhs[i] * rhs[i].
        let lhs_buf: Vec<i32> = lhs.into_vec();
        let rhs_buf: Vec<i32> = rhs.into_vec();
        let expected: Vec<i32> = lhs_buf.iter().zip(&rhs_buf).map(|(&a, &b)| a.wrapping_mul(b)).collect();

        let out_hbm = launch(elementwise_mul_kernel, (&mut device, &lhs_hbm, &rhs_hbm)).await.unwrap();

        let actual: Vec<i32> = out_hbm.to_host::<m![A]>(&mut device.pdma).await.unwrap().into_vec();
        for (i, (&e, &a)) in expected.iter().zip(&actual).enumerate() {
            assert_eq!(e, a, "elementwise_mul mismatch at i={i}: expected {e}, actual {a}");
        }
    }
}

축약 패턴

내적

목표

\(\sum_i x_i y_i\)를 계산합니다. 이 패턴은 어텐션 점수와 유사도 계산에 적용됩니다.

매핑과 데이터 이동

한 피연산자는 파이프라인을 통해 흐르고, sub 컨텍스트는 다른 피연산자를 TRF에 저장합니다. sub 컨텍스트는 Fetch, Collect, .to_trf()를 거쳐 해당 피연산자를 적재합니다.

.contract_outer()는 32바이트 플릿을 짝지어 64바이트 패킷으로 만들고, 고정된 TRF 피연산자를 읽습니다. 짝짓기는 Stream Adapter가 수행하고, 고정 피연산자는 TRF Sequencer가 읽습니다. 두 피연산자는 모두 레인별 곱셈기로 들어갑니다. .contract_packet()은 하드웨어 리덕션 트리를 통해 곱을 공간적으로 리덕션합니다. .contract_time::<m![1]>()은 시간에 걸쳐 누적하여 Slice마다 스칼라 하나를 생성합니다. .contract_lane()은 레인을 접으며, Lane = m![1]일 때는 자명한 연산입니다. .cast()는 f32 누산기를 bf16으로 변환합니다.

축약은 많은 항에 대해 곱셈-누산을 수행하므로 확장된 타입이 필요합니다. 축약은 i4와 i8을 i32로 확장하고, fp8과 bf16을 f32로 확장합니다. .align()과 .contract()에 붙은 어노테이션이 이 확장된 타입을 선택합니다.

디바이스 소스

디바이스 소스: src/kernel/dot_product_kernel.rs.

use furiosa_opt_std::prelude::*;

axes![A = 2048];

pub type Chip = m![1];
pub type Cluster = m![1 # 2];
pub type Slice = m![1 # 256]; // 1 active slice; m![A / 8 # 256] would distribute across all 256
pub type Time = m![1]; // No temporal iteration
pub type Lane = m![1]; // No lane parallelism

#[device(chip = 1)]
pub fn dot_product_kernel(
    device: &mut Device,
    lhs: &HbmTensor<bf16, Chip, m![A]>,
    rhs: &HbmTensor<bf16, Chip, m![A]>,
) -> HbmTensor<bf16, Chip, m![1]> {
    // HBM → DM
    let lhs: DmTensor<bf16, Chip, Cluster, Slice, m![A]> = lhs.to_dm(&mut device.tdma);
    let rhs: DmTensor<bf16, Chip, Cluster, Slice, m![A]> = rhs.to_dm(&mut device.tdma);

    // Sub context: load rhs into TRF
    let rhs: TrfTensor<bf16, Chip, Cluster, Slice, Lane, m![A]> = device
        .sub
        .begin(rhs.view())
        .fetch::<Time, m![A]>()
        .collect::<m![{ Time }, A / 16], m![A % 16]>()
        .to_trf();

    // Main context: stream lhs through the Contraction Engine, reduce along A
    let result: DmTensor<bf16, Chip, Cluster, Slice, m![1 # 8]> = device
        .main
        .begin(lhs.view())
        .fetch::<Time, m![A]>()
        .collect::<m![A / 16], m![A % 16]>()
        // Pair consecutive 32-byte flits into 64-byte packets, halving time steps (A/16 → A/32)
        .contract_outer::<m![A / 32], m![A % 32], _, _, _>(&rhs)
        .contract_packet::<m![1]>()
        .contract_time::<m![1]>()
        .contract_lane::<m![1], m![1 # 8]>(LaneMode::Interleaved)
        .cast::<bf16, m![1 # 16]>() // cast f32 accumulator output back to bf16
        .commit_trim::<m![1 # 8]>()
        .commit();

    // DM → HBM
    result.to_hbm(&mut device.tdma)
}

호스트 프로그램과 오라클

기본 템플릿 호스트 프로그램과 오라클: src/dot_product.rs.

use furiosa_opt_std::prelude::*;
use {{ crate_name }}::kernel::dot_product_kernel::{A, dot_product_kernel};
use rand::SeedableRng;
use rand::rngs::SmallRng;

#[tokio::main]
async fn main() -> Result<(), Error> {
    let mut device = Device::new(dot_product_kernel.topology())?;
    let mut rng = SmallRng::seed_from_u64(42);
    let lhs = HostTensor::<bf16, m![A]>::rand(&mut rng);
    let rhs = HostTensor::<bf16, m![A]>::rand(&mut rng);
    let lhs_hbm = lhs.to_hbm(&mut device.pdma).await?;
    let rhs_hbm = rhs.to_hbm(&mut device.pdma).await?;
    let _out_hbm = device.launch(dot_product_kernel, (&lhs_hbm, &rhs_hbm)).await?;
    println!("Dot Product: kernel ran");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn matches_reference() {
        let mut device = Device::new(dot_product_kernel.topology()).unwrap();

        let mut rng = SmallRng::seed_from_u64(42);
        let lhs = HostTensor::<bf16, m![A]>::rand(&mut rng);
        let rhs = HostTensor::<bf16, m![A]>::rand(&mut rng);

        let lhs_hbm = lhs.to_hbm(&mut device.pdma).await.unwrap();
        let rhs_hbm = rhs.to_hbm(&mut device.pdma).await.unwrap();

        // Reference: sum_i lhs[i] * rhs[i] in f32, then round to bf16.
        let lhs_buf: Vec<bf16> = lhs.into_vec();
        let rhs_buf: Vec<bf16> = rhs.into_vec();
        let expected_f32: f32 = lhs_buf
            .iter()
            .zip(&rhs_buf)
            .map(|(&a, &b)| f32::from(a) * f32::from(b))
            .sum();
        let expected = bf16::from_f32(expected_f32);

        let out_hbm = device.launch(dot_product_kernel, (&lhs_hbm, &rhs_hbm)).await.unwrap();

        let actual_buf: Vec<bf16> = out_hbm.to_host::<m![1]>(&mut device.pdma).await.unwrap().into_vec();
        if let Some(&actual) = actual_buf.first() {
            let diff = (f32::from(actual) - f32::from(expected)).abs();
            let tol = (0.02 * f32::from(expected).abs()).max(0.5);
            assert!(
                diff <= tol,
                "dot_product mismatch: expected {expected:?}, actual {actual:?}, diff {diff} > tol {tol}"
            );
        }
    }
}

GEMV

목표

\(y_i = \sum_j A_{ij} x_j\)를 계산합니다. 이 패턴은 LLM 디코드에 해당합니다.

매핑과 데이터 이동

출력 차원 I를 여러 Slice에 걸쳐 매핑하여 각 Slice가 한 행을 계산하도록 하십시오. 전체 벡터 x를 브로드캐스트하여 모든 Slice가 이 벡터를 자신의 행과 축약할 수 있도록 하십시오. 내적과 달리 각 행은 서로 독립적이지만, 모든 행에는 동일한 완전한 벡터 피연산자가 필요합니다.

디바이스 소스

디바이스 소스: src/kernel/gemv_kernel.rs.

use furiosa_opt_std::prelude::*;

axes![I = 256, J = 2048];

pub type Chip = m![1];
pub type Cluster = m![1 # 2];
pub type Slice = m![I]; // Distribute output dimension across slices
pub type Time = m![J / 32]; // Temporal iterations for reduction dimension
pub type Packet = m![J % 32]; // Packet size for reduction dimension
pub type Lane = m![1];

#[device(chip = 1)]
pub fn gemv_kernel(
    device: &mut Device,
    matrix: &HbmTensor<bf16, Chip, m![I, J]>,
    vector: &HbmTensor<bf16, Chip, m![J]>,
) -> HbmTensor<bf16, Chip, m![I]> {
    // Move data from HBM to DM
    let matrix: DmTensor<bf16, Chip, Cluster, Slice, m![J]> = matrix.to_dm(&mut device.tdma);
    let vector: DmTensor<bf16, Chip, Cluster, Slice, m![J]> = vector.to_dm(&mut device.tdma);

    // Load vector into TRF
    let vector_trf: TrfTensor<bf16, Chip, Cluster, Slice, Lane, m![J]> = device
        .sub
        .begin(vector.view())
        .fetch::<m![1], m![J]>()
        // Collect Engine: split into 32-byte flits.
        .collect::<m![J / 16], m![J % 16]>()
        .to_trf();

    // Compute GEMV: matrix × vector
    // Key difference: `I` maps to slice (preserved), `J` gets reduced
    let result: DmTensor<bf16, Chip, Cluster, Slice, m![1 # 4]> = device
        .main
        .begin(matrix.view())
        .fetch::<m![J / 16], m![J % 16]>()
        .collect::<m![J / 16], m![J % 16]>()
        .contract_outer::<Time, Packet, _, _, _>(&vector_trf)
        .contract_packet::<m![1]>()
        .contract_time::<m![1]>()
        .contract_lane::<m![1], m![1 # 8]>(LaneMode::Interleaved)
        .cast::<bf16, m![1 # 16]>()
        .commit_trim::<m![1 # 4]>()
        .commit();

    // Transfer result to HBM
    result.to_hbm(&mut device.tdma)
}

호스트 프로그램과 오라클

기본 템플릿 호스트 프로그램과 오라클: src/gemv.rs.

use furiosa_opt_std::prelude::*;
use {{ crate_name }}::kernel::gemv_kernel::{I, J, gemv_kernel};
use rand::SeedableRng;
use rand::rngs::SmallRng;

#[tokio::main]
async fn main() -> Result<(), Error> {
    let mut device = Device::new(gemv_kernel.topology())?;
    let mut rng = SmallRng::seed_from_u64(42);
    let matrix = HostTensor::<bf16, m![I, J]>::rand(&mut rng);
    let vector = HostTensor::<bf16, m![J]>::rand(&mut rng);
    let matrix_hbm = matrix.to_hbm(&mut device.pdma).await?;
    let vector_hbm = vector.to_hbm(&mut device.pdma).await?;
    let _out_hbm = launch(gemv_kernel, (&mut device, &matrix_hbm, &vector_hbm)).await?;
    println!("GEMV: kernel ran");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn matches_reference() {
        let mut device = Device::new(gemv_kernel.topology()).unwrap();

        let mut rng = SmallRng::seed_from_u64(42);
        let matrix = HostTensor::<bf16, m![I, J]>::rand(&mut rng);
        let vector = HostTensor::<bf16, m![J]>::rand(&mut rng);

        let matrix_hbm = matrix.to_hbm(&mut device.pdma).await.unwrap();
        let vector_hbm = vector.to_hbm(&mut device.pdma).await.unwrap();

        // Reference: y[i] = sum_j matrix[i, j] * vector[j] in f32, rounded to bf16.
        let mat_buf: Vec<bf16> = matrix.into_vec();
        let vec_buf: Vec<bf16> = vector.into_vec();
        let expected: Vec<bf16> = mat_buf
            .chunks(J::SIZE)
            .map(|row| {
                let acc: f32 = row
                    .iter()
                    .zip(&vec_buf)
                    .map(|(&a, &b)| f32::from(a) * f32::from(b))
                    .sum();
                bf16::from_f32(acc)
            })
            .collect();

        let out_hbm = launch(gemv_kernel, (&mut device, &matrix_hbm, &vector_hbm)).await.unwrap();

        let actual: Vec<bf16> = out_hbm.to_host::<m![I]>(&mut device.pdma).await.unwrap().into_vec();
        for (i, (&e, &a)) in expected.iter().zip(&actual).enumerate() {
            let diff = (f32::from(a) - f32::from(e)).abs();
            let tol = (0.02 * f32::from(e).abs()).max(0.5);
            assert!(
                diff <= tol,
                "gemv mismatch at i={i}: expected {e:?}, actual {a:?}, diff {diff} > tol {tol}"
            );
        }
    }
}

GEMM

목표

\(C_{ij} = \sum_k A_{ik} B_{kj}\)를 계산합니다. A는 J에 걸쳐 브로드캐스트되고, B는 I에 걸쳐 브로드캐스트됩니다.

매핑과 데이터 이동

type Slice = m![I / 32, J / 32]로 두 출력 차원을 모두 매핑하여 각 Slice가 16 × 16 출력 타일을 계산하도록 하십시오. Switch Engine이 각 B 타일을 대응하는 Slice로 라우팅하므로, 각 Slice는 자신의 출력 타일에 속하는 J 부분만 받습니다. .contract_packet::<m![1]>()은 K를 따라 공간적으로 리덕션합니다. .contract_time::<m![I]>()은 I를 유지하면서 시간에 걸쳐 누적합니다. .contract_lane::<m![I], m![J # 8]>(LaneMode::Interleaved)는 출력 패킷에서 I와 J를 유지합니다.

디바이스 소스

디바이스 소스: src/kernel/gemm_kernel.rs.

use furiosa_opt_std::prelude::*;

axes![I = 512, J = 512, K = 64];

pub type Chip = m![1];
pub type Cluster = m![1 # 2];
// Distribute output dimensions `I` and `J` across slices
pub type Slice = m![I / 32, J / 32]; // Each slice handles a 16 × 16 output tile
pub type Lane = m![J % 8];

#[device(chip = 1)]
pub fn gemm_kernel(
    device: &mut Device,
    a: &HbmTensor<bf16, Chip, m![I, K]>,
    b: &HbmTensor<bf16, Chip, m![J, K]>,
) -> HbmTensor<bf16, Chip, m![I, J]> {
    // Move data from HBM to DM
    let a: DmTensor<bf16, Chip, Cluster, Slice, m![I % 32, K]> = a.to_dm(&mut device.tdma);
    let b: DmTensor<bf16, Chip, Cluster, Slice, m![J % 32, K]> = b.to_dm(&mut device.tdma);

    // Load matrix B into TRF
    // Switch Engine distributes B across 256 slices
    // Each slice gets the full `K` dimension but only its (16 × 16) output tile
    // See: Switch Engine topologies for details on distribution
    let b_trf: TrfTensor<bf16, Chip, Cluster, Slice, Lane, m![J / 8 % 4, K]> = device
        .sub
        .begin(b.view())
        .fetch::<m![J % 8, J / 8 % 4], m![K]>()
        .collect::<m![J % 8, J / 8 % 4, K / 16], m![K % 16]>()
        .to_trf();

    // Compute GEMM: A × B
    // Switch Engine ensures matching (`I / 32`, `J / 32`) slice distribution
    // Contraction reduces along `K`, preserves `I` and `J`
    let result: DmTensor<bf16, Chip, Cluster, Slice, m![I % 32, J % 32]> = device
        .main
        .begin(a.view())
        .fetch::<m![I % 32, J / 8 % 4], m![K]>()
        .collect::<m![I % 32, J / 8 % 4, K / 16], m![K % 16]>()
        .contract_outer::<m![I % 32, J / 8 % 4, K / 32], m![K % 32], _, _, _>(&b_trf)
        .contract_packet::<m![1]>()
        .contract_time::<m![I % 32, J / 8 % 4]>()
        .contract_lane::<m![I % 32, J / 8 % 4], m![J % 8]>(LaneMode::Interleaved)
        .cast::<bf16, m![J % 8 # 16]>()
        .commit_trim::<m![J % 8]>()
        .commit();

    // Transfer result to HBM
    result.to_hbm(&mut device.tdma)
}

호스트 프로그램과 오라클

기본 템플릿 호스트 프로그램과 오라클: src/gemm.rs.

use furiosa_opt_std::prelude::*;
use {{ crate_name }}::kernel::gemm_kernel::{I, J, K, gemm_kernel};
use rand::SeedableRng;
use rand::rngs::SmallRng;

#[tokio::main]
async fn main() -> Result<(), Error> {
    let mut device = Device::new(gemm_kernel.topology())?;
    let mut rng = SmallRng::seed_from_u64(42);
    let a = HostTensor::<bf16, m![I, K]>::rand(&mut rng);
    let b = HostTensor::<bf16, m![J, K]>::rand(&mut rng);
    let a_hbm = a.to_hbm(&mut device.pdma).await?;
    let b_hbm = b.to_hbm(&mut device.pdma).await?;
    let _out_hbm = launch(gemm_kernel, (&mut device, &a_hbm, &b_hbm)).await?;
    println!("GEMM: kernel ran");
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn matches_reference() {
        let mut device = Device::new(gemm_kernel.topology()).unwrap();

        let mut rng = SmallRng::seed_from_u64(42);
        let a = HostTensor::<bf16, m![I, K]>::rand(&mut rng);
        let b = HostTensor::<bf16, m![J, K]>::rand(&mut rng);

        let a_hbm = a.to_hbm(&mut device.pdma).await.unwrap();
        let b_hbm = b.to_hbm(&mut device.pdma).await.unwrap();

        // Reference: C[i, j] = sum_k A[i, k] * B[j, k] in f32, rounded to bf16.
        let a_buf: Vec<bf16> = a.into_vec();
        let b_buf: Vec<bf16> = b.into_vec();
        let expected: Vec<bf16> = a_buf
            .chunks(K::SIZE)
            .flat_map(|a_row| {
                b_buf.chunks(K::SIZE).map(move |b_row| {
                    let acc: f32 = a_row
                        .iter()
                        .zip(b_row)
                        .map(|(&a, &b)| f32::from(a) * f32::from(b))
                        .sum();
                    bf16::from_f32(acc)
                })
            })
            .collect();

        let out_hbm = launch(gemm_kernel, (&mut device, &a_hbm, &b_hbm)).await.unwrap();

        let actual: Vec<bf16> = out_hbm.to_host::<m![I, J]>(&mut device.pdma).await.unwrap().into_vec();
        for (idx, (&e, &av)) in expected.iter().zip(&actual).enumerate() {
            let diff = (f32::from(av) - f32::from(e)).abs();
            let tol = (0.05 * f32::from(e).abs()).max(1.0);
            assert!(diff <= tol, "gemm mismatch at idx={idx}: expected {e:?}, actual {av:?}");
        }
    }
}

Blocked GEMM은 온Chip DM 용량을 초과하는 행렬을 처리할 수 있도록 GEMM에 타일링을 더해 확장합니다. 이 사례는 K에 대한 시간 분할과 여러 Chip에 걸친 I 및 J의 공간 분할을 다룹니다. Flash Attention은 하나의 Transformer 어텐션 헤드에 걸쳐 GEMM, Vector Engine softmax, main/sub 프리페치를 결합합니다.