(weights: Dict[str, torch.Tensor])
| 1873 | pattern_info = ['fc', 'gate', 'proj', 'qkv', 'dense'] |
| 1874 | |
| 1875 | def process_kv_scaling_factor(weights: Dict[str, torch.Tensor]): |
| 1876 | new_entries = {} |
| 1877 | names_to_delete = set() |
| 1878 | |
| 1879 | # If k, v cache scaling factors are stored separately, combine them into kv cache scaling factor. |
| 1880 | for name, param in weights.items(): |
| 1881 | if name.endswith('.k_cache_scaling_factor'): |
| 1882 | v_name = name.replace('k_cache_scaling_factor', |
| 1883 | 'v_cache_scaling_factor') |
| 1884 | assert v_name in weights, f"{v_name} not found" |
| 1885 | kv_name = name.replace('k_cache_scaling_factor', |
| 1886 | 'kv_cache_scaling_factor') |
| 1887 | new_entries[kv_name] = torch.max(weights[name], weights[v_name]) |
| 1888 | names_to_delete.update([name, v_name]) |
| 1889 | weights.update(new_entries) |
| 1890 | for k in names_to_delete: |
| 1891 | del weights[k] |
| 1892 | |
| 1893 | new_entries = [] |
| 1894 | # The unified converter generate_tllm_weights() already generates these rcp weights, but legacy |
| 1895 | # converters do not. Handle it here. |
| 1896 | for name, param in weights.items(): |
| 1897 | if name.endswith('.kv_cache_scaling_factor'): |
| 1898 | rcp_name = name.replace('kv_cache_scaling_factor', |
| 1899 | 'kv_cache_rcp_scaling_factor') |
| 1900 | if rcp_name not in weights: |
| 1901 | new_entries.append((rcp_name, torch.reciprocal(param))) |
| 1902 | weights.update(new_entries) |
| 1903 | |
| 1904 | process_kv_scaling_factor(weights) |
| 1905 |
no test coverage detected