fix positions device mismatch: move to rope cache device in forward_attention

This commit is contained in:
2026-05-31 21:54:56 +00:00
parent 2676476013
commit 0d2b5ceb93

View File

@@ -432,6 +432,9 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
ratio = compressor.ratio if compressor is not None else 0
scale = 1.0 / math.sqrt(hd)
pfx = f"model.layers.{li}.self_attn"
# Ensure positions is on the same device as rope caches
if positions.device != rope_cos.device:
positions = positions.to(rope_cos.device)
# 1. Q projection: q_a → q_a_norm → q_b → q_b_norm
q_a = do_nvfp4_linear(x_normed, w, pfx, 'q_a_proj')