CuTeDSL at Perplexity

在 Perplexity，我們使用自家製的運行時間最佳化服務引擎 (ROSE) 來提供各種模型服務，範圍從 embeddings 到萬億參數的 LLM。我們的運行時間引擎位於所有 API（包括 Sonar、Search 和 Embeddings）的後端，旨在賦能研究人員探索最適合的模型與架構，同時迅速將最佳候選模型部署到生產環境中。為了達成此目標，我們必須能夠以極快的速度實作、部署並最佳化先進架構的推論。在這篇網誌文章中，我們探討了整合至推論引擎中的 CuTeDSL 核心如何讓我們能夠建構所需的 GPU 核心，讓模型在 NVIDIA GPU 上發揮顛峰效能。

由 Perplexity 各個團隊建構的自訂模型，透過 AI 推論團隊建構的 ROSE 推論引擎，在自家製的 NVIDIA Hopper 和 Blackwell GPU 上進行託管。ROSE 最初的建構目的是為了更好地為語言解碼和分類提供自訂的 Llama 模型，並具備與 NVIDIA Triton 推論伺服器相容的介面。隨著時間推進，它發展成為一個成熟的推論引擎，為廣泛的團隊提供大型 LLM 以及排序、分類、評分和其他基於 Transformer 的模型，以驅動 Perplexity 的產品。ROSE 的設計目標很簡單：隨著時間的推移，我們投入心力以達到顛峰效能，但首要之務始終是適應我們模型不斷演進的架構與使用情境。

從宏觀來看，ROSE 由一個 engine 組成，該引擎使模型能夠適應客戶可與之互動的介面。這些引擎具有高度的可重複使用性，能夠處理請求排程與批次處理，並從底層模型擷取解碼的 token、embeddings 或分數。裝置的初始化、行程間與節點間通訊通道以及權重載入全由它們處理。對於 LLM 而言，引擎負責批次處理、分塊與取樣。它還負責針對完整注意力和線性注意力以及前綴匹配配置 KV 儲存空間。Embedding 引擎則較為簡單，因為它們僅執行線上批次處理。執行過程經過最佳化，使引擎能在模型於加速器上並行執行的同時，準備下一輪的工作。

ROSE 提供了豐富的自訂層集合用於建模，包裝了使用 CuTeDSL、Triton、CUDA、CUTLASS 和 cuBLAS 建構的 GPU 核心。對於矩陣乘法和注意力實作，我們主要依賴來自 CUTLASS 和 cuBLAS 的 NVIDIA 核心。管線的其餘部分，包括 embedding、norm、MoE 和激發核心，都是 ROSE 核心函式庫的一部分。為了以最佳方式支援我們執行推論的硬體，ROSE 配備了自己的裝置間通訊原語，用於 MoE 路由、分派和合併。雖然其他核心相對簡單，但必須經過徹底的專門化，才能在它們執行的所有不同隱藏維度和裝置上達到顛峰效能。此外，我們依賴編譯時期專門化來適應各種模型的各種參數與微調（例如 RMS norm 的 weight bias 與無 weight bias、激發函數等）。因此，我們決定使用 DSL 來實作它們：我們最初是用 Triton 建構 ROSE 的。不過，我們自此已採用 CuTeDSL 作為 GPU 編程的主要環境。

從數學角度來看，推論核心通常並不複雜。因此，在開發核心時，複雜度通常不在於表達要執行什麼運算，而在於硬體應該如何準確執行。為了建構最佳核心，我們需要一種語言，能夠針對底層硬體提供正確的抽象。合適的語言使我們能夠指定高層運算來執行運算，並在廣闊的配置空間中的所有點上進行積極的專門化與最佳化。

CuTeDSL 是一個基於 Python 的 NVIDIA GPU DSL，建構於 CuTe 佈局代數和 MLIR 之上。它能及時 (JIT) 編譯為高度最佳化的 PTX，同時保留對低階硬體原語的控制。它旨在提供與 CUTLASS C++ 相同的特性與效能，而沒有會阻礙開發速度的昂貴編譯時間。在我們調查或使用過的語言中（包括 CUDA、cuTile、TVM IR 和 Triton），我們認為它是最適合我們核心開發的選擇。

總體而言，雖然不像 Triton 或 cuTile 那樣高階且簡潔，但我們發現 CuTe 佈局代數及其組合（透過 cute.Tensor 公開）是合理的替代方案。此外，CuTeDSL 提供對硬體原語的不受限存取，為我們提供了在所有邊緣情況下達到顛峰效能所需的控制層級。使用 CUDA 時，編譯時期專門化的機會有限，因為每個額外的參數都會為編譯時間帶來不合理的成本，或為建置系統增加複雜度。相比之下，透過及時編譯 (JIT) 的 DSL，我們可以更積極地專門化，而不會犧牲迭代速度。與使用 CUDA 作為中間步驟的 TVM IR 相比，我們發現基於 MLIR 且降級至 PTX 的管線比產生 CUDA 作為中間步驟的管線更為強固且更容易除錯。此外，透過使用自訂運算來包裝 PTX 指令，我們發現我們可以充分利用底層硬體的所有功能。

CUDA 核心的大部分編譯成本通常花在核心的樣板展開階段，因為從 PTX 階段開始的最佳化和編譯可以在幾秒鐘內完成。對於具有許多編譯時期參數的非簡單核心，編譯的這個階段可能會暴增至數分鐘。這個瓶頸因平行編譯的機會只在稍後的 IR 和 PTX 層級出現（而 C++ 樣板處理則是循序發生）而更加突顯。雖然存在透過將樣板實例化移動到獨立的編譯單元來解決此問題的技巧，但它們需要對建置系統進行重大調整並使用笨拙的巨集。相比之下，及時編譯 (JIT) 的 CuTeDSL 核心不需要複雜的建置規則，也不會阻礙編譯時間，因為開發週期可以集中在配置空間中的單一點上，而隨後的仔細單元測試可以涵蓋核心的整個設計空間。

儘管 JIT 編譯將成本從開發時期轉移到執行時期，但額外負荷是完全可以接受的。當部署模型時，它只會啟動整個配置空間中的少數幾個點，在推論引擎的啟動/預熱期間產生秒級的額外負荷。相比之下，預先編譯核心的所有配置可能會跨越數分鐘。此外，透過快取 JIT 編譯的成品，可以進一步減輕額外負荷。

當使用基於樣板的分派和專門化（如 CuTe 和 CUTLASS 的情況）時，追蹤錯誤根源的困難度更加突顯了 CUDA 沈重的編譯成本。有了 DSL，錯誤可以直接指向產生它的程式碼行，從而大大簡化除錯。此外，透過中間 MLIR 階段或依賴附加到產生的 PTX 的除錯資訊，將產生的程式碼對應到產生它的 DSL 陳述式要簡單得多。多虧了 DSL 的這些特性，我們發現自己在開發核心時的生產力高得多。

雖然核心在各個模型之間有許多相似之處，但它們仍然必須配備一系列可能對效能產生影響的配置旋鈕。相關參數包括：

隱藏維度/表頭維度/注意力表頭數量

輸入/輸出資料類型

量化方案

編譯時期常數（偏置、激發函數等）

編譯時期專門化對效能和除錯都很有用。特別是在管線化處理器中（如 GPU 的情況），不存在的分支比永遠會走或永遠不會走的分支更好。此外，針對特定配置旋鈕集合編譯核心會產生較少的 PTX/SASS 指令，在調查期間使用 cuda-gdb 時更容易分析和除錯。

對於給定的批次大小和配置，可以估計對 DRAM 的預期流量和算術運算次數。隨後，可以根據每個案例做出最佳決策，以進行最佳網格配置、同步化以及特定硬體功能的利用。CuTeDSL 允許我們定義和重複使用用於深度專門化的元件，而不會產生無法接受的編譯時期額外負荷。

一般而言，推論核心的效能必須針對兩個獨立的使用情境進行最佳化：

預填充 (Prefill)：小批次大小（少數分塊、序列），每個序列有許多 token。必須為 embeddings 和 LLM 的輸入 token 最大化吞吐量。延遲在數十到數百毫秒之間。

解碼 (Decode)：小到大批次大小，每個序列一個（無推測解碼）到幾個 token（草稿-目標推測、多 token 預測）。延遲接近核心啟動延遲。

對於 GEMM 和注意力，預填充階段在大量的 token 上執行，效能受限於裝置的運算吞吐量。解碼涉及對小批次大小進行相同的運算，執行時間主要由將啟動和權重從 DRAM 載入到串流多處理器和張量核心的成本所主導，反而受限於裝置的記憶體頻寬。儘管批次大小或批次中的 token 總數是一個相對清晰的指標，但預填充和解碼之間的切換閾值更為微妙，因為記憶體頻寬也由模型的隱藏維度和啟動的資料類型決定。因此，如果演算法差異或網格/區塊指派會對其中任何一個有利，核心就會被分割成預填充和解碼實作。

在預填充和解碼核心之間普遍不同的參數是區塊中的 warp 數量（執行緒數除以每個 warp 32 個執行緒的 warp 大小）。將此參數設為編譯時期常數是理想的做法，因為它有助於暫存器配置。在 CUDA 中，這通常是透過透過 __launch_bounds__ 指定執行緒數量並提供 warp 數量作為樣板參數來完成：

cpp

template <size_t NUM_WARPS>
__global__ void 
__launch_bounds__(NUM_WARPS * WARP_SIZE, 1)
some_kernel(...) {

}

區塊中所有執行緒可用的暫存器數量是一個硬體限制，在 Hopper 和 Blackwell 上為 64K。如果事先知道執行緒數量，編譯器就能準確知道它可以為每個執行緒指派多少暫存器，並且能夠成功編譯出更積極最佳化的核心。然而，這會付出編譯時期的成本，因為現在必須針對不同的 warp 大小重新編譯核心。有了 CuTeDSL，這不再是個問題，因為我們不需要預先支付整個配置空間的編譯成本。該 DSL 還期望區塊大小是一個常數運算式，隱式執行 CUDA 中由 __latch_bounds__ 暗示的最佳化。

除了使 warp 數量成為常數之外，理想的做法是在預填充和解碼之間變動它。為了充分利用 GPU 的記憶體頻寬，始終為每個區塊指派大量的 warp 是不夠的，因為只有透過使用盡可能多的可用串流多處理器 (SM) 才能達到完整的頻寬。因此，對於給定的 token 數量以及已知的資料類型和隱藏維度，我們先挑選理想的每執行緒向量維度（bfloat16 為 8，float32 為 4），以確保產生的程式碼使用 128 位元載入/儲存指令（在 Blackwell 上為 256 位元）。知道了向量大小後，我們就可以計算核心所需的 warp 數量。我們首先嘗試最大化 SM 的數量：如果 warp 數量少於 SM，則選擇每個區塊單一 warp。這通常是核心的解碼變體。如果有足夠的 warp，我們會逐漸將 warp 數量增加到 8、16 或 32，並進行效能評測以找出上限，在此上限之後通常不再看到效能提升。

雖然 Triton 在底層執行了一些此類最佳化，但我們發現透過仔細使用 CuTe 佈局代數，我們可以在不受高階編譯器缺點限制的情況下達到類似的結果。特別是當涉及子 warp 水平歸納時，我們發現 CuTeDSL 為我們提供了產生最佳程式碼所需的細粒度控制。

在某些情況下，為了將工作分發到盡可能多的區塊，可以使用 cooperative_groups::this_grid().sync() 來同步化所有參與的區塊。但有以下限制：由於區塊一旦開始執行就無法從 SM 取消排程，為了保證向前進展，區塊數量不能超過裝置上可用的 SM 數量，以確保在任何時間點，所有區塊都可以更新和輪詢參與同步化的全域記憶體位置。此限制由 cudaLaunchCooperativeKernel 強制執行。網格屏障的成本在 Hopper 上大約是 3us，再加上在某些情況下透過隱式記憶體屏障清除所有先前記憶體操作的成本。總體而言，這相當於將核心分割為兩個啟動的額外負荷，從而導致效能敏感的設計決策。儘管程式化相依啟動 (PDL) 可以從 Hopper 開始降低啟動額外負荷，但對於解碼核心來說，避免這種額外負荷是值得的。

如果其中一個工作負載需要的區塊數量大於 SM 數量，則應分割核心，並依賴核心啟動的隱式屏障。這主要是針對預填充的情況。如果對所有工作負載而言最理想的區塊數量適合放入網格中，則偏好使用網格屏障，例如操作較少 token 的解碼核心的情況。

在 CuTeDSL 中，這些決策可以在主機端啟動程式碼中做出，並且可以根據需要組合和交錯核心與網格屏障，如下方的範例所示。透過這種混合策略，可以將解碼核心最佳化以達到最小延遲，並將預填充核心最佳化以達到最大吞吐量。與基於 CUDA 的解決方案不同，這種邏輯可以更簡潔地定義，且編譯時期額外負荷更少。

python

@cute.kernel
def dispatch_kernel_decode(...):
	route_tokens(...)
	grid_barrier(grid_barrier_ptr)
	copy_tokens_to_experts(...)

@cute.kernel
def dispatch_kernel_prefill_route(...):
	route_tokens(...)
	
@cute.kernel
def dispatch_kernel_prefill_copy(...):
	copy_tokens_to_experts(...)
	
@cute.jit
def dispatch(...):
	num_route_blocks = ...
	num_copy_blocks = ...
	
	if num_copy_blocks < NUM_SMS:
		dispatch_kernel_decode(...).launch(
			grid=(min(max(num_route_blocks, num_copy_blocks), NUM_SMS), 1, 1)
		)
	else:
		dispach_kernel_prefill_route(...).launch(
			grid=(num_route_blocks, 1, 1)
		)
		dispatch_kernel_prefill_copy(...).launch(
			grid=(num_copy_blocks, 1, 1)
		)

就我們的 MoE 全對全 (all-to-all) 核心而言，如下方詳細說明，具有帶或不帶網格屏障建構核心並在編譯時進行最佳化的靈活性，透過 CuTeDSL 協助我們實現了顯著的效能提升。

我們在 CuTeDSL 中建構的第一個核心是專門化的 RMS norm 核心，針對小維度進行了最佳化。QK norm 用於近期 Qwen 模型的自我注意力層中，針對 64、128 或 256 個元素以每個表頭為基礎進行套用。與在超過數千個隱藏維度上運作、需要透過共用記憶體和暫存器進行水平歸納的常規 RMS norm 不同，QK norm 僅需跨 warp 進行歸納，並且有可能將多個表頭擠壓到單一 warp 中。這成為一個使用情境，我們可以在編譯時進行積極最佳化，以為所有表頭維度和資料類型找到最佳的 token 到 warp 分配。

RMS norm 是一個相當直接的操作，儘管它需要水平總和來計算向量的平方範數：

cpp

r = w * (x * ((x * x).sum(dim=-1) / x.shape[-1] + EPS))

由於注意力表頭通常很小且每執行緒運算高度向量化，因此單一 warp 可以計算平方範數。為了最大化記憶體頻寬，每個執行緒應處理 128 位元的輸入，這相當於 8 個 bfloat16 值。因此，半個 warp 足以處理一個表頭，並且可以使用由 shuffle_sync_bfly 執行的水平歸納來決定平方範數。從表頭維度和資料類型開始，我們可以決定每個 warp 的表頭數量，並在編譯時指定引數和歸納輪數。

然後，我們可以藉助自動向量化的載入/儲存，並使用用於 warp 歸納的低階架構特定運算，以相對較少的程式碼行建構專門化的核心。核心決定要處理哪個表頭（Q 或 K）並識別要套用的對應權重。它從表頭載入整個向量並首先在本地加總。隨後，根據每個表頭的執行緒數量，計算平方範數。最後，決定自動向量化的 RMS norm 並寫回記憶體。特定於模型的參數（例如 WEIGHT_BIAS 和 EPS）會被常數摺疊。

python

@cute.kernel
def _qk_norm_kernel(
    q: cute.Tensor,
    k: cute.Tensor,
    q_weight: cute.Tensor,
    k_weight: cute.Tensor,
    WEIGHT_BIAS: cutlass.Constexpr[float],
    EPS: cutlass.Constexpr[float],
) -> None:
    # Find the block and thread indices.
    bidx, _, _ = cute.arch.block_idx()
    bdim, _, _ = cute.arch.block_dim()
    tidx, _, _ = cute.arch.thread_idx()

    _, num_q_heads, (threads_per_head, vecs_per_thread, vec_size) = q.shape
    _, num_k_heads, _ = k.shape
    head_dim = threads_per_head * vecs_per_thread * vec_size

    # Determine the thread and the vector to work on.
    vec_id = (bidx * bdim + tidx) // threads_per_head
    thread_id = tidx % threads_per_head

    # Determine the token and head indices.
    token = vec_id // (num_q_heads + num_k_heads)
    head_offset = vec_id % (num_q_heads + num_k_heads)
    is_q = head_offset < num_q_heads

    # Determine which tensor to work with.
    thread_coord = (thread_id, None, None)
    head_coord = (token, head_offset if is_q else head_offset - num_q_heads, thread_coord)

    # Load the head into registers.
    x_ref = q[head_coord] if is_q else k[head_coord]
    w = (q_weight if is_q else k_weight)[thread_coord].load()

    # Compute the squared sum locally.
    x = x_ref.load()
    x_f32 = x.to(cute.Float32)
    x_square = x_f32 * x_f32
    x_square_sum = x_square.reduce(cute.ReductionOp.ADD, init_val=0.0)

    # Sum up across the warp.
    mask = activemask()
    start = cute.arch.WARP_SIZE // (cute.arch.WARP_SIZE // threads_per_head * 2)
    for offset in [start >> i for i in range(5) if (start >> i) > 0]:
        x_square_sum += cute.arch.shuffle_sync_bfly(
            x_square_sum,
            offset=offset,
            mask=mask,
        )

    # RMS norm.
    rms = cute.math.rsqrt(x_square_sum / head_dim + EPS)
    x_norm = ((w + WEIGHT_BIAS) * (x * rms)).to(x.element_type)
    x_ref.store(x_norm)

與常規的 RMS norm 核心相比，此專門化變體為預填充和解碼帶來了 2-3 倍更好的吞吐量和延遲。

MoE 分派和合併將輸入 token 洗牌到它們被路由到的專家，在 MoE GEMM 核心執行相關的矩陣乘法之後執行加權累積。效能特別令人感興趣，因為各種產品都依賴各種大小的 MoE 模型。根據部署類型，我們區分要最佳化的不同核心：

TP=1：用於單一 GPU 的本地核心。分派將 token 洗牌到相關專家，合併則將其累積。主要用於 H100 上的 Qwen3 30B 部署。

DP=1, TP ≠ 1：張量並行、單節點核心。輸入被複製，分派選擇目前排名的 token，合併則跨所有排名進行累積。跨 NVLink 或 NVLS 進行累積。用於大型預填充器。

DP ≠ 1, NVLink：具備選用張量並行的資料並行核心。每個排名將 token 傳送到代管特定專家的每個其他排名。路由資訊被聚合，token 透過 NVIDIA NVLink 傳輸。用於資料並行的 H200 或 GB200 部署。

DP ≠ 1, InfiniBand：跨 NVLink 和 NVIDIA Quantum InfiniBand 運作的點對點資料並行和張量並行核心。用於多節點 Hopper 部署。

本地核心根據路由表計算指派給每個專家的 token 數量，並在屬於每個專家的 token 清單中為每個 token 指派一個唯一的索引。路由步驟之後，token 會被複製到選用填補的緩衝區中，以傳遞給 MoE GEMM 核心。然後，合併核心使用 token 位置在暫存器中執行累積，將最終的 token 寫入輸出緩衝區。

分派階段在預填充和解碼之間有所不同：對於少數 token，要讀出的路由資訊很少，因此我們在單一 block 中處理它，並向網格的其餘部分發出訊號，由個別 block 複製 token 區塊。token 分派是透過以原子方式遞增共用記憶體中的計數器來決定的。由於預填充必須處理更多路由資訊，因此路由在獨立核心中跨多個 block 進行分割。首先使用共用記憶體中的計數器在本地指派索引，然後跨全域記憶體中的所有 block 進行聚合，以決定每個 token 的位置。此核心的網格被夾取到 SM 的數量，以便能夠使用網格級別的屏障。隨後的核心執行複製，跨越大量 block 和 warp 啟動，以充分利用可用的記憶體頻寬。

合併核心相對簡單，因為它執行向量化累積，依賴路由資訊和分派所計算的位移。CuTeDSL 在這裡證明了其價值，因為運算可以透過簡單的 CuTe 張量來表達，這些張量會針對所有資料類型和隱藏維度編譯為高效的程式碼。

在張量並行部署中，路由資訊和輸入 token 會被複製。透過稍作更改以僅路由和選擇指派給位於目前排名的專家的 token，本地分派核心可以完全重複使用。藉由一些參數和編譯時期專門化，本地分派核心得以重複使用。

由於排名數量通常少於 token 被路由到的專家數量，在平均情況下，可以假設每個 token 至少有一個複本被路由到每個排名。在這種情況下，可以重複使用本地合併核心來聚合 token 的本地複本，接著進行 all-reduce 以跨參與的排名進行加總。若是預填充核心，我們會啟動獨立的核心，然而解碼會將這兩個操作融合到單一啟動中。

對於 all-reduce，如果只有兩個排名參與，我們使用單純的 NVLink；當有 4 個或更多排名參與時，則切換到 NVLS，以使吞吐量增加 10-20%。多虧了 DSL，我們可以在 MoE 和常規 all-reduce 核心之間共用大部分的歸納邏輯和對等緩衝區交換。這是一個對硬體進行緊密控制很有幫助的情況，因為 NVLS 需要透過 multimem.ld_reduce 和 multimem.st 進行 DSL 運算。歸納本身的實作主要遵循 NCCL 測試中的 all_reduce 範例。

資料並行核心依賴 NVLink 來為多個裝置聚合 token。在分派階段，在交換 token 之前，對等體首先在本地進行路由，以識別每個發送的 token 在每個排名上的目標位移。隨後，它們開始將 token 複製到為每個發送者配置的接收緩衝區中，以避免衝突。每個 token 的 release-store 寫入用於發出完成信號。在接收端，路由資訊被聚合，以便將所有 token 壓縮到單一緩衝區中以進行 MoE GEMM。接收者透過 acquire-load 輪詢對等體寫入的旗標，以等待寫入完成。

在合併端，token 透過 NVLink 複製到來源排名，並附帶指示收到的每個 token 旗標。合併核心類似於本地和單節點變體，但在定址和一些額外的同步化邏輯上有些許差異。

我們建構了一個 PTX 包裝器函式庫，以公開系統和裝置層級範圍的 release-acquire 載入和儲存來實作此核心。透過改善同步化旗標的粒度並配置更多 block 來透過 NVLink 發送和接收 token，我們比先前的 CUDA 核心實現了 2 倍的效能提升。

我們的點對點核心先前在搭配 ConnectX-7 InfiniBand 轉接卡的 Hopper GPU 上實現了最先進的效能。我們自此將合併核心移植到 CuTeDSL，以進一步提升效能。由於 CUDA 版本已經嚴重依賴編譯時期專門化，它們證明了是我們儲存庫中的編譯時期瓶頸，迫使我們省略了一些參數。透過遷移到 JIT 編譯的 DSL，我們針對更多參數進行了專門化並調整了網格配置，在預填充和解碼之間引入了進一步的差異，從而獲得了大約 3% 的延遲提升。

從 ROSE 核心的角度來看，Blackwell 支援分為兩類：與架構無關和與架構相關的核心。依賴張量核心的架構特定核心需要明確的工作來進行適應。然而，由於使用了 DSL，與架構無關的核心可以輕鬆適應。

Blackwell GPU 擁有最多 160 個 SM，而 Hopper 則為 132 個。對於大多數核心而言，這是一個編譯時期常數，主機端啟動器程式碼中的邏輯將其納入考量，以決定要執行的最佳（預填充對解碼）核心。此外，Blackwell 還引入了 256 位元載入指令，可以透過增加向量寬度（CuTe 張量的最後一個維度）來輕鬆發出，以改善記憶體頻寬利用率。

Blackwell 引入了對大量量子化方案（MXFP8、MXFP4、NVFP4）的硬體支援。對這些方案的支援主要是透過調整現有的 FP8 區塊級縮放核心來實作的。JIT 支援在此至關重要，因為支援新模式時不會使編譯時期配置空間暴增。過去，我們發現 Triton 在表達跨越少數元素（超過一半或四分之一的 warp）的水平歸納時顯得笨拙。對於 NVFP4 和 MXFP8/MXFP4 分別計算 16 或 32 個元素的絕對最大值時，對歸納的更細粒度控制對於達到最佳效能至關重要。

作為一個相對較新的語言，LLM 在其預訓練期間接觸 CuTeDSL 程式碼的機會很少。此外，我們目前的程式碼風格（其核心針對啟動和 norm 進行了最佳化）與現有的 CUTLASS 或 Quack 範例略有不同，後者針對 TMA 和架構特定的矩陣乘法指令進行了高度最佳化。達到 LLM 可以協助編寫核心的階段是一個漸進的過程。

我們手動建立了一個初始的核心和公用程式函數語料庫，遵循一致的模式，然後才能依賴代理進行核心轉換。隨著語料庫的發展，我們可以執行從 CUDA 到 CuTeDSL 的粗略翻譯，隨後進行大量的手動重工。特別是，我們發現模型在指標到張量轉換和向量化方面遇到困難，在試圖讓測試通過編譯時，會卡在嘗試修復類型錯誤上。它們對核心基於網格的性質表現出很少的理解，混淆了執行緒和網格同步化。最令人沮喪的是，有一種誘惑是產生幻覺，去創造 CUDA 特定函式庫的不存在翻譯，而不是針對我們自己的 CuTeDSL 替代方案進行建構。隨著時間的推移，隨著 CuTeDSL 核心語料庫的增長，我們發現 LLM 可以建構大部分的骨架程式碼並填補空白，讓我們能夠專注於高層設計。

與 CUDA 相比，我們發現 CuTeDSL 提供了好得多的編譯時期效能，使我們能夠針對更多參數進行專門化並實現效能提升。雖然由於缺少與 Triton 相比的某些抽象，核心稍微有些冗長，但我們發現複雜度是合理的，因為對硬體的更細粒度控制允許我們建構更快的核心。

在撰寫本文時，CUTLASS C++ 和 CuTeDSL 核心之間的效能差距微乎其微，這為將高效能 GEMM 核心與裝置間通訊、正規化或激發進行組合和融合開啟了令人興奮的機會。透過致力於這項 DSL，我們將能夠從我們的裝置中提取更多效能，以改善我們從 embeddings 到 LLM 等各種工作負載的吞吐量。

https://github.com/Dao-AILab/quack

https://github.com/NVIDIA/cutlass

https://github.com/NVIDIA/nccl-tests

https://docs.nvidia.com/cuda/parallel-thread-execution/

BibTeX

@online{perplexity2026CutedslAtPerplexity,
  author = {{Perplexity} and Nandor Licker and Bex Brooker and Jesse Dwyer and Phi Hoang},
  title = {{CuTeDSL at Perplexity}},
  date = {2026-05-06},
  url = {https://www.perplexity.ai/hub/blog/cutedsl-at-perplexity}
}

Perplexity AI 的推論

為什麼選擇 CuTeDSL？

開發者體驗

大型配置空間

預填同解碼專門化

網格同步化

CuTeDSL 核心

QK Norm

專家混合模型 (MoE) 分派/合併

本地

單節點

資料並行

點對點

Blackwell 支援

AI 輔助核心生成

結論與未來工作

參考資料

引用