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

Debug _merge_tokens shapes

Browse source
This commit is contained in:
Daniel Maddern 2026-08-19 22:21:30 +07:00
parent ffb3c9e729
commit b9dcdda116
1 changed files with 2 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
src/h3_blackwell_runtime/qwen3vl_vision.py
View file

@ -676,6 +676,8 @@ class Qwen3VL32BVision(nn.Module):
"""Flatten 2x2 spatial-merge blocks (the reference's pre-merge flatten)."""
t, h, w = (int(v) for v in grid_thw.tolist()[0])
merge = 2
if x.ndim != 2 or x.shape[0] != t * h * w:
print(f"[merge_tokens] x.shape={tuple(x.shape)}, grid={(t, h, w)}, expected N={t*h*w}")
return x.view(t, h // merge, merge, w // merge, merge, -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]]: