[Bugfix][DCP] Cast LSE to fp32 in a2a combine to fix bf16 bitcast crash (#47801)
Signed-off-by: Shawn Tsai <shawnyht@gmail.com>
shawn committed
Jul 8, 2026 at 08:41 UTC
f05603fa287ab020acc5faafd49c5decf0762533
1 file changed
+4
vllm/v1/attention/ops/dcp_alltoall.py
+4
@@ -426,6 +426,10 @@ def dcp_a2a_lse_reduce(
426
if H % world_size != 0:
427
raise ValueError(f"H={H} must be divisible by DCP world size {world_size}.")
428
H_per_rank = H // world_size
429
+ # The pack kernel bit-casts the LSE as fp32; some MLA backends return it in
430
+ # the activation dtype (bf16/fp16), so enforce the documented fp32 contract.
431
+ if cp_attn_lse.dtype != torch.float32:
432
+ cp_attn_lse = cp_attn_lse.to(torch.float32)
433
lse_pack_dim = _dcp_a2a_lse_pack_dim(cp_attn_out.dtype)
434
435
send_buffer, recv_buffer = _dcp_a2a_send_recv_buffers(