(expert_weights: torch.Tensor, block_size: torch.Tensor, dtype: torch.dtype)
| 275 | return out, scale |
| 276 | |
| 277 | def per_expert_blockwise_quantize(expert_weights: torch.Tensor, block_size: torch.Tensor, dtype: torch.dtype) -> Tuple[torch.Tensor, torch.Tensor]: |
| 278 | assert expert_weights.ndim == 3 |
| 279 | num_experts = expert_weights.shape[0] |
| 280 | output_weights = [] |
| 281 | scales = [] |
| 282 | for e in range(num_experts): |
| 283 | weight, scale = blockwise_quantize(expert_weights[e], block_size, dtype) |
| 284 | output_weights.append(weight) |
| 285 | scales.append(scale) |
| 286 | return torch.stack(output_weights), torch.stack(scales) |
| 287 | |
| 288 | def per_expert_k_quantize(expert_weights: torch.Tensor, method: Literal["q2_k", "q3_k"]) -> torch.Tensor: |
| 289 | assert expert_weights.ndim == 3 |
no test coverage detected