A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 07:15 UTC
Nemotron Nano VL: Streamline pixel shuffle (#37580)
Signed-off-by: milesial <milesial@users.noreply.github.com>
milesial committed
Apr 10, 2026 at 00:31 UTC
270e8a41025436c86de96b76317a0029b30416c9
1 file changed
+11
-36
vllm/model_executor/models/nano_nemotron_vl.py
+11
-36
@@ -1005,38 +1005,27 @@ class NemotronH_Nano_VL_V2(
1005
)
1006
1007
def pixel_shuffle(self, x, scale_factor=0.5):
1008
- n, w, h, c = x.size()
1009
- # N, W, H, C --> N, W, H * scale, C // scale
1010
- x = x.view(
1011
- n,
1012
- w,
1013
- int(h * scale_factor),
1014
- int(c / scale_factor),
1015
- )
1016
- # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
1017
- x = x.permute(0, 2, 1, 3).contiguous()
1018
- # N, H * scale, W, C // scale -->
1019
- # N, H * scale, W * scale, C // (scale ** 2)
1020
- x = x.view(
1021
- n,
1022
- int(h * scale_factor),
1023
- int(w * scale_factor),
1024
- int(c / (scale_factor * scale_factor)),
1025
- )
1008
+ n, h, w, c = x.size()
1009
+ r = int(1 / scale_factor)
1010
+ new_h = h // r
1011
+ new_w = w // r
1012
+ new_c = c * r * r
1013
+
1014
+ x = x.view(n, new_h, r, new_w, r, c)
1015
if self.ps_version == "v1":
1016
warnings.warn(
1017
"In ps_version 'v1', the height and width have not "
1018
"been swapped back, which results in a transposed image.",
1019
stacklevel=2,
1020
)
1021
+ x = x.permute(0, 3, 1, 2, 4, 5).reshape(n, new_w, new_h, new_c)
1022
else:
1033
- x = x.permute(0, 2, 1, 3).contiguous()
1023
+ x = x.permute(0, 1, 3, 2, 4, 5).reshape(n, new_h, new_w, new_c)
1024
return x
1025
1026
def pixel_shuffle_dynamic_res(
1027
self, x: torch.Tensor, *, imgs_sizes: list[tuple[int, int]]
1028
) -> torch.Tensor:
1039
- scale_factor = self.downsample_ratio
1029
patch_dim = self.patch_size
1030
seq_lens = calc_seq_lens(imgs_sizes, patch_dim)
1031
splits = torch.split(x, seq_lens, dim=-2)
@@ -1045,22 +1034,8 @@ class NemotronH_Nano_VL_V2(
1034
h = imgs_sizes[i][0] // patch_dim
1035
w = imgs_sizes[i][1] // patch_dim
1036
sv = sv.reshape(sv.shape[0], h, w, -1)
1048
-
1049
- n, h, w, c = sv.size()
1050
-
1051
- sv = sv.view(n, h, int(w * scale_factor), int(c / scale_factor))
1052
- sv = sv.permute(0, 2, 1, 3).contiguous()
1053
- sv = sv.view(
1054
- n,
1055
- int(w * scale_factor),
1056
- int(h * scale_factor),
1057
- int(c / (scale_factor * scale_factor)),
1058
- )
1059
-
1060
- if self.ps_version == "v2":
1061
- sv = sv.permute(0, 2, 1, 3).contiguous()
1062
-
1063
- sv = sv.reshape(sv.shape[0], -1, sv.shape[-1])
1037
+ sv = self.pixel_shuffle(sv, scale_factor=self.downsample_ratio)
1038
+ sv = sv.flatten(1, 2)
1039
out.append(sv)
1040
1041
x = torch.cat(out, dim=-2)