flashinfer.trtllm_low_latency_gemm.prepare_low_latency_gemm_weights¶

flashinfer.trtllm_low_latency_gemm.prepare_low_latency_gemm_weights(w: Tensor, permutation_indices_cache: Dict[Size, Tensor]) → Tensor¶

Helper method to prepare the input weight tensor for low-latency TRTLLM GEMM. It includes shuffling and converting to block layout.

Parameters:

w (torch.Tensor) – The weight tensor to shuffle, shape (n, k), fp8 e4m3.
permutation_indices_cache (dict) – Some location to cache permutation indices. Calculating them is expensive.

Returns:

block_layout_shuffled_weights – The shuffled and block-layout weight tensor, shape (k // 128, n, 128), fp8 e4m3.

Return type:

torch.Tensor