| 354 | } |
| 355 | |
| 356 | void code2x8_matvec( |
| 357 | const torch::Tensor& A, |
| 358 | const torch::Tensor& B, |
| 359 | torch::Tensor& C, |
| 360 | const torch::Tensor& codebook, |
| 361 | bool use_bfloat16 |
| 362 | ) { |
| 363 | const at::cuda::OptionalCUDAGuard device_guard(device_of(A)); |
| 364 | int prob_m = C.size(0); |
| 365 | int prob_k = B.size(0); |
| 366 | if (use_bfloat16) { |
| 367 | code2x8_matvec_cuda<true>( |
| 368 | A.data_ptr(), |
| 369 | B.data_ptr(), |
| 370 | C.data_ptr(), |
| 371 | codebook.data_ptr(), |
| 372 | prob_m, |
| 373 | prob_k |
| 374 | ); |
| 375 | } else { |
| 376 | code2x8_matvec_cuda<false>( |
| 377 | A.data_ptr(), |
| 378 | B.data_ptr(), |
| 379 | C.data_ptr(), |
| 380 | codebook.data_ptr(), |
| 381 | prob_m, |
| 382 | prob_k |
| 383 | ); |
| 384 | } |
| 385 | } |
| 386 | |
| 387 | torch::Tensor code2x8_matmat( |
| 388 | const torch::Tensor& input, |