| 519 | } |
| 520 | |
| 521 | void code1x8_matvec( |
| 522 | const torch::Tensor& A, |
| 523 | const torch::Tensor& B, |
| 524 | torch::Tensor& C, |
| 525 | const torch::Tensor& codebook, |
| 526 | bool use_bfloat16 |
| 527 | ) { |
| 528 | const at::cuda::OptionalCUDAGuard device_guard(device_of(A)); |
| 529 | int prob_m = C.size(0); |
| 530 | int prob_k = B.size(0); |
| 531 | if (use_bfloat16) { |
| 532 | code1x8_matvec_cuda<true>( |
| 533 | A.data_ptr(), |
| 534 | B.data_ptr(), |
| 535 | C.data_ptr(), |
| 536 | codebook.data_ptr(), |
| 537 | prob_m, |
| 538 | prob_k |
| 539 | ); |
| 540 | } else { |
| 541 | code1x8_matvec_cuda<false>( |
| 542 | A.data_ptr(), |
| 543 | B.data_ptr(), |
| 544 | C.data_ptr(), |
| 545 | codebook.data_ptr(), |
| 546 | prob_m, |
| 547 | prob_k |
| 548 | ); |
| 549 | } |
| 550 | } |
| 551 | |
| 552 | torch::Tensor code1x8_matmat( |
| 553 | const torch::Tensor& input, |