@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 16070ed7e9

[Minimax-M3][Perf] Use triton_mrope for vision tower + int64 offset fix for triton_mrope (#58526)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Kimi <noreply@moonshot.cn> Co-authored-by: Isotr0py <Isotr0py@outlook.com>

Thien Tran committed Sep 25, 2026 at 09:30 UTC 16070ed7e9b051cb0b1958d4289912157026d2f4

4 files changed +280 -179

tests/kernels/core/test_mrope.py

+73 -1

@@ -6,7 +6,10 @@ import pytest

6 import torch

8 from vllm.model_executor.layers.rotary_embedding import get_rope

9 -from vllm.model_executor.layers.rotary_embedding.mrope import apply_interleaved_rope

9 +from vllm.model_executor.layers.rotary_embedding.mrope import (

10 + apply_interleaved_rope,

11 + triton_mrope,

12 +)

13 from vllm.platforms import current_platform

14 from vllm.transformers_utils.config import get_config

15 from vllm.utils.torch_utils import set_random_seed

@@ -174,6 +177,75 @@ def test_mrope(

177 torch.testing.assert_close(query_native, query_cuda, atol=atol, rtol=rtol)

178 torch.testing.assert_close(key_native, key_cuda, atol=atol, rtol=rtol)

179

180 + # Strided q/k views of a packed buffer (the MiniMax M3 ViT usage):

181 + # bitwise-identical to the contiguous run, neighboring columns untouched.

182 + pad = 16

183 + packed = torch.zeros(

184 + num_tokens,

185 + query.shape[1] + key.shape[1] + pad,

186 + dtype=dtype,

187 + device=device,

188 + )

189 + packed[:, : query.shape[1]] = query

190 + packed[:, query.shape[1] : query.shape[1] + key.shape[1]] = key

191 + sentinel = packed[:, -pad:].clone()

192 + q_strided = packed[:, : query.shape[1]]

193 + k_strided = packed[:, query.shape[1] : query.shape[1] + key.shape[1]]

194 + mrope_helper_class.forward_cuda(positions, q_strided, k_strided)

195 + torch.testing.assert_close(q_strided, query_cuda, atol=0, rtol=0)

196 + torch.testing.assert_close(k_strided, key_cuda, atol=0, rtol=0)

197 + torch.testing.assert_close(packed[:, -pad:], sentinel, atol=0, rtol=0)

198 +

199 +

200 +@pytest.mark.skipif(

201 + not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests."

202 +)

203 +def test_mrope_row_offset_int64():

204 + # MiniMax M3 ViT worst case: a max-size video item packs ~768k patches,

205 + # and num_tokens * packed-qkv row stride (3840) exceeds 2^31 at 559,241

206 + # rows. Row-base math must be int64 or those rows wrap to negative

207 + # offsets (illegal memory access at startup profiling).

208 + set_random_seed(0)

209 + n_qh, n_kh, hd, rd = 16, 16, 80, 78

210 + mrope_section = [13, 13, 13]

211 + row_stride = (n_qh + n_kh + n_kh) * hd # packed qkv row

212 + num_tokens = 2**31 // row_stride + 8

213 + free, _ = current_platform.mem_get_info()

214 + if free < 2 * num_tokens * row_stride * 2:

215 + pytest.skip("not enough free device memory")

216 +

217 + packed = torch.empty(num_tokens, row_stride, dtype=torch.bfloat16, device=device)

218 + q_fill = torch.randn(n_qh * hd, dtype=torch.bfloat16, device=device)

219 + k_fill = torch.randn(n_kh * hd, dtype=torch.bfloat16, device=device)

220 + packed[:, : n_qh * hd] = q_fill

221 + packed[:, n_qh * hd : (n_qh + n_kh) * hd] = k_fill

222 + half_rd = rd // 2

223 + cos = torch.randn(3, num_tokens, half_rd, dtype=torch.float32, device=device)

224 + sin = torch.randn(3, num_tokens, half_rd, dtype=torch.float32, device=device)

225 + q = packed[:, : n_qh * hd]

226 + k = packed[:, n_qh * hd : (n_qh + n_kh) * hd]

227 + triton_mrope(q, k, cos, sin, mrope_section, hd, rd, False, True)

228 +

229 + t_end = mrope_section[0]

230 + h_end = t_end + mrope_section[1]

231 + for r in (0, num_tokens - 8, num_tokens - 1):

232 + cos_row = torch.cat(

233 + [cos[0, r, :t_end], cos[1, r, t_end:h_end], cos[2, r, h_end:]]

234 + )

235 + sin_row = torch.cat(

236 + [sin[0, r, :t_end], sin[1, r, t_end:h_end], sin[2, r, h_end:]]

237 + )

238 + for fill, out in ((q_fill, q[r]), (k_fill, k[r])):

239 + row = fill.view(-1, hd).float()

240 + r1, r2 = row[:, :half_rd], row[:, half_rd:rd]

241 + exp = torch.cat(

242 + [r1 * cos_row - r2 * sin_row, r2 * cos_row + r1 * sin_row, row[:, rd:]],

243 + dim=1,

244 + )

245 + torch.testing.assert_close(

246 + out, exp.view(-1).to(torch.bfloat16), atol=1e-2, rtol=1.6e-2

247 + )

248 +

249

250 @pytest.mark.skipif(

251 not current_platform.is_cuda_alike(), reason="Skipping CUDA/ROCm only tests."

vllm/model_executor/layers/rotary_embedding/mrope.py

+12 -7

@@ -20,6 +20,8 @@ def _triton_mrope_forward(

20 cos,

21 sin,

22 num_tokens,

23 + q_token_stride,

24 + k_token_stride,

25 n_qh: tl.constexpr,

26 n_kh: tl.constexpr,

27 hd: tl.constexpr,

@@ -39,9 +41,10 @@ def _triton_mrope_forward(

41 # and supports cos and sin cache with shape (3, num_tokens, rotary_dim // 2)

42 # instead of (3, bsz, seq_len, head_dim), also supports interleaved rotary

43 pid = tl.program_id(0)

42 - # locate start address

43 - q_ptr = q_ptr + pid * (n_qh * hd)

44 - k_ptr = k_ptr + pid * (n_kh * hd)

44 + # locate start address (int64: num_tokens * row_stride exceeds 2^31 for

45 + # worst-case ViT forwards, e.g. MiniMax M3 packs ~768k patches/item)

46 + q_ptr = q_ptr + pid.to(tl.int64) * q_token_stride

47 + k_ptr = k_ptr + pid.to(tl.int64) * k_token_stride

49 # ####################################################################

50 # get the cos(mθ_{i...d/2}) and sin(mθ_{i...d/2}) for token position

@@ -186,6 +189,10 @@ def triton_mrope(

189 adjacent (GPT-J) layout.

190

191 """

192 + # The kernel indexes each token row through an explicit row stride and

193 + # requires dense rows.

194 + assert q.dim() == 2 and q.stride(-1) == 1

195 + assert k.dim() == 2 and k.stride(-1) == 1

196 n_row, n_q_head_head_dim = q.shape

197 n_q_head = n_q_head_head_dim // head_size

198 n_kv_head = k.shape[1] // head_size

@@ -193,10 +200,6 @@ def triton_mrope(

200 pad_n_q_head = triton.next_power_of_2(n_q_head)

201 pad_n_kv_head = triton.next_power_of_2(n_kv_head)

202

196 - # ensure tensors passed into the kernel are contiguous.

197 - # It will be no-op if they are already contiguous

198 - q = q.contiguous()

199 - k = k.contiguous()

203 cos = cos.contiguous()

204 sin = sin.contiguous()

205

@@ -211,6 +214,8 @@ def triton_mrope(

214 cos,

215 sin,

216 n_row,

217 + q.stride(0),

218 + k.stride(0),

219 n_q_head,

220 n_kv_head,

221 head_size,

vllm/model_executor/layers/rotary_embedding/mrope_vit_setup.py new

+153

@@ -0,0 +1,153 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import torch

5 +

6 +from vllm.triton_utils import tl, triton

7 +from vllm.utils.torch_utils import async_tensor_h2d

8 +

9 +

10 +# num_tokens/num_grids vary per request (incl. the ==1 class) and grids_ptr

11 +# is an offset slice of the packed H2D buffer, so its 16B alignment varies

12 +# with grid count; without these the kernel recompiles mid-serving.

13 +@triton.jit(

14 + do_not_specialize=["num_tokens", "num_grids"],

15 + do_not_specialize_on_alignment=["grids_ptr"],

16 +)

17 +def _vit_mrope_setup_kernel(

18 + inv_freq_t_ptr,

19 + inv_freq_h_ptr,

20 + inv_freq_w_ptr,

21 + cos_out_ptr,

22 + sin_out_ptr,

23 + cu_seqlens_ptr, # (G + 1,) int32 token boundaries per grid

24 + grids_ptr, # (G, 3) int32: t, h, w per grid

25 + num_tokens,

26 + num_grids,

27 + half_t: tl.constexpr,

28 + half_h: tl.constexpr,

29 + half_w: tl.constexpr,

30 + spatial_merge: tl.constexpr,

31 + BLOCK_N: tl.constexpr,

32 + BLOCK_D: tl.constexpr,

33 +):

34 + """Build triton_mrope's (3, N, half_rot) cos/sin t/h/w planes in one

35 + launch: compute each token's (t, h, w) position arithmetically and the

36 + rotation in fp32 from the per-axis inv_freq. Only each axis's mrope

37 + section columns are written; the rest is left uninitialized (the

38 + consumer's per-axis masked loads read only the section columns)."""

39 + pid = tl.program_id(0)

40 + tok = pid * BLOCK_N + tl.arange(0, BLOCK_N)

41 + tok_mask = tok < num_tokens

42 +

43 + # Grid (segment) index per token; cu_seqlens is nondecreasing.

44 + seg = tl.zeros([BLOCK_N], dtype=tl.int32)

45 + for g in range(num_grids):

46 + end = tl.load(cu_seqlens_ptr + g + 1)

47 + seg += (tok >= end).to(tl.int32)

48 +

49 + seg_start = tl.load(cu_seqlens_ptr + seg, mask=tok_mask, other=0)

50 + grid_h = tl.load(grids_ptr + seg * 3 + 1, mask=tok_mask, other=1)

51 + grid_w = tl.load(grids_ptr + seg * 3 + 2, mask=tok_mask, other=1)

52 +

53 + local = tok - seg_start

54 + rem = local % (grid_h * grid_w)

55 + tpos = local // (grid_h * grid_w)

56 + # Within a frame tokens are ordered by spatial-merge block (bh, bw, mh,

57 + # mw), matching the reference ViT position-id permute.

58 + bh = rem // (spatial_merge * grid_w)

59 + r = rem % (spatial_merge * grid_w)

60 + bw = r // (spatial_merge * spatial_merge)

61 + r2 = r % (spatial_merge * spatial_merge)

62 + hpos = bh * spatial_merge + r2 // spatial_merge

63 + wpos = bw * spatial_merge + r2 % spatial_merge

64 +

65 + d = tl.arange(0, BLOCK_D)

66 + half_rot = half_t + half_h + half_w

67 + row = (tok * half_rot)[:, None] + d[None, :]

68 + plane = num_tokens * half_rot

69 +

70 + # Per-axis compute+store to cap live registers; the fp32 trig sequence

71 + # wants few elements per thread (hence small BLOCK_N, many warps).

72 + in_t = d < half_t

73 + inv_t = tl.load(inv_freq_t_ptr + d, mask=in_t, other=0.0)

74 + freq = tpos.to(tl.float32)[:, None] * inv_t[None, :]

75 + mask = tok_mask[:, None] & in_t[None, :]

76 + tl.store(cos_out_ptr + row, tl.cos(freq), mask=mask)

77 + tl.store(sin_out_ptr + row, tl.sin(freq), mask=mask)

78 +

79 + in_h = (d >= half_t) & (d < half_t + half_h)

80 + inv_h = tl.load(inv_freq_h_ptr + d - half_t, mask=in_h, other=0.0)

81 + freq = hpos.to(tl.float32)[:, None] * inv_h[None, :]

82 + mask = tok_mask[:, None] & in_h[None, :]

83 + tl.store(cos_out_ptr + plane + row, tl.cos(freq), mask=mask)

84 + tl.store(sin_out_ptr + plane + row, tl.sin(freq), mask=mask)

85 +

86 + in_w = (d >= half_t + half_h) & (d < half_rot)

87 + inv_w = tl.load(inv_freq_w_ptr + d - half_t - half_h, mask=in_w, other=0.0)

88 + freq = wpos.to(tl.float32)[:, None] * inv_w[None, :]

89 + mask = tok_mask[:, None] & in_w[None, :]

90 + tl.store(cos_out_ptr + 2 * plane + row, tl.cos(freq), mask=mask)

91 + tl.store(sin_out_ptr + 2 * plane + row, tl.sin(freq), mask=mask)

92 +

93 +

94 +def vit_mrope_setup(

95 + inv_freq_t: torch.Tensor,

96 + inv_freq_h: torch.Tensor,

97 + inv_freq_w: torch.Tensor,

98 + grid_thw: list[list[int]],

99 + spatial_merge_size: int,

100 +) -> tuple[torch.Tensor, torch.Tensor]:

101 + """(3, total_N, half_rot_dim) fp32 cos/sin t/h/w planes for triton_mrope.

102 +

103 + One fused kernel: positions are computed arithmetically from the grids

104 + and the rotation comes from the per-axis inv_freq, so there is no

105 + cos/sin cache to size or grow.

106 +

107 + Args:

108 + inv_freq_t: t-axis inverse frequencies, (t_dim // 2,).

109 + inv_freq_h: h-axis inverse frequencies, (h_dim // 2,).

110 + inv_freq_w: w-axis inverse frequencies, (w_dim // 2,).

111 + grid_thw: per-grid [t, h, w] patch counts (host metadata).

112 + spatial_merge_size: spatial merge factor of the ViT.

113 +

114 + """

115 + device = inv_freq_t.device

116 + cu = [0]

117 + for t, h, w in grid_thw:

118 + cu.append(cu[-1] + t * h * w)

119 + n = cu[-1]

120 + g = len(grid_thw)

121 + # grid_thw is per-request host metadata, so one H2D is unavoidable;

122 + # pack cu_seqlens and grids into a single pinned async copy.

123 + buf = async_tensor_h2d(

124 + cu + [d for grid in grid_thw for d in grid], device=device, dtype=torch.int32

125 + )

126 + cu_seqlens = buf[: g + 1]

127 + grids = buf[g + 1 :].view(g, 3)

128 +

129 + half_t = inv_freq_t.numel()

130 + half_h = inv_freq_h.numel()

131 + half_w = inv_freq_w.numel()

132 + half_rot = half_t + half_h + half_w

133 + cos = torch.empty(3, n, half_rot, device=device, dtype=torch.float32)

134 + sin = torch.empty(3, n, half_rot, device=device, dtype=torch.float32)

135 + _vit_mrope_setup_kernel[(triton.cdiv(n, 32),)](

136 + inv_freq_t,

137 + inv_freq_h,

138 + inv_freq_w,

139 + cos,

140 + sin,

141 + cu_seqlens,

142 + grids,

143 + n,

144 + g,

145 + half_t=half_t,

146 + half_h=half_h,

147 + half_w=half_w,

148 + spatial_merge=spatial_merge_size,

149 + BLOCK_N=32,

150 + BLOCK_D=triton.next_power_of_2(half_rot),

151 + num_warps=8,

152 + )

153 + return cos, sin

vllm/models/minimax_m3/common/vision_tower.py

+42 -171

@@ -21,20 +21,16 @@ from vllm.model_executor.layers.linear import (

21 RowParallelLinear,

22 )

23 from vllm.model_executor.layers.quantization import QuantizationConfig

24 -from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb

24 +from vllm.model_executor.layers.rotary_embedding.mrope import triton_mrope

25 +from vllm.model_executor.layers.rotary_embedding.mrope_vit_setup import (

26 + vit_mrope_setup,

27 +)

28 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

29 from vllm.model_executor.models.utils import maybe_prefix

30 from vllm.model_executor.models.vision import (

31 get_vit_attn_backend,

32 is_vit_use_data_parallel,

33 )

31 -from vllm.platforms import current_platform

32 -

33 -# ROCm caps a kernel-launch gridDim.y at 65536. The HIP flash-attn Triton

34 -# rotary kernel launches grid.y = cdiv(seqlen, BLOCK_M), so it fails with

35 -# hipErrorInvalidValue once cdiv(seqlen, BLOCK_M) > 65536. Used below to decide

36 -# when RoPE must be applied per video segment instead of in one launch.

37 -_HIP_MAX_GRID_DIM_Y = 65536

36 class MiniMaxVLPatchEmbed(nn.Module):

@@ -125,56 +121,13 @@ class MiniMaxVLAttention(nn.Module):

121 head_size=self.head_dim,

122 prefix=f"{prefix}.attn",

123 )

128 - # ApplyRotaryEmb handles the internal cos/sin repeat and partial

129 - # rotation (ro_dim = half_rot_dim * 2 < head_dim for MiniMax).

130 - # enable_fp32_compute=True runs the rotation in fp32 (q/k upcast,

131 - # fp32 cos/sin), matching the reference ``_minimax_rope_applier``.

132 - self.apply_rotary_emb = ApplyRotaryEmb(

133 - enforce_enable=True, enable_fp32_compute=True

134 - )

135 -

136 - def _apply_rotary_emb(

137 - self,

138 - qk_reshaped: torch.Tensor,

139 - rotary_cos: torch.Tensor,

140 - rotary_sin: torch.Tensor,

141 - seq_len: int,

142 - rotary_segment_lengths: list[int] | None,

143 - ) -> torch.Tensor:

144 - # Default fast path (all NVIDIA inputs, and ROCm short clips/images):

145 - # a single rotary kernel launch. ``rotary_segment_lengths`` is only

146 - # populated on ROCm (see ``MiniMaxVLVisionTransformer.forward``), so

147 - # the per-segment path below is ROCm-only and never touches the

148 - # NVIDIA/CUDA code path.

149 - if not current_platform.is_rocm() or rotary_segment_lengths is None:

150 - return self.apply_rotary_emb(qk_reshaped, rotary_cos, rotary_sin)

151 -

152 - # ROCm only: the HIP flash-attn Triton rotary kernel fails with

153 - # hipErrorInvalidValue once grid.y = cdiv(seqlen, BLOCK_M) exceeds

154 - # _HIP_MAX_GRID_DIM_Y (65536). BLOCK_M is 8 for rotary_dim <= 128

155 - # (MiniMax-M3 vision: rotary_dim=78), giving a hard limit of

156 - # 65536 * BLOCK_M tokens — measured exactly as 524288 OK / 524289 fail.

157 - # Only long videos cross it; since vision_segment_max_frames caps each

158 - # segment at a few frames (<< limit), applying RoPE per segment keeps

159 - # every sub-call in range. Splitting on segment boundaries is

160 - # mathematically exact because rotary_cos/sin are precomputed per token.

161 - # Images and short clips stay on the single-kernel fast path above.

162 - rotary_dim = rotary_cos.shape[-1] * 2

163 - block_m = 8 if rotary_dim <= 128 else 4

164 - hip_rotary_max_seqlen = _HIP_MAX_GRID_DIM_Y * block_m

165 - if seq_len <= hip_rotary_max_seqlen or len(rotary_segment_lengths) <= 1:

166 - return self.apply_rotary_emb(qk_reshaped, rotary_cos, rotary_sin)

167 -

168 - qk_segments = qk_reshaped.split(rotary_segment_lengths, dim=1)

169 - cos_segments = rotary_cos.split(rotary_segment_lengths, dim=0)

170 - sin_segments = rotary_sin.split(rotary_segment_lengths, dim=0)

171 - return torch.cat(

172 - [

173 - self.apply_rotary_emb(qk_s, cos_s, sin_s)

174 - for qk_s, cos_s, sin_s in zip(qk_segments, cos_segments, sin_segments)

175 - ],

176 - dim=1,

177 - )

124 + # Partial 3D RoPE geometry: rot_dim is split evenly across t/h/w

125 + # (same formula as MiniMaxVLVisionTransformer; rot_dim may be <

126 + # head_dim). mrope_section is in half-dim units for triton_mrope.

127 + rope_dims = 2 * (self.head_dim // 2)

128 + axis_dim = 2 * ((rope_dims // 3) // 2)

129 + self.rotary_dim = 3 * axis_dim

130 + self.mrope_section = [axis_dim // 2] * 3

131

132 def forward(

133 self,

@@ -183,33 +136,34 @@ class MiniMaxVLAttention(nn.Module):

136 rotary_cos: torch.Tensor,

137 rotary_sin: torch.Tensor,

138 max_seqlen: torch.Tensor,

186 - rotary_segment_lengths: list[int] | None = None,

139 sequence_lengths: torch.Tensor | None = None,

140 ) -> torch.Tensor:

141 # x: (N, 1, embed_dim) [seq=N, batch=1, chan=embed_dim]

142 x_qkv, _ = self.qkv_proj(x) # (N, 1, 3 * heads_per_part * head_dim)

143 seq_len, batch_size, _ = x_qkv.shape

144

193 - # Rearrange to (b=1, N, 3, heads, head_dim) — same as Qwen2_5_VisionAttention

194 - qkv = rearrange(

195 - x_qkv,

196 - "s b (three head d) -> b s three head d",

197 - three=3,

198 - head=self.num_heads_per_partition,

199 - )

200 - qk, v = qkv[:, :, :2], qkv[:, :, 2] # (b,N,2,h,d) and (b,N,h,d)

201 -

202 - # Stack q/k → (2*b, N, heads, head_dim) for joint RoPE application.

203 - # rotary_cos/sin: (N, half_rot_dim) — ApplyRotaryEmb expands internally

204 - # and rotates only the first 2*half_rot_dim dims, passing the rest through.

205 - qk_reshaped = rearrange(qk, "b s two h d -> (two b) s h d", two=2).contiguous()

206 - qk_rotated = self._apply_rotary_emb(

207 - qk_reshaped, rotary_cos, rotary_sin, seq_len, rotary_segment_lengths

208 - )

209 - qk_rotated = qk_rotated.view(

210 - 2, batch_size, seq_len, self.num_heads_per_partition, self.head_dim

145 + qkv = x_qkv.view(seq_len, 3, self.num_heads_per_partition, self.head_dim)

146 + v = qkv[:, 2].unsqueeze(0) # (b=1, N, heads, head_dim) strided view

147 +

148 + # In-place partial 3D RoPE on strided q/k views of the packed qkv

149 + # tensor: one kernel launch, no repack/upcast copies. fp32 cos/sin

150 + # make the kernel compute in fp32, matching the reference

151 + # ``_minimax_rope_applier`` precision.

152 + q = qkv[:, 0].reshape(seq_len, -1)

153 + k = qkv[:, 1].reshape(seq_len, -1)

154 + triton_mrope(

155 + q,

156 + k,

157 + rotary_cos,

158 + rotary_sin,

159 + self.mrope_section,

160 + self.head_dim,

161 + self.rotary_dim,

162 + mrope_interleaved=False,

163 + is_neox_style=True,

164 )

212 - q, k = qk_rotated.unbind(dim=0) # each (b=1, N, heads, head_dim)

165 + q = q.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)

166 + k = k.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)

167

168 # Flash attention → (b, N, heads, head_dim)

169 context = self.attn(

@@ -272,7 +226,6 @@ class MiniMaxVLEncoderLayer(nn.Module):

226 rotary_cos: torch.Tensor,

227 rotary_sin: torch.Tensor,

228 max_seqlen: torch.Tensor,

275 - rotary_segment_lengths: list[int] | None = None,

229 sequence_lengths: torch.Tensor | None = None,

230 ) -> torch.Tensor:

231 # x: (N, 1, hidden_size)

@@ -282,7 +235,6 @@ class MiniMaxVLEncoderLayer(nn.Module):

235 rotary_cos,

236 rotary_sin,

237 max_seqlen,

285 - rotary_segment_lengths,

238 sequence_lengths,

239 )

240 residual = x

@@ -324,7 +276,6 @@ class MiniMaxVLEncoder(nn.Module):

276 rotary_cos: torch.Tensor,

277 rotary_sin: torch.Tensor,

278 max_seqlen: torch.Tensor,

327 - rotary_segment_lengths: list[int] | None = None,

279 sequence_lengths: torch.Tensor | None = None,

280 ) -> torch.Tensor:

281 for layer in self.layers:

@@ -334,7 +285,6 @@ class MiniMaxVLEncoder(nn.Module):

285 rotary_cos,

286 rotary_sin,

287 max_seqlen,

337 - rotary_segment_lengths,

288 sequence_lengths,

289 )

290 return x

@@ -427,82 +377,6 @@ class MiniMaxVLVisionTransformer(nn.Module):

377 # out_hidden_size needed by run_dp_sharded_mrope_vision_model

378 self.out_hidden_size = embed_dim

379

430 - # ── RoPE helpers ─────────────────────────────────────────────────────

431 -

432 - def _get_3d_rope_embed(

433 - self, grid_t: int, grid_h: int, grid_w: int, spatial_merge_size: int

434 - ) -> torch.Tensor:

435 - """Compute 3D RoPE frequencies for a single (T, H, W) grid.

436 -

437 - Returns (T*H*W, half_rot_dim) on the same device as inv_freq buffers.

438 - Mirrors the reference ``_get_3d_rope_embed`` exactly.

439 - """

440 - tokens_per_frame = grid_h * grid_w

441 -

442 - tpos_ids = (

443 - torch.arange(grid_t, device=self.inv_freq_t.device)

444 - .unsqueeze(1)

445 - .expand(-1, tokens_per_frame)

446 - .flatten()

447 - )

448 -

449 - hpos_ids = (

450 - torch.arange(grid_h, device=self.inv_freq_h.device)

451 - .unsqueeze(1)

452 - .expand(-1, grid_w)

453 - .reshape(

454 - grid_h // spatial_merge_size,

455 - spatial_merge_size,

456 - grid_w // spatial_merge_size,

457 - spatial_merge_size,

458 - )

459 - .permute(0, 2, 1, 3)

460 - .unsqueeze(0)

461 - .expand(grid_t, -1, -1, -1, -1)

462 - .flatten()

463 - )

464 - wpos_ids = (

465 - torch.arange(grid_w, device=self.inv_freq_w.device)

466 - .unsqueeze(0)

467 - .expand(grid_h, -1)

468 - .reshape(

469 - grid_h // spatial_merge_size,

470 - spatial_merge_size,

471 - grid_w // spatial_merge_size,

472 - spatial_merge_size,

473 - )

474 - .permute(0, 2, 1, 3)

475 - .unsqueeze(0)

476 - .expand(grid_t, -1, -1, -1, -1)

477 - .flatten()

478 - )

479 -

480 - max_t = max(grid_t, 1)

481 - max_hw = max(grid_h, grid_w)

482 -

483 - seq_t = torch.arange(

484 - max_t, device=self.inv_freq_t.device, dtype=self.inv_freq_t.dtype

485 - )

486 - seq_hw = torch.arange(

487 - max_hw, device=self.inv_freq_h.device, dtype=self.inv_freq_h.dtype

488 - )

489 -

490 - freqs_t = torch.outer(seq_t, self.inv_freq_t) # (max_t, t_dim/2)

491 - freqs_h = torch.outer(seq_hw, self.inv_freq_h) # (max_hw, h_dim/2)

492 - freqs_w = torch.outer(seq_hw, self.inv_freq_w) # (max_hw, w_dim/2)

493 -

494 - return torch.cat(

495 - [freqs_t[tpos_ids], freqs_h[hpos_ids], freqs_w[wpos_ids]], dim=-1

496 - ) # (T*H*W, half_rot_dim)

497 -

498 - def _get_rope_embed_3d(

499 - self, grid_thw: list[list[int]], spatial_merge_size: int

500 - ) -> torch.Tensor:

501 - embeds = [

502 - self._get_3d_rope_embed(t, h, w, spatial_merge_size) for t, h, w in grid_thw

503 - ]

504 - return torch.cat(embeds, dim=0) # (total_N, half_rot_dim)

505 -

380 # ── Frame-limit helper (mirrors the reference) ───────────────────────

381

382 def _apply_max_frames_limit(self, grid_thw: list[list[int]]) -> list[list[int]]:

@@ -557,21 +431,19 @@ class MiniMaxVLVisionTransformer(nn.Module):

431 hidden.device,

432 )

433

560 - # 3D RoPE: (total_N, half_rot_dim); ApplyRotaryEmb expands internally

561 - freqs = self._get_rope_embed_3d(limited, self.spatial_merge_size)

562 - freqs = freqs.to(device=hidden.device)

563 - # Keep cos/sin in fp32; ApplyRotaryEmb(enable_fp32_compute=True) runs the

564 - # rotation in fp32 to match the reference precision.

565 - rotary_cos, rotary_sin = freqs.cos(), freqs.sin()

434 + # 3D RoPE cos/sin: (3, total_N, half_rot_dim) fp32 t/h/w planes for

435 + # triton_mrope from one fused setup kernel. Kept fp32 so the kernel

436 + # computes the rotation in fp32, matching the reference.

437 + rotary_cos, rotary_sin = vit_mrope_setup(

438 + self.inv_freq_t,

439 + self.inv_freq_h,

440 + self.inv_freq_w,

441 + limited,

442 + self.spatial_merge_size,

443 + )

444

445 # Encoder expects (N, 1, hidden_size) — add batch dim

446 hidden = hidden.unsqueeze(1)

569 - # On ROCm, the flash_attn Triton rotary kernel can fail with

570 - # hipErrorInvalidValue when seqlen is very large, e.g. 192k video

571 - # tokens; pass per-segment lengths so RoPE can be applied in chunks.

572 - # On other platforms leave it None -> single-kernel fast path, so the

573 - # NVIDIA/CUDA code path is unchanged.

574 - rotary_segment_lengths = lens if current_platform.is_rocm() else None

447

448 hidden = self.encoder(

449 hidden,

@@ -579,7 +451,6 @@ class MiniMaxVLVisionTransformer(nn.Module):

451 rotary_cos,

452 rotary_sin,

453 max_seqlen,

582 - rotary_segment_lengths,

454 sequence_lengths=sequence_lengths,

455 )

456 hidden = hidden.squeeze(1) # back to (total_N, hidden_size)