MarlinSparseLayout

class torchao.dtypes.MarlinSparseLayout

MarlinSparseLayout 是一個用於處理稀疏張量格式的佈局類，它專門為 Marlin 稀疏核心設計。該佈局用於最佳化具有 2:4 稀疏模式的仿射量化張量的儲存和計算。

該佈局確保張量資料以與 Marlin 稀疏核心操作相容的格式進行預處理和儲存。它提供了預處理輸入張量和管理量化張量佈局的方法。

pre_process(input: Tensor) → Tensor

預處理輸入張量，使其符合 Marlin 稀疏核心的正確格式。

• 1. 輸入張量被轉置，因為線性層以轉置格式保留權重。

• 2. 張量注入 2:4 稀疏性。

• 3. 再次轉置，因為量化過程將計算 dim=-1 的尺度。

引數:

input (torch.Tensor) – 要預處理的輸入張量。

返回:

預處理後的張量。

返回型別:

torch.Tensor