Add prefill attention value debug check

This commit is contained in:
2026-05-19 17:55:35 +00:00
parent 4f02113aa0
commit 0773c9608c

View File

@@ -751,6 +751,13 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
o[num_decode_tokens:] = causal_prefill_attention(
q_prefill, kv_rope_prefill, self.scale,
)
# Debug: check attention output
import sys as _sys
_amax = o[num_decode_tokens:].amax().item()
_nan = torch.isnan(o[num_decode_tokens:]).any().item()
_std = o[num_decode_tokens:].float().std().item()
if _amax > 100 or _nan or _std < 0.001:
print(f"[BLACKWELL] PREFILL CHECK: amax={_amax:.4f} NaN={_nan} std={_std:.6f} cr={self.compress_ratio}", file=_sys.stderr, flush=True)
# Write into the output buffer
if self.n_local_heads < self.padded_heads: