From 99a7550a51d434e468c9fd3765d5c94014c0fe79 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Wed, 19 Aug 2026 22:31:52 +0700 Subject: [PATCH] Debug _merge_tokens prod --- src/h3_blackwell_runtime/qwen3vl_vision.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/h3_blackwell_runtime/qwen3vl_vision.py b/src/h3_blackwell_runtime/qwen3vl_vision.py index e4b5a9b..1391f01 100644 --- a/src/h3_blackwell_runtime/qwen3vl_vision.py +++ b/src/h3_blackwell_runtime/qwen3vl_vision.py @@ -674,12 +674,14 @@ class Qwen3VL32BVision(nn.Module): @staticmethod def _merge_tokens(x: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: """Flatten 2x2 spatial-merge blocks (the reference's pre-merge flatten).""" - # grid_thw is a 2D long tensor [[t, h, w], ...]; take the first image. first = grid_thw[0] t = int(first[0].item()) h = int(first[1].item()) w = int(first[2].item()) merge = 2 + print(f"[merge_tokens] x.shape={tuple(x.shape)}, t={t}, h={h}, w={w}, " + f"target=({t}, {h//merge}, {merge}, {w//merge}, {merge}, {x.shape[-1]}), " + f"prod={t*(h//merge)*merge*(w//merge)*merge*x.shape[-1]}", flush=True) return x.view((t, h // merge, merge, w // merge, merge, x.shape[-1])).permute(0, 1, 3, 2, 4, 5).reshape(-1, -1) def forward(self, flatten_patches: torch.Tensor, grid_thw: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: