Skip to content

Commit

Permalink
fix: compatible with torch 2.2 (#478)
Browse files Browse the repository at this point in the history
torch 2.2 no member function `getCurrentCUDABlasLtHandle`
torch 2.3 and 2.4 works well
  • Loading branch information
zhyncs authored Aug 27, 2024
1 parent 1a6b17e commit ac41d1b
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion include/flashinfer/bmm_fp8.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,7 @@ void bmm_fp8_internal_cublaslt(const AT* A, const BT* B, DT* D, int batch_size,
auto workspace = allocator.allocate(workspace_size);
cublasLtMatmulHeuristicResult_t heuristic_result = {};
int returned_result = 0;
auto lt_handle = at::cuda::getCurrentCUDABlasLtHandle();
auto lt_handle = reinterpret_cast<cublasLtHandle_t>(at::cuda::getCurrentCUDABlasHandle());
TORCH_CUDABLAS_CHECK(cublasLtMatmulAlgoGetHeuristic(
lt_handle, matmul_desp.descriptor(), a_desp.descriptor(), b_desp.descriptor(),
d_desp.descriptor(), d_desp.descriptor(), preference.descriptor(), 1, &heuristic_result,
Expand Down

0 comments on commit ac41d1b

Please sign in to comment.