[Bugfix][Perf] Indexer upcast WK to BF16 for fusion

[benchislett]

Purpose

Alternative fix to #38870 which maintains the fusion.

Performance:

• B200 TP8 FP8 BS1 8k/1k
• Compares this PR (top), multi-stream ([Performance] DeepSeek V3.2 multi-stream indexer overlap #35968), and baseline ([Bugfix] Fix DSV32 weight loading #38870)

Upcast+Fused WK: (Decode): 11.90 ms
Upcast+Fused WK: (TTFT):   375.0 ms

Multi-Stream:    (Decode): 12.74 ms
Multi-Stream:    (TTFT):   376.5 ms

Separate:        (Decode): 12.59 ms
Separate:        (TTFT):   378.2 ms

Testing

My setup is broken, GSM8k is giving 0.00 for me even with #38870. Will try to fix my setup and rerun, but have moderate confidence in this fix. Would be handy if someone else could try running this in the meantime.

I'm okay with merging #38870 if we need to fix ASAP this weekend.

[zyongye]

I would certainly want to run the OG ckpt as it is with multi-stream.

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to handle DeepSeek-V2 checkpoints where 'wk' weights are stored in FP8 with separate scales. By dequantizing these weights to BF16 during the loading process, the model maintains fusion with 'weights_proj'. The implementation includes a new helper function, '_try_load_fp8_indexer_wk', and a buffering system to synchronize weights and scales. Feedback was provided regarding the potential for memory leaks in the buffering logic if the loading process is interrupted or if keys are missing from the checkpoint.

[gemini-code-assist]

The logic to buffer tensors until both weight and scale are present is susceptible to memory leaks if the checkpoint loading process is interrupted or if certain keys are missing from the checkpoint. Consider adding a timeout or a mechanism to clear the buffer if loading fails.

[benchislett]

They would go out of scope once the weight loading function returns. The references are not stored on self.

[benchislett]

I would certainly want to run the OG ckpt as it is with multi-stream.

This is the multi-stream baseline I compared against. It seems like the multi-streaming with the FP8 kernels is disabling torch.compile and undoing one of the elementwise op fusions, causing a slowdown. Could you explain what/why you want to run that I have not compared here?

[zyongye]

I would certainly want to run the OG ckpt as it is with multi-stream.

This is the multi-stream baseline I compared against. It seems like the multi-streaming with the FP8 kernels is disabling torch.compile and undoing one of the elementwise op fusions, causing a slowdown. Could you explain what/why you want to run that I have not compared here?

I am hoping that for the original model ckpt we want to let it run as its original precision, at least on one platform. Changing the precision will cause unknown problem later on and hard to trace back the bug.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @benchislett.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[benchislett]

I'm of the opinion that upcasting for the fusion is worth the trouble. The perf improvement seems nontrivial, and it does help avoid having to support both multi-stream and fusion for this op. @robertgshaw2-redhat to weight in

[zyongye]

If they really care about perf they should run the nvfp4 checkpoint. I feel like we want to have accuracy guarantee before perf on the OG checkpoint.

[benchislett]

There are many cases where we still care about perf for non-NVFP4 checkpoints. And as I understand, upcasting from FP8 to BF16 shouldn't hurt accuracy if done properly

[mgoin]

I think it is fair enough to support this, as we do for MLA upconvert as well before we had fp8_bmm available. Specific numerics may change, but accuracy should certainly not be worse

[benchislett]

Reran on 4xGB200, same perf results and GSM8k is good

Details

4xGB200, TP4, DSV3.2 FP8, Fused WK

GSM8k:

Results:
Accuracy: 0.955
Invalid responses: 0.000
Total latency: 236.064 s
Questions per second: 5.587
Total output tokens: 120664
Output tokens per second: 511.148

1k1k BS1:

============ Serving Benchmark Result ============
Successful requests:                     10        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  135.16    
Total input tokens:                      9704      
Total generated tokens:                  11079     
Request throughput (req/s):              0.07      
Output token throughput (tok/s):         81.97     
Peak output token throughput (tok/s):    86.00     
Peak concurrent requests:                2.00      
Total token throughput (tok/s):          153.76    
---------------Time to First Token----------------
Mean TTFT (ms):                          407.79    
Median TTFT (ms):                        158.75    
P99 TTFT (ms):                           2362.70   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          11.84     
Median TPOT (ms):                        11.80     
P99 TPOT (ms):                           11.98     
---------------Inter-token Latency----------------
Mean ITL (ms):                           11.84     
Median ITL (ms):                         11.81     
P99 ITL (ms):                            12.47     
==================================================

1k1k BS8:

============ Serving Benchmark Result ============
Successful requests:                     80        
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  193.45    
Total input tokens:                      82493     
Total generated tokens:                  82702     
Request throughput (req/s):              0.41      
Output token throughput (tok/s):         427.50    
Peak output token throughput (tok/s):    496.00    
Peak concurrent requests:                11.00     
Total token throughput (tok/s):          853.92    
---------------Time to First Token----------------
Mean TTFT (ms):                          224.54    
Median TTFT (ms):                        187.26    
P99 TTFT (ms):                           574.97    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          17.96     
Median TPOT (ms):                        18.10     
P99 TPOT (ms):                           18.59     
---------------Inter-token Latency----------------
Mean ITL (ms):                           17.96     
Median ITL (ms):                         17.14     
P99 ITL (ms):                            18.11     
==================================================

4xGB200, TP4, DSV3.2 FP8, No-Fused-WK

GSM8k:

Results:
Accuracy: 0.956
Invalid responses: 0.000
Total latency: 238.044 s
Questions per second: 5.541
Total output tokens: 120447
Output tokens per second: 505.987

1k1k BS1:

============ Serving Benchmark Result ============
Successful requests:                     10        
Failed requests:                         0         
Maximum request concurrency:             1         
Benchmark duration (s):                  139.43    
Total input tokens:                      9704      
Total generated tokens:                  11079     
Request throughput (req/s):              0.07      
Output token throughput (tok/s):         79.46     
Peak output token throughput (tok/s):    82.00     
Peak concurrent requests:                2.00      
Total token throughput (tok/s):          149.06    
---------------Time to First Token----------------
Mean TTFT (ms):                          178.87    
Median TTFT (ms):                        174.89    
P99 TTFT (ms):                           211.54    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          12.43     
Median TPOT (ms):                        12.39     
P99 TPOT (ms):                           12.57     
---------------Inter-token Latency----------------
Mean ITL (ms):                           12.43     
Median ITL (ms):                         12.40     
P99 ITL (ms):                            13.07     
==================================================

1k1k BS8:

============ Serving Benchmark Result ============
Successful requests:                     80        
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  200.50    
Total input tokens:                      82493     
Total generated tokens:                  82702     
Request throughput (req/s):              0.40      
Output token throughput (tok/s):         412.48    
Peak output token throughput (tok/s):    472.00    
Peak concurrent requests:                11.00     
Total token throughput (tok/s):          823.92    
---------------Time to First Token----------------
Mean TTFT (ms):                          233.02    
Median TTFT (ms):                        196.72    
P99 TTFT (ms):                           586.79    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          18.61     
Median TPOT (ms):                        18.71     
P99 TPOT (ms):                           19.20     
---------------Inter-token Latency----------------
Mean ITL (ms):                           18.62     
Median ITL (ms):                         17.74     
P99 ITL (ms):                            18.77     
==================================================

[benchislett]

I also manually inspected a few of the FP8 weight scales for the WK layers. They're all fairly small factors (e.g. 0.0002), so we should be staying comfortably in the range where BF16 is well-represented.