[Bugfix] Fix broken MTP weight loading for FP8 KV Scales

[benchislett]

Purpose

maybe_remap_kv_scale_name is a part of the DS weight loading logic, but it missing from MTP. This causes a crash when loading, for example, nvidia/DeepSeek-R1-0528-FP4 with MTP which uses this feature by default.

Test Plan

LM Eval GSM8k of nvidia/DeepSeek-R1-0528-FP4 on 4xB200:

Test Result

local-completions (base_url=http://0.0.0.0:8049/v1/completions,model=nvidia/DeepSeek-R1-0528-FP4,tokenized_requests=False,tokenizer_backend=None,num_concurrent=128,timeout=120,max_retries=5), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9530|±  |0.0058|
|     |       |strict-match    |     5|exact_match|↑  |0.9484|±  |0.0061|

[gemini-code-assist]

Code Review

This pull request addresses a bug that caused a crash when loading models with FP8 KV scales using Multi-Token Prediction (MTP). The fix involves incorporating the maybe_remap_kv_scale_name function into the MTP weight loading process, which was previously missing. This change aligns the MTP weight loading logic with the existing implementation for the base model, ensuring that KV scale names are correctly remapped. The implementation is correct and directly resolves the reported issue. The code is clean and there are no further issues.