flashinfer.fp4_quantization.shuffle_matrix_sf_a

flashinfer.fp4_quantization.shuffle_matrix_sf_a(input_tensor: Tensor, epilogue_tile_m: int, num_elts_per_sf: int = 16)

Cuda 实现 trtllm-gen 的 shuffleMatrixSfA,但有一个注意事项。shuffleMatrixSfA 期望输入是 128x4 布局,然后应用与 shuffleMatrixA 相同的混洗,并以 128x4 布局写出。此函数期望输入是线性布局。这样做是因为 NVFP4 检查点的缩放因子被量化并且是线性布局。此函数不添加填充。