A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 23:00 UTC
[Bugfix] Fix the fp8_mqa_logits dim mismatch (#32652)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Jan 20, 2026 at 18:48 UTC
c4e5bdf61b00882749f9a4184c0576b4aa4ed512
2 files changed
+3
-3
vllm/model_executor/models/deepseek_v2.py
+1
-1
@@ -686,7 +686,7 @@ def sparse_attn_indexer(
686
fp8_mqa_logits_func = rocm_fp8_mqa_logits
687
logits = fp8_mqa_logits_func(
688
q_fp8[chunk.token_start : chunk.token_end],
689
- (k_fp8, k_scale.view(torch.float32)),
689
+ (k_fp8, k_scale.view(torch.float32).flatten()),
690
weights[chunk.token_start : chunk.token_end],
691
chunk.cu_seqlen_ks,
692
chunk.cu_seqlen_ke,
vllm/utils/deep_gemm.py
+2
-2
@@ -249,8 +249,8 @@ def fp8_mqa_logits(
249
q: Query tensor of shape [M, H, D]. Casted to
250
`torch.float8_e4m3fn` by caller.
251
kv: Tuple `(k_fp8, k_scales)` where `k_fp8` has shape [N, D] with
252
- dtype `torch.float8_e4m3fn` and `k_scales` has shape [N] (or
253
- [N, 1]) with dtype `torch.float32`.
252
+ dtype `torch.float8_e4m3fn` and `k_scales` has shape [N])
253
+ with dtype `torch.float32`.
254
weights: weights of shape [M, H], dtype `torch.float32`.
255
cu_seqlen_ks: Start indices (inclusive) for valid K per query position,
256
shape [M], dtype int32.