submission 611052

  ⋯ 4 unchanged lines
   
  torch.backends.cuda.matmul.allow_tf32 = True
   
+ # cuBLASLt with autotuned algorithm selection via load_inline
+ _cuda_src = r"""
+ #include <cuda_runtime.h>
+ #include <cublasLt.h>
+ #include <cuda_fp16.h>
+ #include <torch/extension.h>
+ #include <unordered_map>
+  
+ static cublasLtHandle_t ltHandle = nullptr;
+ static void* workspace = nullptr;
+ static const size_t workspaceSize = 64 * 1024 * 1024;
+  
+ struct CachedPlan {
+     cublasLtMatmulDesc_t matmulDesc;
+     cublasLtMatrixLayout_t Adesc, Bdesc, Cdesc;
+     cublasLtMatmulAlgo_t algo;
+ };
+  
+ static std::unordered_map<uint64_t, CachedPlan> planCache;
+  
+ static inline uint64_t make_key(int M, int K, int N) {
+     return ((uint64_t)M << 40) | ((uint64_t)K << 20) | (uint64_t)N;
+ }
+  
+ static void ensure_init() {
+     if (!ltHandle) {
+         cublasLtCreate(&ltHandle);
+         cudaMalloc(&workspace, workspaceSize);
+     }
+ }
+  
+ static void do_autotune(int M, int K, int N) {
+     ensure_init();
+     uint64_t key = make_key(M, K, N);
+     if (planCache.count(key)) return;
+  
+     CachedPlan plan;
+     cublasLtMatmulDescCreate(&plan.matmulDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+     cublasLtMatrixLayoutCreate(&plan.Adesc, CUDA_R_16F, K, M, K);
+     cublasLtMatrixLayoutCreate(&plan.Bdesc, CUDA_R_16F, N, K, N);
+     cublasLtMatrixLayoutCreate(&plan.Cdesc, CUDA_R_16F, N, M, N);
+  
+     cublasOperation_t opN = CUBLAS_OP_N;
+     cublasLtMatmulDescSetAttribute(plan.matmulDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opN, sizeof(opN));
+     cublasLtMatmulDescSetAttribute(plan.matmulDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opN, sizeof(opN));
+  
+     cublasLtMatmulPreference_t pref;
+     cublasLtMatmulPreferenceCreate(&pref);
+     cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize));
+  
+     const int kMaxAlgos = 32;
+     cublasLtMatmulHeuristicResult_t heu[kMaxAlgos];
+     int nResults = 0;
+     cublasLtMatmulAlgoGetHeuristic(ltHandle, plan.matmulDesc, plan.Bdesc, plan.Adesc,
+                                     plan.Cdesc, plan.Cdesc, pref, kMaxAlgos, heu, &nResults);
+  
+     __half *dA, *dB, *dC;
+     cudaMalloc(&dA, (size_t)M * K * sizeof(__half));
+     cudaMalloc(&dB, (size_t)K * N * sizeof(__half));
+     cudaMalloc(&dC, (size_t)M * N * sizeof(__half));
+  
+     cudaEvent_t t0, t1;
+     cudaEventCreate(&t0);
+     cudaEventCreate(&t1);
+  
+     float alpha = 1.0f, beta = 0.0f;
+     float bestTime = 1e30f;
+     int bestIdx = 0;
+  
+     for (int i = 0; i < nResults; i++) {
+         for (int w = 0; w < 10; w++)
+             cublasLtMatmul(ltHandle, plan.matmulDesc, &alpha,
+                 dB, plan.Bdesc, dA, plan.Adesc, &beta,
+                 dC, plan.Cdesc, dC, plan.Cdesc,
+                 &heu[i].algo, workspace, workspaceSize, 0);
+         cudaDeviceSynchronize();
+  
+         float total = 0;
+         for (int r = 0; r < 3; r++) {
+             cudaEventRecord(t0, 0);
+             for (int j = 0; j < 100; j++)
+                 cublasLtMatmul(ltHandle, plan.matmulDesc, &alpha,
+                     dB, plan.Bdesc, dA, plan.Adesc, &beta,
+                     dC, plan.Cdesc, dC, plan.Cdesc,
+                     &heu[i].algo, workspace, workspaceSize, 0);
+             cudaEventRecord(t1, 0);
+             cudaEventSynchronize(t1);
+             float ms;
+             cudaEventElapsedTime(&ms, t0, t1);
+             total += ms / 100;
+         }
+         if (total / 3 < bestTime) {
+             bestTime = total / 3;
+             bestIdx = i;
+         }
+     }
+  
+     plan.algo = heu[bestIdx].algo;
+     planCache[key] = plan;
+  
+     cudaEventDestroy(t0);
+     cudaEventDestroy(t1);
+     cudaFree(dA);
+     cudaFree(dB);
+     cudaFree(dC);
+     cublasLtMatmulPreferenceDestroy(pref);
+ }
+  
+ void matmul_cublaslt(torch::Tensor A, torch::Tensor B, torch::Tensor C) {
+     int M = A.size(0);
+     int K = A.size(1);
+     int N = B.size(1);
+     uint64_t key = make_key(M, K, N);
+     auto it = planCache.find(key);
+     if (it == planCache.end()) {
+         do_autotune(M, K, N);
+         it = planCache.find(key);
+     }
+     const CachedPlan& p = it->second;
+     float alpha = 1.0f, beta = 0.0f;
+     cublasLtMatmul(ltHandle, p.matmulDesc, &alpha,
+         B.data_ptr(), p.Bdesc, A.data_ptr(), p.Adesc, &beta,
+         C.data_ptr(), p.Cdesc, C.data_ptr(), p.Cdesc,
+         &p.algo, workspace, workspaceSize, 0);
+ }
+ """
+  
+ _module = None
+  
+ def _get_module():
+ 	global _module
+ 	if _module is None:
+ 		from torch.utils.cpp_extension import load_inline
+ 		_module = load_inline(
+ 			name="matmul_cublaslt",
+ 			cpp_sources="void matmul_cublaslt(torch::Tensor A, torch::Tensor B, torch::Tensor C);",
+ 			cuda_sources=_cuda_src,
+ 			functions=["matmul_cublaslt"],
+ 			extra_ldflags=["-lcublasLt", "-lcublas"],
+ 			verbose=False,
+ 		)
+ 	return _module
+  
+ # Eagerly compile
+ _get_module()
+  
  def custom_kernel(data: input_t) -> output_t:
  	a, b, c = data
- 	torch.mm(a, b, out=c)
+ 	_get_module().matmul_cublaslt(a, b, c)
  	return c