daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Debug _merge_tokens prod

Browse source
This commit is contained in:
Daniel Maddern 2026-08-19 22:31:52 +07:00
parent ece7b2497e
commit 99a7550a51
1 changed files with 3 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/h3_blackwell_runtime/qwen3vl_vision.py
View file

@ -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]]: