A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Perf] Apply torch.compile for `per_block_cast_to_fp8` (#24611)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 22, 2025 at 21:42 UTC
9949aa2ef1cc77721dea9077e0f76b9deb6fd066
1 file changed
+2
-2
vllm/utils/deep_gemm.py
+2
-2
@@ -135,7 +135,7 @@ DEFAULT_BLOCK_SIZE = [128, 128]
135
136
137
# Taken from https://github.com/deepseek-ai/DeepGEMM/blob/dd6ed14acbc7445dcef224248a77ab4d22b5f240/deep_gemm/utils/math.py#L38
138
-# TODO(wentao): optimize this function, using triton or cuda kernel
138
+@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
139
def per_block_cast_to_fp8(
140
x: torch.Tensor,
141
block_size: list[int] = DEFAULT_BLOCK_SIZE,
@@ -187,4 +187,4 @@ __all__ = [
187
"is_deep_gemm_e8m0_used",
188
"is_deep_gemm_supported",
189
"should_use_deepgemm_for_fp8_linear",
190
-]
190
+]
\ No newline at end of file