[Kernels][DP/EP] Optimize Silu Kernel for R1 (#24054)

Signed-off-by: elvircrn <elvircrn@gmail.com>
This commit is contained in:
Elvir Crnčević
2025-09-13 09:17:27 +02:00
committed by GitHub
parent dbeee3844c
commit 98229db244
6 changed files with 1269 additions and 128 deletions

View File

@@ -32,6 +32,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
#define stride_tag
#endif
ops.def(
"silu_mul_fp8_quant_deep_gemm_cuda(Tensor input, Tensor counts, Tensor! "
"y_q, Tensor! y_s, int group_size, "
"bool use_ue8m0, int num_parallel_tokens) -> ()");
ops.impl("silu_mul_fp8_quant_deep_gemm_cuda", torch::kCUDA,
&silu_mul_fp8_quant_deep_gemm_cuda);
ops.def("weak_ref_tensor(Tensor input) -> Tensor");
ops.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);