| 421 | } |
| 422 | |
| 423 | torch::Tensor code2x8_dequant( |
| 424 | const torch::Tensor& codes, |
| 425 | const torch::Tensor& codebooks, |
| 426 | const torch::Tensor& scales |
| 427 | ) { |
| 428 | auto use_bfloat16 = check_use_bfloat16(codebooks); |
| 429 | auto in_features = codes.size(1) * 8; |
| 430 | auto out_features = scales.size(0); |
| 431 | |
| 432 | auto weight = torch::empty({out_features, in_features}, |
| 433 | torch::TensorOptions() |
| 434 | .dtype(codebooks.dtype()) |
| 435 | .device(codebooks.device()) |
| 436 | ); |
| 437 | code2x8_dequant_cuda( |
| 438 | codes.data_ptr(), |
| 439 | weight.data_ptr(), |
| 440 | codebooks.data_ptr(), |
| 441 | out_features, |
| 442 | in_features, |
| 443 | use_bfloat16 |
| 444 | ); |
| 445 | weight *= scales.index({"...", 0, 0}); |
| 446 | |
| 447 | return weight; |
| 448 | } |
| 449 | |
| 450 | torch::Tensor code2x8_matmat_dequant( |
| 451 | const torch::Tensor& input, |
nothing calls this directly
no test coverage detected