Perplexity 的 CuTeDSL

在 Perplexity，我們使用內部的 Runtime-Optimized Serving Engine (ROSE) 來提供從嵌入模型到兆級參數 LLM 的各類模型服務。作為所有 API（包括 Sonar、搜尋和嵌入）背後的支撐，我們的執行時期環境旨在賦能研究人員探索最合適的模型與架構，並快速將最佳候選模型部署到生產環境中。為達成此目標，我們需要能夠以極快的速度實作、部署及優化最先進架構的推論。在本篇部落格文章中，我們將探討 CuTeDSL 核心如何整合至我們的推論引擎，讓我們能夠建構必要的 GPU 核心，將 NVIDIA GPU 上的模型效能推向極致。

由 Perplexity 各個團隊開發的自訂模型，透過 AI 推論團隊建構的 ROSE 推論引擎，託管在內部的 NVIDIA Hopper 和 Blackwell GPU 上。ROSE 最初是為了更好地為語言解碼和分類提供客製化的 Llama 模型服務而建構的，並具備與 NVIDIA Triton 推論伺服器相容的介面。隨著時間的推移，它已發展成為一套完整的推論引擎，為 Perplexity 的產品服務廣泛的團隊，提供大型 LLM 以及排序、分類、評分和其他基於 Transformer 的模型服務。ROSE 的設計目標很簡單：雖然我們投入心力以求達到極致效能，但主要優先事項仍是適應不斷演變的模型架構與應用場景。

從高層架構來看，ROSE 由一個能將模型調整為用戶端可互動介面的引擎組成。這些引擎具有高度可重用性，能處理請求排程與批次處理，並從底層模型檢索解碼後的 Token、嵌入或分數。裝置初始化、處理序間與節點間的通訊通道以及權重載入，皆由引擎負責處理。對於 LLM，引擎負責批次處理、分塊與採樣；同時也負責處理 KV 儲存配置（針對完整與線性注意力機制）以及前綴匹配。嵌入引擎則較為簡單，僅執行線上批次處理。執行過程經過優化，使引擎能在模型於加速器上運行時，平行準備下一輪的工作。

ROSE 提供了豐富的自訂模型層集合，封裝了以 CuTeDSL、Triton、CUDA、CUTLASS 和 cuBLAS 建構的 GPU 核心。對於矩陣乘法與注意力實作，我們主要依賴來自 CUTLASS 和 cuBLAS 的 NVIDIA 核心。管道的其餘部分（包括嵌入、歸一化、MoE 和激活函數核心）則屬於 ROSE 的核心庫。為了最好地支援我們運行推論的硬體，ROSE 配備了專屬的裝置間通訊基元，用於 MoE 路由、分派與合併。雖然其他核心相對簡單，但它們必須經過徹底的專門化，才能在所運行的所有不同隱藏維度和裝置上實現極致效能。此外，我們依賴編譯時期的專門化來適應各種模型的調整與配置（例如 RMS 歸一化的權重偏差與無權重偏差、激活函數等）。因此，我們決定使用 DSL 來實作這些核心：我們最初使用 Triton 建構 ROSE，但自此之後，我們已採用 CuTeDSL 作為 GPU 程式設計的主要環境。

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

CuTeDSL 是一個基於 Python 的 NVIDIA GPU DSL，建構在 CuTe 佈局代數與 MLIR 之上。它能即時編譯為高度優化的 PTX，同時保有對低階硬體基元的控制權。其目標是提供與 CUTLASS C++ 相同的功能與效能，且無損於阻礙開發速度的昂貴編譯時間。在我們評估或使用過的語言中（包括 CUDA、cuTile、TVM IR 和 Triton），我們認為它最適合我們的核心開發工作。

總體而言，雖然不像 Triton 或 cuTile 那樣高階與簡潔，但我們發現 CuTe 佈局代數及其組合（透過 cute.Tensor 公開）是合理的替代方案。此外，CuTeDSL 提供對硬體基元的無限制存取，讓我們擁有在所有極端情況下實現極致效能所需的控制等級。使用 CUDA 時，編譯時期的專門化機會有限，因為每個額外的參數都會增加不合理的編譯時間成本，或增加建構系統的複雜性。相比之下，使用即時編譯的 DSL，我們可以進行更激進的專門化而不會犧牲迭代速度。與使用 CUDA 作為中間步驟的 TVM IR 不同，我們發現基於 MLIR 並下放到 PTX 的管道，比發出 CUDA 作為中間步驟的管道更穩健且更容易除錯。此外，透過使用自訂運算來封裝 PTX 指令，我們發現可以充分利用底層硬體的所有功能。

CUDA 核心大部分的編譯成本通常消耗在核心的模板展開階段，因為從 PTX 階段開始的優化與編譯可以在幾秒鐘內完成。隨著編譯時期參數的增加，對於非平凡的核心，這一編譯階段的時間可能會暴增至數分鐘。這個瓶頸更因為平行編譯的機會僅在後期的 IR 和 PTX 層級出現而加劇，而 C++ 模板處理則是以順序方式進行。雖然存在透過將模板實例化移至單獨編譯單元來解決此問題的技巧，但這需要對建構系統進行重大調整並使用笨重的巨集。相比之下，以即時 (JIT) 編譯的 CuTeDSL 核心不需要複雜的建構規則，也不會阻礙編譯時間，因為開發週期可以專注於配置空間中的單一點，而後續仔細的單元測試則可以涵蓋核心的整個設計空間。

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

CUDA 的沉重編譯成本更因在進行基於模板的分派與專門化時難以追蹤錯誤來源而加劇（例如 CuTe 和 CUTLASS 的情況）。使用 DSL，錯誤可以直接指向產生錯誤的程式碼行，大幅簡化了除錯過程。此外，透過中間 MLIR 階段，或依賴附加於所生成 PTX 的除錯資訊，將生成的程式碼對應回其原始的 DSL 語句也簡單得多。多虧了 DSL 的這些特性，我們發現在開發核心時生產力提高了不少。

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

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

輸入/輸出資料類型

量化方案

編譯時期常數（偏差、激活函數等）

編譯時期專門化對於效能與除錯都很有用。特別是在流水線處理器（如 GPU）中，不存在的分支比總是執行或從不執行的分支更好。此外，針對特定配置旋鈕編譯核心，會產生較少的 PTX/SASS 指令，這些指令在調查過程中執行 cuda-gdb 時更容易分析與除錯。

對於給定的批次大小與配置，預期的 DRAM 流量與算術運算數量是可以估算的。隨後，可以根據個別情況針對最佳網格分配、同步以及某些硬體功能的使用做出決定。CuTeDSL 允許我們定義與重複使用組件進行深度專門化，而不會產生無法接受的編譯時期開銷。

一般而言，推論核心的效能必須針對兩種不同的使用場景進行優化：

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

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

對於 GEMM 和注意力機制，預填充階段在大量 Token 上執行，效能受限於裝置的計算吞吐量。解碼涉及在小批次大小上執行相同的操作，執行時間主要由將激活函數與權重從 DRAM 載入到串流多處理器與 Tensor 核心的成本所主導，因此受限於裝置的記憶體頻寬。儘管批次大小或批次中的總 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，再加上在某些情況下透過隱式記憶體屏障重新整理先前所有記憶體操作的成本。總體而言，這與將核心拆分為兩個啟動的開銷相當，這是一個效能敏感的設計決策。儘管從 Hopper 開始，程式化相依啟動 (PDL) 可以減少啟動開銷，但對於解碼核心而言，避免此開銷是值得的。

如果其中一個工作負載需要的區塊數量超過 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 歸一化核心，針對小維度進行了優化。近期 Qwen 模型自注意力層中使用的 QK 歸一化，是基於每個頭、針對 64、128 或 256 個元素進行應用的。與在超過數千個隱藏維度上運作、需要透過共享記憶體和暫存器進行水平縮減的常規 RMS 歸一化不同，QK 歸一化僅需要在 Warp 之間進行縮減，且有可能將多個頭壓縮到單個 Warp 中。這成為了一個我們可以在編譯時期積極優化，以找到所有頭維度和資料類型的最佳 Token 到 Warp 分配的應用場景。

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

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 歸一化被確定並寫回記憶體。特定於模型的參數（例如 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 歸一化核心相比，此專門化變體為預填充和解碼提供了 2-3 倍的吞吐量與延遲改進。

MoE 分派與合併將輸入 Token 洗牌到其路由到的專家，並在 MoE GEMM 核心執行相關矩陣乘法後執行加權累積。效能特別受到關注，因為各種產品都依賴不同大小的 MoE 模型。根據部署類型，我們區分了不同的核心進行優化：

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

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

DP ≠ 1，NVLink：具有選擇性張量並行的資料並行核心。每個 Rank 向託管特定專家的每個其他 Rank 發送 Token。路由資訊被聚合，Token 透過 NVIDIA NVLink 傳輸。用於資料並行 H200 或 GB200 部署。

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

本地核心根據路由表計算分配給每個專家的 Token 數量，並為屬於每個專家的 Token 列表中的每個 Token 分配唯一索引。路由步驟後，Token 被複製到可選的填充緩衝區中，以傳遞給 MoE GEMM 核心。合併核心隨後利用 Token 位置在暫存器中執行累積，將最終 Token 寫入輸出緩衝區。

分派階段在預填充與解碼之間有所不同：對於少量 Token，讀取的路由資訊很少，因此我們在單個區塊中處理它並向網格其餘部分發出訊號，個別區塊複製 Token 的區塊。Token 分配由原子遞增共享記憶體中的計數器來決定。由於預填充必須處理更多路由資訊，因此路由在單獨的核心中拆分為多個區塊。索引首先使用共享記憶體中的計數器在本地分配，然後在全域記憶體中跨所有區塊聚合，以確定每個 Token 的位置。此核心的網格被限制為 SM 的數量，以便能夠使用網格級屏障。隨後的核心執行複製，在更大數量的區塊與 Warp 上啟動，以充分利用可用記憶體頻寬。

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

在張量並行部署中，路由資訊與輸入 Token 會被複製。本地分派核心可以完全重複使用，只需稍作更改，僅路由與選擇分配給目前 Rank 上專家的 Token。透過一些參數與編譯時期專門化，本地分派核心得以重複使用。

由於 Rank 的數量通常少於 Token 路由到的專家數量，一般情況下可以假設每個 Token 至少有一個副本路由到每個 Rank。在這種情況下，本地合併核心可以重複使用來聚合 Token 的本地副本，隨後進行全域歸約 (all-reduce) 以在參與的 Rank 之間加總。對於預填充核心，我們啟動單獨的核心，但解碼將這兩個操作融合為單次啟動。

對於全域歸約，如果只有兩個 Rank 參與，我們使用普通 NVLink；當涉及 4 個或更多 Rank 時，我們切換到 NVLS 以獲得 10-20% 的吞吐量提升。多虧了 DSL，我們可以共享 MoE 與常規全域歸約核心之間的大部分歸約邏輯與對等緩衝區交換。這是在硬體上進行精細控制有益的情況，因為 NVLS 需要針對 multimem.ld_reduce 和 multimem.st 的 DSL 操作。歸約本身的實作在很大程度上遵循 NCCL 測試中的 all_reduce 範例。

資料並行核心依賴 NVLink 來聚合多個裝置的 Token。在分派階段，交換 Token 之前，對等點首先在本地路由，以識別它們發送的每個 Token 在每個 Rank 上的目標偏移量。隨後，它們開始將 Token 複製到為每個發送者配置的接收緩衝區中，以避免衝突。使用每 Token 發布-儲存 (release-store) 寫入來標記完成。在接收端，路由資訊被聚合，以便能夠將所有 Token 濃縮為單個緩衝區以用於 MoE GEMM。接收器使用獲取-載入 (acquire-load) 輪詢對等點寫入的旗標，以等待寫入完成。

在合併端，Token 透過 NVLink 複製到來源 Rank，並以每 Token 旗標指示接收。合併核心與本地與單節點變體相似，在定址與一些額外的同步邏輯上略有不同。

我們建構了一個 PTX 包裝器庫，以公開系統與裝置層級的發布-獲取負載與儲存，來實作此核心。透過提高同步旗標的顆粒度，並分配更多區塊以透過 NVLink 發送與接收 Token，我們與之前的 CUDA 核心相比實現了 2 倍的改進。

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

從 ROSE 核心的角度來看，Blackwell 支援分為兩類：與架構無關的核心與與架構相關的核心。依賴 Tensor 核心的與架構相關的核心需要明確的工作來適應。然而，多虧了 DSL 的使用，與架構無關的核心可以輕鬆適應。

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

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

作為一種相對較新的語言，LLM 在其預訓練期間很少接觸到 CuTeDSL 程式碼。此外，我們目前程式碼的風格（針對激活函數與歸一化優化過的核心）與現有的 CUTLASS 或 Quack 範例略有不同，後者針對 TMA 和特定架構的矩陣乘法指令進行了深度優化。LLM 能夠協助編寫核心是一個循序漸進的過程。

在我們能夠依賴代理進行核心轉換之前，我們手動建立了一個初始的核心與實用函數語料庫，遵循一致的模式。隨著語料庫的發展，我們能夠進行從 CUDA 到 CuTeDSL 的初步轉換，隨後進行了大量手動重構。特別是，我們發現模型在指標到 Tensor 的轉換和向量化方面遇到了困難，在試圖編譯測試時陷入了修正型別錯誤的困境。它們對核心的網格本質顯示出很少的理解，混淆了討論串與網格同步。最令人沮喪的是，它們傾向於幻覺出非現有的 CUDA 特定庫轉換，而不是根據我們自己的 CuTeDSL 替換項進行建構。隨著時間的推移，隨著 CuTeDSL 核心語料庫的增長，我們發現 LLM 可以建構大部分骨架程式碼並填寫空白，讓我們能夠專注於高層設計。

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

在撰寫本文時，CUTLASS C++ 與 CuTeDSL 核心之間的效能差距微乎其微，為組合與融合高效能 GEMM 核心與裝置間通訊、歸一化或激活函數開啟了令人興奮的機會。透過致力於此 DSL，我們將能夠從裝置中提取更多效能，以改善我們從嵌入到 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 歸一化

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

本地

單節點

資料並行

點對點

Blackwell 支援

AI 輔助核心生成

結論與未來工作

參考文獻

引用