Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Outer

Outer 단계는 두 피연산자를 같은 형태로 브로드캐스트한 뒤 요소별로 곱합니다.

“Outer“라는 이름은 선형대수의 외적(outer product)에서 유래했습니다. 벡터 u(길이 n)와 v(길이 m)에 대해 u v^T는 (u v^T)[i, j] = u[i] × v[j]를 만족하는 n × m 행렬입니다. 이 행렬은 u를 열 축(길이 m)을 따라 브로드캐스트하고, v를 행 축(길이 n)을 따라 브로드캐스트한 뒤, 요소별로 곱해서 얻습니다. Outer 단계의 세 하위 단계는 바로 이 의미론을 하드웨어로 구현한 것이며, 차례로 실행됩니다.

  • Stream Adapter는 Collect Engine에서 오는 스트리밍 피연산자를 처리하고 브로드캐스트합니다.
  • TRF Sequencer는 TRF SRAM에서 오는 TRF 피연산자를 처리하고 브로드캐스트합니다.
  • Multiplier는 피연산자 타입을 넓힌 뒤(i4 / i8은 i32로, f8 / bf16은 f32로) 정렬된 두 피연산자를 요소별로 곱합니다.

출력은 결합 매핑 [Chip, Cluster, Slice, Lane, Time, Packet]을 따르는 곱셈 결과 텐서 하나이며, Packet Reducer가 리덕션 덧셈을 수행할 수 있는 상태입니다.

인터페이스

CollectTensor에서 .contract_outer(&trf)를 호출하면 Outer 단계가 실행됩니다.

impl<
    'l,
    const T: Tu,
    P: CanApplyContractOuter,
    D: Scalar + ContractionCast,
    Chip: M,
    Cluster: M,
    Slice: M,
    Time: M,
    Packet: M,
    B: Backend,
> TuTensor<'l, T, P, D, Chip, Cluster, Slice, Time, Packet, B>
{
    /// Runs the Outer stage: stashes the two un-broadcast operands (widened to the accumulator type)
    /// and the layouts [`super::ContractTimeTensor::contract_lane`] needs to fuse them into a [`LazyContraction`]. No
    /// materializing alternative, no per-backend branch -- every backend fuses the same way.
    ///
    /// The impl block binds [`ContractionCast`], so this method does not exist on a `TuTensor` whose
    /// element type the Multiplier has no MAC for.
    #[primitive(TuTensor::contract_outer)]
    pub fn contract_outer<OutTime: M, OutPacket: M, Lane: M, TrfElement: M, TrfD>(
        self,
        trf_tensor: &TrfTensor<TrfD, Chip, Cluster, Slice, Lane, TrfElement, B>,
    ) -> ContractOuterTensor<'l, T, <D as ContractionCast>::Output, D, Chip, Cluster, Slice, Lane, OutTime, OutPacket, B>
    where
        D: Cast<<D as ContractionCast>::Output>,
        // The weight (TRF) type must form a valid contraction-engine operand pair with the
        // stream type `D`: same type, or a mixed integer precision within a
        // family (i4/i5 x i4/i5, i8/i9 x i8/i9). Both operands widen to the
        // stream's accumulator for the multiply.
        TrfD: Scalar + ContractionWeight<D> + Cast<<D as ContractionCast>::Output>,
    {
        type Out<D> = <D as ContractionCast>::Output;

        // Skipping the broadcast transpose does not skip its validity contract -- a malformed
        // contraction would otherwise slip past these asserts and panic far downstream instead.
        stream_adapter::verify_stream_adapter::<D, Lane, Time, Packet, OutTime, OutPacket>();
        trf_sequencer::verify_trf_sequencer::<TrfD, Lane, TrfElement, OutTime, OutPacket>();

        // The operands keep their own compact layouts: lhs is `[Chip, Cluster, Slice, Time, Packet]`
        // (`self.inner`), rhs is `[Chip, Cluster, Slice, Lane, TrfElement]` (`trf_tensor`). A
        // bare-buffer backend reads its strides from these; `MathStorage` ignores them (its axes live
        // in the storage).
        let lhs_map = <m![{ Chip }, { Cluster }, { Slice }, { Time }, { Packet }]>::to_value();
        let rhs_map = <m![{ Chip }, { Cluster }, { Slice }, { Lane }, { TrfElement }]>::to_value();
        let pre_reduce = <m![{ Chip }, { Cluster }, { Slice }, { Lane }, { OutTime }, { OutPacket }]>::to_value();

        // Widen each operand to the `Out<D>` accumulator up front, at the operand's own (compact) size,
        // not pre_reduce's -- the fold at `contract_lane` then runs entirely in `Out<D>`. Parity-identical
        // to the old per-cell widen (same `Cast::cast` per element), it just never allocates the
        // pre_reduce-shaped broadcast this stage used to build.
        //
        let contraction = LazyContraction {
            lhs: self.inner.map(|v| -> Out<D> { v.cast() }).inner,
            rhs: trf_tensor.inner.map(|v| -> Out<D> { v.cast() }).inner,
            lhs_map,
            rhs_map,
            pre_reduce,
        };
        ContractOuterTensor::new(self.device, contraction)
    }
}

두 경로는 각자의 어댑터(스트리밍 경로는 Stream Adapter, TRF 경로는 TRF Sequencer)를 거친 뒤 같은 형태의 Lane / Time / Packet을 공급하고, Multiplier는 정렬된 위치끼리 요소별로 곱합니다.

스트리밍 피연산자의 Time / Packet은 출력의 OutTime / OutPacket에 매핑됩니다. OutPacket은 패킹을 통해 Time의 가장 안쪽에 있는 크기 1 또는 2의 인수를 흡수합니다. OutTime은 Time의 나머지 인수를 유지하며, 브로드캐스트를 통해 가장 안쪽 위치에 브로드캐스트 인수가 추가됩니다. 브로드캐스트 인수는 TRF 피연산자의 Lane / Element(스트리밍 피연산자가 TRF 매핑에 맞춰 복제되는 부분)에서 나오며, OutTime / OutPacket에는 나타나지만 입력과 TRF 어디에도 없는 순수 출력 축에서도 나옵니다(예: einsum AB, BC -> ABCD에서 브로드캐스트되는 D).

TrfTensor의 형태는 [Chip, Cluster, Slice, Lane, Element]이며, Chip / Cluster / Slice / Lane은 공간적으로 병렬입니다. Chip / Cluster / Slice는 그대로 출력으로 전달되고, Lane은 레인별 데이터를 1–8개의 하드웨어 레인에 나눕니다. Element(.to_trf()가 정한 레인별 레이아웃)는 TRF Sequencer가 OutTime / OutPacket을 채우도록 형태를 바꿉니다.

Stream Adapter

Stream Adapter는 패킹과 브로드캐스트라는 두 연산으로 스트리밍 Time / Packet을 계산 형태(Lane / OutTime / OutPacket)로 변환합니다. 컴파일러는 사용자가 지정한 OutTime / OutPacket, TRF 피연산자, 순수 출력 브로드캐스트 축으로부터 세 자유 변수(PackSize, LaneBroadcast, TimeBroadcast)를 도출하며, 그 결과 매핑은 다음과 같습니다.

Lane      = LaneBroadcast
OutTime   = [Time / PackSize, TimeBroadcast]
OutPacket = [Time % PackSize, Packet] # (64 / D::SIZE)

패킹

Collect Engine은 32 B 플릿을 생성하고, Outer 단계는 PackSize × 32 B(RNGD에서는 32 또는 64) 크기의 패킷을 내보냅니다. 패킹은 연속된 플릿 PackSize ∈ {1, 2}개를 패킷 하나로 합칩니다.

PackTime   = [Time / PackSize]
PackPacket = [Time % PackSize, Packet] # (PackSize × 32 / D::SIZE)

PackSize는 OutPacket을 입력 Packet과 대조해 정해집니다. OutPacket이 Time의 가장 안쪽에 있는 크기 2의 인수를 흡수하면 PackSize = 2이고, 그렇지 않으면 PackSize = 1입니다. 동등하게 PackSize = OutPacket::SIZE * D::SIZE / 32이므로, 사용자가 OutPacket(32 B 또는 64 B)을 선택하면 패킹의 Collect 플릿 수가 그에 따라 결정됩니다.

하드웨어는 내부적으로 항상 64 B 패킷으로 동작합니다. PackSize = 1이면 사용하지 않는 32 B 절반은 0으로 채워지며, 이 값은 논리적인 OutPacket 타입으로 전파되지 않습니다. 따라서 후속 단계(Packet Reducer, Lane Folder)는 PackSize × 32 B 페이로드만 보게 되어 더미 사이클을 피할 수 있습니다. 자세한 내용은 Lane Folder Sequential 참고 사항을 참조하십시오.

브로드캐스트

패킹이 끝나면 Stream Adapter는 LaneBroadcast(TRF의 Lane 매핑, ∈ {1, 2, 4, 8})로 데이터를 공간적으로 브로드캐스트하고, TimeBroadcast로 시간적으로 브로드캐스트합니다. TimeBroadcast는 입력 Time에 없는 TRF Element의 인수를 다루며, OutTime에 있지만 입력과 TRF 어디에도 없는 순수 출력 축도 함께 다룹니다. 두 경우 모두 같은 브로드캐스트 메커니즘이 패킷을 복제합니다. 각 목적지는 같은 OutPacket을 받습니다.

Lane      = LaneBroadcast
OutTime   = [PackTime, TimeBroadcast]
OutPacket = PackPacket

TimeBroadcast 인수는 OutTime의 가장 안쪽 위치를 차지합니다. 따라서 바깥쪽 PackTime 인수를 순회하기 전에 같은 OutPacket이 이 인수들에 걸쳐 반복 전송됩니다.

예제

아래 예제는 두 연산을 모두 사용합니다. 패킹은 Time의 가장 안쪽에 있는 크기 2의 인수 L을 Packet에 흡수하고(PackSize = 2), Lane 브로드캐스트는 그렇게 만들어진 패킷을 N = 8개 레인에 분배하며, Time 브로드캐스트는 TRF에만 있는 B = 5 축을 따라 스트리밍 데이터를 타일링합니다.

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![M = 32, N = 8, K = 16, L = 2, B = 5];

fn stream_adapter_example<'l, const T: Tu>(
    input: CollectTensor<'l, { T }, bf16, m![1], m![1 # 2], m![1 # 256], m![M, L], m![K]>,
    trf: &TrfTensor<bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![B, L, K]>,
) -> ContractOuterTensor<'l, { T }, f32, bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![M, B], m![L, K]> {
    // Packing (PackSize = 2):
    //   L = 2 (innermost Time) absorbed into Packet.
    //   PackTime = [M = 32], PackPacket = [L = 2, K = 16] = 32 bf16 = 64B.
    // Lane Broadcast: same packet to all N = 8 lanes.
    // Time Broadcast: B = 5 (TRF-only) added at innermost OutTime.
    //   OutTime = [M, B = 5], OutPacket = [L = 2, K = 16].
    input.contract_outer::<m![M, B], m![L, K], _, _, _>(trf)
}

let mut device = Device::new(Topology { chips: 1, pes: 8 }).unwrap();

let a: CollectTensor<'_, _, bf16, m![1], m![1 # 2], m![1 # 256], m![M, L], m![K]> = CollectTensor::new(&mut device.main, Tensor::zero());
let b: TrfTensor<bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![B, L, K]> = TrfTensor::zero();
let _o = stream_adapter_example(a, &b);
}

제약 조건

  • OutPacket::SIZE * Storage::SIZE ∈ {32, 64} 바이트(RNGD 기준)이며, 여기서 Storage는 확장 전 피연산자 dtype입니다(예: bf16 = 2 B이며, 결과 텐서가 담는 확장된 f32 누산기가 아닙니다). 이 크기는 PackSize = 1일 때 32이고 PackSize = 2일 때 64입니다. 사용자가 이 크기를 선택하면 패킹의 Collect 플릿 수가 그에 따라 결정됩니다.
  • PackSize ∈ {1, 2}(패킹 참조).
  • Lane::SIZE ∈ {1, 2, 4, 8}입니다.

성능

PackSize는 MAC 활용률을 결정합니다. PackSize = 2는 64 B 전체를 채워 모든 MAC을 사용합니다. PackSize = 1은 32 B만 채우므로 0으로 패딩된 절반은 항상 0과 곱해지고 실효 처리량이 절반으로 줄어듭니다.

PackSize = 2는 Packet당 2사이클이 걸리지만(32 B 플릿 두 개가 64 B Packet 하나로 합쳐짐), 이는 파이프라인 병목이 아닙니다. 업스트림이 최대 fetch 속도로 매 사이클 32 B 플릿 하나를 공급하므로, Stream Adapter는 플릿이 도착하는 즉시 소비하고 다운스트림 소비 속도에 맞춰 2사이클마다 Packet 하나를 내보냅니다.

Time Broadcasting은 fetch 비용을 여러 사이클에 분산합니다. 브로드캐스트 인자는 다시 fetch하지 않고 같은 스트리밍 패킷을 여러 사이클에 걸쳐 재사용하므로, 해당 인자에 대한 대역폭 비용이 없어집니다.

Fetch 대역폭(fetch당 최대 32 B/사이클)이 Stream Adapter 전체의 상한을 결정합니다. 활용률을 극대화하려면 여러 Slice에 걸쳐 fetch 패턴을 인터리브하십시오.

TRF Sequencer

TRF Sequencer는 TrfTensor를 읽고 그 Element를 Packet Reducer에 맞게 OutTime / OutPacket으로 재구성합니다. TRF 저장 레이아웃(레인, 뱅크, 행, 더블 버퍼링, 캐시)은 Register Files를 참조하십시오.

매핑은 다음과 같습니다:

OutTime   = (sequencing over [Element / ReadSize] with broadcasts)
OutPacket = [PacketBroadcast, Element % ReadSize]

OutPacket은 매 사이클 TRF 읽기 한 번으로 채워집니다. 즉, 매 사이클 완전한 OutPacket 하나(두 뱅크에 걸쳐 레인당 640비트, 한 뱅크만 읽을 때는 320비트)가 생성됩니다. 8개 레인에 걸친 Slice당 총량(바이트 단위)은 Contraction Engine으로를 참조하십시오. 읽기는 Element의 가장 안쪽 연속 부분을 가져와 64 B OutPacket을 채우도록 복제합니다. 컴파일러는 Element % ReadSize == OutPacket % ReadSize와 ReadSize * D::SIZE ≤ 64 바이트를 만족하는 가장 큰 ReadSize를 선택합니다. ReadSize가 넓으면 레인당 두 TRF 뱅크에 걸치고, 좁으면 한 뱅크만 사용합니다.

OutTime은 Element / ReadSize(및 선택적인 브로드캐스트)를 순차 처리하여 여러 사이클에 걸쳐 채워집니다. TRF Sequencer는 다른 모든 Sequencer와 동일한 중첩 루프 구성을 사용합니다.

예제

이 예시에서는 ReadSize가 64 B 읽기 한 번으로 Element 전체를 포괄하므로 Element / ReadSize는 자명하며, Sequencer는 브로드캐스트만 순회합니다:

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![M = 32, N = 8, K = 32];

fn trf_sequencer_full_read<'l, const T: Tu>(
    input: CollectTensor<'l, T, bf16, m![1], m![1 # 2], m![1 # 256], m![M, K / 16], m![K % 16]>,
    trf: &TrfTensor<bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![K]>,
) -> ContractOuterTensor<'l, T, f32, bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![M], m![K]> {
    // Element         = K
    // ReadSize        = 32
    // PacketBroadcast = 1
    // OutTime         = M      (sequencing over [K / 32] (= 1), broadcast M)
    // OutPacket       = K      (= [1, K % 32])
    input.contract_outer::<m![M], m![K], _, _, _>(trf)
}

let mut device = Device::new(Topology { chips: 1, pes: 8 }).unwrap();

let a: CollectTensor<'_, _, bf16, m![1], m![1 # 2], m![1 # 256], m![M, K / 16], m![K % 16]> = CollectTensor::new(&mut device.main, Tensor::zero());
let b: TrfTensor<bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![K]> = TrfTensor::zero();
let _o = trf_sequencer_full_read(a, &b);
}

이 예시에서는 ReadSize가 Element의 일부만 포괄하므로 Element / ReadSize는 자명하지 않으며, Sequencer는 브로드캐스트와 함께 바깥쪽 Element 인자를 순회합니다:

#![allow(unused)]
fn main() {
#![feature(adt_const_params)]
extern crate furiosa_opt_std;
use furiosa_opt_std::prelude::*;
axes![M = 32, N = 8, K = 16, L = 2, O = 2];

fn trf_sequencer_partial_read<'l, const T: Tu>(
    input: CollectTensor<'l, T, bf16, m![1], m![1 # 2], m![1 # 256], m![O, M, L], m![K]>,
    trf: &TrfTensor<bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![O, K]>,
) -> ContractOuterTensor<'l, T, f32, bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![O, M], m![L, K]> {
    // Element         = [O, K]
    // ReadSize        = 16
    // PacketBroadcast = L
    // OutTime         = [O, M]    (sequencing over [O, K] / 16 (= O), broadcast M)
    // OutPacket       = [L, K]    (= [L, [O, K] % 16])
    input.contract_outer::<m![O, M], m![L, K], _, _, _>(trf)
}

let mut device = Device::new(Topology { chips: 1, pes: 8 }).unwrap();

let a: CollectTensor<'_, _, bf16, m![1], m![1 # 2], m![1 # 256], m![O, M, L], m![K]> = CollectTensor::new(&mut device.main, Tensor::zero());
let b: TrfTensor<bf16, m![1], m![1 # 2], m![1 # 256], m![N], m![O, K]> = TrfTensor::zero();
let _o = trf_sequencer_partial_read(a, &b);
}

제약 조건

  • 하드웨어 차원: Chip::SIZE, Cluster::SIZE, Slice::SIZE는 하드웨어 구성과 일치해야 합니다(Sequencer 참조).
  • 주소 정렬: Element % ReadSize가 64 B 전체를 포괄하면 읽기가 레인당 두 TRF 뱅크에 걸치므로, Sequencer의 기준 주소와 모든 스트라이드는 64 B에 정렬되어야 합니다.

아키텍처

Sequencer는 다른 모든 Sequencer와 동일한 중첩 루프 구성을 사용하여 여러 사이클에 걸쳐 Element의 바깥쪽 인자(즉, Element / ReadSize)를 순회하므로, 하나의 TrfTensor는 내용을 모두 소진하기 전까지 Element / ReadSize 사이클을 거칩니다. PacketBroadcast 인자는 한 사이클 안에서 같은 행을 복제하여, 추가 TRF 읽기 대역폭을 소비하지 않고 본래의 ReadSize를 넘어 64 B OutPacket을 채웁니다.

성능

처리량은 레인당 사이클마다 완전한 OutPacket 하나입니다. 두 뱅크를 모두 읽으면 레인당 640비트이고, 한 뱅크만 읽으면 레인당 320비트입니다. Slice별 총 바이트 수는 Register Files: To Contraction Engine을 참조하십시오. TRF 읽기 캐시와 뱅크 교대(Register Files: 더블 버퍼링 참조) 덕분에 브로드캐스트 재사용과 좁은 읽기가 일어나는 동안에도 동시 서브 컨텍스트 저장이 막히지 않습니다.

Multiplier

Multiplier는 Stream Adapter와 TRF Sequencer에서 정렬된 두 피연산자를 받아, 하류 누산기에서 오버플로가 일어나지 않도록 각 입력 요소를 축약 출력 타입으로 확장한 다음 요소별로 곱합니다. 그 출력은 결합 매핑 [Chip, Cluster, Slice, Lane, Time, Packet]의 단일 텐서이며, Packet Reducer의 입력이 됩니다. 매 Time 사이클마다 모든 Lane이 곱셈 결과로 채운 완전한 packet을 병렬로 생성합니다.

피연산자 타입

다음은 Multiplier가 허용하는 피연산자 쌍과 각 쌍이 확장되는 타입입니다.

Stream(활성화)가중치(TRF)누산기
i4 또는 i5i4 또는 i5i32
i8 또는 i9i8 또는 i9i32
f8e4m3f8e4m3f32
f8e5m2f8e5m2f32
bf16bf16f32

정수 피연산자는 원래 형식(i4/i8)일 수도 있고, 영점을 차감해 준비한 형식(i5/i9, 영점 차감에서 생성)일 수도 있으며, 같은 계열 안에서는 두 피연산자가 일치하지 않아도 됩니다. 서로 다른 계열(i4와 i8 불가)이나 종류(정수와 부동소수점 불가)를 섞을 수 없으며, 부동소수점 쌍은 타입이 정확히 같아야 합니다.

이 표에 없는 요소 타입에 대한 축약은 대신 Vector Engine에서 실행됩니다. 요소별로 곱한 다음 Intra-Slice Reduce로 리덕션하며, 이때 계산은 i32 / f32로 수행됩니다.

원문 © FuriosaAI · Programming Tensor Contraction Processors · Apache License 2.0 · NOTICE
yeokja 한국어 번역 · Anthropic 사의 claude-sonnet-5 모델을 활용하여 번역되었으며 학습을 모두 비허용한 상태로 작업하였습니다 · 이 문서는 원문을 수정한 비공식 번역본이며 Apache License 2.0으로 제공합니다.