@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9ef98d527e

Fork 0

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 15, 16:00 UTC

[Model][MiniMaxText01] Support MiniMaxText01 model inference (#13454)

Signed-off-by: qscqesze <475517977@qq.com> Co-authored-by: qingjun <qingjun@minimaxi.com> Co-authored-by: qscqesze <475517977@qq.com>

Gerald committed Apr 2, 2025 at 04:23 UTC 9ef98d527ee1b943afd39d9dcad0ffe089ab241f

11 files changed +2440 -130

docs/source/models/supported_models.md

@@ -503,6 +503,11 @@ See [this page](#generative-models) for more information on how to use generativ

503 * `xverse/XVERSE-7B-Chat`, `xverse/XVERSE-13B-Chat`, `xverse/XVERSE-65B-Chat`, etc.

504 * ✅︎

505 * ✅︎

506 +- * `MiniMaxText01ForCausalLM`

507 + * MiniMax-Text

508 + * `MiniMaxAI/MiniMax-Text-01`, etc.

509 + *

510 + * ✅︎

511 - * `Zamba2ForCausalLM`

512 * Zamba2

513 * `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc.

tests/kernels/test_lightning_attn.py new

+286

@@ -0,0 +1,286 @@

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

2 +

3 +import pytest

4 +import torch

5 +

6 +from vllm.model_executor.layers.lightning_attn import (

7 + linear_decode_forward_triton)

8 +from vllm.platforms import current_platform

9 +

10 +NUM_HEADS = [4, 8]

11 +HEAD_SIZES = [64]

12 +BATCH_SIZES = [1, 2]

13 +SEQ_LENGTHS = [16]

14 +DTYPES = [torch.float32]

15 +

16 +

17 +def reference_lightning_attention(q, k, v, ed, block_size, kv_history):

18 + """Reference implementation of lightning attention core algorithm

19 +

20 + The difference from the main implementation is that this processes

21 + each step sequentially, instead of using parallelized triton kernels

22 + """

23 + B, H, S, D = q.shape

24 + E = v.shape[-1]

25 + dtype = q.dtype

26 + output = torch.zeros((B, H, S, E), dtype=dtype, device=q.device)

27 +

28 + # Use clone() to ensure an independent copy

29 + if kv_history is None:

30 + kv_cache = torch.zeros((B, H, D, E), dtype=dtype, device=q.device)

31 + else:

32 + kv_cache = kv_history.clone()

33 +

34 + # More efficient implementation

35 + # Convert decay factors to matrix form

36 + if ed.dim() == 1:

37 + decay = torch.exp(-ed).view(1, -1, 1, 1)

38 + else:

39 + decay = torch.exp(-ed)

40 +

41 + for b in range(B):

42 + for step in range(S):

43 + # Process all heads at once for this position

44 + q_bs = q[b, :, step] # [H, D]

45 + k_bs = k[b, :, step] # [H, D]

46 + v_bs = v[b, :, step] # [H, E]

47 +

48 + # Calculate KV outer products for all heads

49 + for h in range(H):

50 + # Calculate KV outer product

51 + kv_outer = torch.outer(k_bs[h], v_bs[h])

52 +

53 + # Update KV cache with decay

54 + # Note: Using the same order as in the Triton kernel

55 + kv_cache[b, h] = decay[0, h, 0, 0] * kv_cache[b, h] + kv_outer

56 +

57 + # Calculate attention output

58 + output[b, h, step] = torch.matmul(q_bs[h], kv_cache[b, h])

59 +

60 + # Match the shape returned by the actual implementation

61 + # The actual implementation returns a tensor of shape [B, H, 2, D, E]

62 + # where dimension 2 contains both KV and KV history

63 + kv_reshaped = kv_cache.unsqueeze(2) # [B, H, 1, D, E]

64 + final_kv_cache = torch.cat([kv_reshaped, kv_reshaped],

65 + dim=2) # [B, H, 2, D, E]

66 +

67 + return output, final_kv_cache

68 +

69 +

70 +def reference_linear_decode(q, k, v, kv_caches, slope_rate, slot_idx):

71 + """Reference implementation: linear attention decode function"""

72 + B, H, _, D = q.shape

73 + output = torch.zeros(B, H * D, dtype=q.dtype, device=q.device)

74 +

75 + # Calculate decay factors once (more efficient)

76 + decay = torch.exp(-slope_rate).view(-1, 1, 1) # [H, 1, 1]

77 +

78 + # Process each batch

79 + for b in range(B):

80 + slot_id = slot_idx[b].item()

81 +

82 + # Skip padding positions

83 + if slot_id == -1:

84 + continue

85 +

86 + # Process all heads at once for this batch

87 + q_b = q[b, :, 0] # [H, D]

88 + k_b = k[b, :, 0] # [H, D]

89 + v_b = v[b, :, 0] # [H, D]

90 +

91 + # Process each attention head

92 + for h in range(H):

93 + # Get current query, key and value

94 + q_bh = q_b[h]

95 + k_bh = k_b[h]

96 + v_bh = v_b[h]

97 +

98 + # Get cache

99 + kv_cache_old = kv_caches[b, h]

100 +

101 + # Calculate new key-value outer product

102 + kv_outer = torch.outer(k_bh, v_bh)

103 +

104 + # Apply decay and update cache

105 + kv_new = kv_outer + decay[h, 0, 0] * kv_cache_old

106 +

107 + # Calculate output

108 + out_h = torch.matmul(q_bh, kv_new)

109 +

110 + # Update output and cache

111 + output[b, h * D:(h + 1) * D] = out_h

112 + kv_caches[b, h] = kv_new

113 +

114 + return output

115 +

116 +

117 +@pytest.mark.parametrize("batch_size", BATCH_SIZES)

118 +@pytest.mark.parametrize("num_heads", NUM_HEADS)

119 +@pytest.mark.parametrize("head_size", HEAD_SIZES)

120 +@pytest.mark.parametrize("dtype", DTYPES)

121 +@torch.inference_mode()

122 +def test_linear_decode_forward_triton(

123 + batch_size: int,

124 + num_heads: int,

125 + head_size: int,

126 + dtype: torch.dtype,

127 +):

128 + torch.set_default_device("cuda")

129 + torch.manual_seed(42)

130 + torch.cuda.manual_seed_all(42)

131 + current_platform.seed_everything(42)

132 + base = 0.01

133 + q = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)

134 + k = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)

135 + v = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)

136 +

137 + kv_caches = base * torch.randn(batch_size,

138 + num_heads,

139 + head_size,

140 + head_size,

141 + dtype=dtype,

142 + device="cuda")

143 +

144 + kv_caches_copy = kv_caches.clone()

145 +

146 + slope_rate = torch.zeros(num_heads, device="cuda")

147 + for h in range(num_heads):

148 + slope_rate[h] = 0.1 * (h + 1)

149 +

150 + slot_idx = torch.arange(batch_size, device="cuda")

151 +

152 + triton_output = linear_decode_forward_triton(q, k, v, kv_caches,

153 + slope_rate, slot_idx)

154 +

155 + reference_output = reference_linear_decode(q, k, v, kv_caches_copy,

156 + slope_rate, slot_idx)

157 + torch.testing.assert_close(triton_output,

158 + reference_output,

159 + rtol=1e-1,

160 + atol=1e-1)

161 + torch.testing.assert_close(kv_caches, kv_caches_copy, rtol=1e-1, atol=1e-1)

162 +

163 + assert triton_output.shape == (batch_size, num_heads * head_size)

164 +

165 +

166 +@pytest.mark.parametrize("num_heads", NUM_HEADS)

167 +@pytest.mark.parametrize("head_size", HEAD_SIZES)

168 +@pytest.mark.parametrize("dtype", DTYPES)

169 +@torch.inference_mode()

170 +def test_linear_decode_forward_triton_with_padding(

171 + num_heads: int,

172 + head_size: int,

173 + dtype: torch.dtype,

174 +):

175 + torch.set_default_device("cuda")

176 + torch.manual_seed(42)

177 + torch.cuda.manual_seed_all(42)

178 + current_platform.seed_everything(42)

179 +

180 + batch_size = 4

181 + base = 0.01

182 + q = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)

183 + k = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)

184 + v = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)

185 +

186 + kv_caches = base * torch.randn(batch_size,

187 + num_heads,

188 + head_size,

189 + head_size,

190 + dtype=dtype,

191 + device="cuda")

192 +

193 + kv_caches_copy = kv_caches.clone()

194 +

195 + slope_rate = torch.zeros(num_heads, device="cuda")

196 + for h in range(num_heads):

197 + slope_rate[h] = 0.1 * (h + 1)

198 +

199 + slot_idx = torch.tensor([0, 1, -1, 2], device="cuda")

200 +

201 + triton_output = linear_decode_forward_triton(q, k, v, kv_caches,

202 + slope_rate, slot_idx)

203 +

204 + reference_output = reference_linear_decode(q, k, v, kv_caches_copy,

205 + slope_rate, slot_idx)

206 +

207 + padding_mask = (slot_idx

208 + != -1).unsqueeze(1).expand(-1, num_heads * head_size)

209 +

210 + triton_masked = triton_output[padding_mask]

211 + reference_masked = reference_output[padding_mask]

212 +

213 + atol, rtol = 1.5e-1, 1.5e-1

214 +

215 + valid_indices = slot_idx != -1

216 +

217 + for i in range(batch_size):

218 + if valid_indices[i] > 0:

219 + torch.testing.assert_close(kv_caches[i],

220 + kv_caches_copy[i],

221 + rtol=rtol,

222 + atol=atol)

223 +

224 + torch.testing.assert_close(triton_masked,

225 + reference_masked,

226 + rtol=rtol,

227 + atol=atol)

228 +

229 + assert triton_output.shape == (batch_size, num_heads * head_size)

230 +

231 +

232 +@pytest.mark.parametrize("batch_size", BATCH_SIZES)

233 +@pytest.mark.parametrize("num_heads", NUM_HEADS)

234 +@pytest.mark.parametrize("head_size", HEAD_SIZES)

235 +@pytest.mark.parametrize("seq_len", SEQ_LENGTHS)

236 +@pytest.mark.parametrize("dtype", DTYPES)

237 +@torch.inference_mode()

238 +def test_lightning_attention_reference(

239 + batch_size: int,

240 + num_heads: int,

241 + head_size: int,

242 + seq_len: int,

243 + dtype: torch.dtype,

244 +):

245 + torch.set_default_device("cuda")

246 + torch.manual_seed(42)

247 + torch.cuda.manual_seed_all(42)

248 + current_platform.seed_everything(42)

249 +

250 + base = 0.01

251 + q = base * torch.randn(

252 + batch_size, num_heads, seq_len, head_size, dtype=dtype)

253 + k = base * torch.randn(

254 + batch_size, num_heads, seq_len, head_size, dtype=dtype)

255 + v = base * torch.randn(

256 + batch_size, num_heads, seq_len, head_size, dtype=dtype)

257 +

258 + ed = torch.zeros(num_heads, device="cuda")

259 + for h in range(num_heads):

260 + ed[h] = 0.1 * (h + 1)

261 +

262 + kv_history = base * torch.randn(batch_size,

263 + num_heads,

264 + head_size,

265 + head_size,

266 + dtype=dtype,

267 + device="cuda")

268 +

269 + kv_history_clone = kv_history.clone()

270 +

271 + ref_output, ref_kv_cache = reference_lightning_attention(

272 + q, k, v, ed, 256, kv_history)

273 +

274 + from vllm.model_executor.layers.lightning_attn import lightning_attention

275 + actual_output, actual_kv_cache = lightning_attention(

276 + q, k, v, ed, 256, kv_history_clone)

277 +

278 + atol, rtol = 1.5e-1, 1.5e-1

279 + torch.testing.assert_close(ref_output, actual_output, rtol=rtol, atol=atol)

280 + torch.testing.assert_close(ref_kv_cache,

281 + actual_kv_cache,

282 + rtol=rtol,

283 + atol=atol)

284 +

285 + assert ref_output.shape == (batch_size, num_heads, seq_len, head_size)

286 + assert ref_kv_cache.shape == actual_kv_cache.shape

tests/models/registry.py

@@ -176,6 +176,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

176 trust_remote_code=True),

177 "MiniCPM3ForCausalLM": _HfExamplesInfo("openbmb/MiniCPM3-4B",

178 trust_remote_code=True),

179 + "MiniMaxText01ForCausalLM": _HfExamplesInfo("MiniMaxAI/MiniMax-Text-01",

180 + trust_remote_code=True),

181 "MistralForCausalLM": _HfExamplesInfo("mistralai/Mistral-7B-Instruct-v0.1"),

182 "MixtralForCausalLM": _HfExamplesInfo("mistralai/Mixtral-8x7B-Instruct-v0.1"), # noqa: E501

183 "QuantMixtralForCausalLM": _HfExamplesInfo("mistral-community/Mixtral-8x22B-v0.1-AWQ"), # noqa: E501

vllm/config.py

+25 -17

@@ -971,26 +971,34 @@ class ModelConfig:

971 return sum(not bc.attention.no_op

972 for bc in block_configs[start:end])

973 else:

974 - # Hybrid model

974 + # Hybrid model Jamba

975 layers_block_type_value = getattr(self.hf_config,

976 "layers_block_type", None)

977 - if layers_block_type_value is None:

978 - raise ValueError("The model is an hybrid without a "

979 - "layers_block_type in the hf_config, "

980 - "cannot determine the num of "

981 - f"{block_type.value} layers")

982 -

983 - if hasattr(self.hf_text_config,

984 - "model_type") and (self.hf_text_config.model_type

985 - == "zamba2"):

986 - if attn_block_type:

987 - return sum(t == "hybrid"

988 - for t in layers_block_type_value[start:end])

989 - else:

990 - return self.get_num_layers(parallel_config)

977 + if layers_block_type_value is not None:

978 + if hasattr(self.hf_text_config,

979 + "model_type") and (self.hf_text_config.model_type

980 + == "zamba2"):

981 + if attn_block_type:

982 + return sum(t == "hybrid"

983 + for t in layers_block_type_value[start:end])

984 + else:

985 + return self.get_num_layers(parallel_config)

986 + return sum(t == block_type.value

987 + for t in layers_block_type_value[start:end])

988 +

989 + # Hybrid model Minimax

990 + attn_type_list = getattr(self.hf_config, "attn_type_list", None)

991 + if attn_type_list:

992 + return sum(t == 1 for t in attn_type_list[start:end])

993 +

994 + if layers_block_type_value is None and attn_type_list is None:

995 + raise ValueError(

996 + "The model is an hybrid without a"

997 + "layers_block_type or an attn_type_list in the hf_config,"

998 + "cannot determine the num of "

999 + f"{block_type.value} layers")

1000

992 - return sum(t == block_type.value

993 - for t in layers_block_type_value[start:end])

1001 + return sum(t == 1 for t in attn_type_list[start:end])

1002

1003 def get_multimodal_config(self) -> "MultiModalConfig":

1004 """

vllm/engine/async_llm_engine.py

+5 -2

@@ -303,8 +303,11 @@ class _AsyncLLMEngine(LLMEngine):

303 ctx.seq_group_metadata_list = seq_group_metadata_list

304 ctx.scheduler_outputs = scheduler_outputs

305

306 - finished_requests_ids = self.scheduler[

307 - virtual_engine].get_and_reset_finished_requests_ids()

306 + if not scheduler_outputs.is_empty():

307 + # this will cause mamba_cache/minimax_cache failed

308 + # to release finished_requests_ids of the last steps

309 + finished_requests_ids = self.scheduler[

310 + virtual_engine].get_and_reset_finished_requests_ids()

311

312 # Maybe switch from async mode to sync mode

313 if not allow_async_output_proc and len(ctx.output_queue) > 0:

vllm/model_executor/layers/lightning_attn.py new

+651

@@ -0,0 +1,651 @@

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

2 +import torch

3 +import triton

4 +import triton.language as tl

5 +from einops import rearrange

6 +

7 +

8 +@triton.jit

9 +def _fwd_diag_kernel(Q, K, V, Out, S, b: tl.constexpr, h: tl.constexpr, n,

10 + d: tl.constexpr, e: tl.constexpr, BLOCK: tl.constexpr,

11 + NUM_BLOCK, CBLOCK: tl.constexpr):

12 + # This kernel computes the diagonal blocks of the attention matrix

13 + # Each diagonal block represents attention

14 + # where queries attend to keys in the same block

15 + off = tl.program_id(0)

16 + off_bh = off // NUM_BLOCK # batch-head index

17 + off_block = off % NUM_BLOCK # block index within the sequence

18 + off_cblock = tl.program_id(1) # sub-block index within a block

19 +

20 + off_h = off_bh % h # head index

21 +

22 + # Calculate base offsets for the current batch and head

23 + qk_offset = off_bh * n * d

24 + v_offset = off_bh * n * e

25 + o_offset = off_bh * n * e

26 +

27 + # Calculate offsets for the current block

28 + block_offset = off_block * BLOCK

29 + qk_block_offset = block_offset * d

30 + v_block_offset = block_offset * e

31 + o_block_offset = block_offset * e

32 +

33 + # Calculate offsets for the current sub-block

34 + cblock_offset = off_cblock * CBLOCK

35 + q_cblock_offset = cblock_offset * d

36 + o_cblock_offset = cblock_offset * e

37 +

38 + # Calculate pointers to the query, key, value, and output tensors

39 + Q_block_ptr = (Q + qk_offset + qk_block_offset + q_cblock_offset +

40 + tl.arange(0, CBLOCK)[:, None] * d +

41 + tl.arange(0, d)[None, :])

42 + K_trans_block_ptr = (K + qk_offset + qk_block_offset +

43 + tl.arange(0, CBLOCK)[None, :] * d +

44 + tl.arange(0, d)[:, None])

45 + V_block_ptr = (V + v_offset + v_block_offset +

46 + tl.arange(0, CBLOCK)[:, None] * e +

47 + tl.arange(0, e)[None, :])

48 + O_block_ptr = (Out + o_offset + o_block_offset + o_cblock_offset +

49 + tl.arange(0, CBLOCK)[:, None] * e +

50 + tl.arange(0, e)[None, :])

51 +

52 + # Load the decay rate for the current head

53 + S_block_ptr = S + off_h

54 + s = tl.load(S_block_ptr)

55 +

56 + i = off_cblock

57 + q_index = tl.arange(0, CBLOCK) + i * CBLOCK

58 +

59 + # Load query values

60 + q = tl.load(Q_block_ptr,

61 + mask=block_offset + q_index[:, None] < n,

62 + other=0.0).to(tl.float32)

63 +

64 + # Initialize output accumulator

65 + qkv = tl.zeros([CBLOCK, e], dtype=tl.float32)

66 +

67 + # Process all sub-blocks up to and

68 + # including the current one (causal attention)

69 + for j in range(i + 1):

70 + kv_index = tl.arange(0, CBLOCK) + j * CBLOCK

71 + diff = q_index[:, None] - kv_index[None, :]

72 + s_index = s * diff

73 + # Apply causal mask: only attend to positions before the current one

74 + s_index = tl.where(diff >= 0, -s_index, float("-inf"))

75 + decay = tl.exp(s_index)

76 +

77 + # Load key and value

78 + k_trans = tl.load(

79 + K_trans_block_ptr,

80 + mask=block_offset + kv_index[None, :] < n,

81 + other=0.0,

82 + ).to(tl.float32)

83 + v = tl.load(

84 + V_block_ptr,

85 + mask=block_offset + kv_index[:, None] < n,

86 + other=0.0,

87 + ).to(tl.float32)

88 +

89 + # Compute attention scores and apply decay

90 + qk = tl.dot(q, k_trans) * decay

91 +

92 + # Compute weighted values and accumulate

93 + qkv += tl.dot(qk, v)

94 +

95 + # Move to the next sub-block

96 + K_trans_block_ptr += CBLOCK * d

97 + V_block_ptr += CBLOCK * e

98 +

99 + # Store the result

100 + tl.store(

101 + O_block_ptr,

102 + qkv.to(O_block_ptr.dtype.element_ty),

103 + mask=block_offset + q_index[:, None] < n,

104 + )

105 +

106 +

107 +@triton.jit

108 +def _fwd_kv_parallel(

109 + K,

110 + V,

111 + K_decay,

112 + KV,

113 + b: tl.constexpr,

114 + h: tl.constexpr,

115 + n,

116 + d: tl.constexpr,

117 + e: tl.constexpr,

118 + BLOCK: tl.constexpr,

119 + NUM_BLOCK,

120 + D_FBLOCK: tl.constexpr,

121 + E_FBLOCK: tl.constexpr,

122 + NUM_FBLOCK: tl.constexpr,

123 + CBLOCK: tl.constexpr,

124 + NUM_CBLOCK: tl.constexpr,

125 +):

126 + # This kernel computes the key-value outer

127 + # products for each block in parallel

128 + off_bh = tl.program_id(0) # batch-head index

129 + off_block = tl.program_id(1) # block index

130 +

131 + off_h = off_bh % h # head index

132 +

133 + block_offset = off_block * BLOCK

134 +

135 + # Calculate offsets for the current block

136 + k_block_offset = block_offset * d

137 + v_block_offset = block_offset * e

138 + kv_block_offset = off_block * d * e

139 +

140 + # Calculate base offsets for the current batch and head

141 + k_offset = off_bh * n * d

142 + v_offset = off_bh * n * e

143 + kv_offset = off_bh * NUM_BLOCK * d * e

144 +

145 + # Calculate pointers to the key, value, and key-value tensors

146 + K_trans_block_ptr = (K + k_offset + k_block_offset +

147 + tl.arange(0, CBLOCK)[None, :] * d +

148 + tl.arange(0, D_FBLOCK)[:, None])

149 + V_block_ptr = (V + v_offset + v_block_offset +

150 + tl.arange(0, CBLOCK)[:, None] * e +

151 + tl.arange(0, E_FBLOCK)[None, :])

152 + KV_block_ptr = (KV + kv_offset + kv_block_offset +

153 + tl.arange(0, D_FBLOCK)[:, None] * e +

154 + tl.arange(0, E_FBLOCK)[None, :])

155 +

156 + # Load the decay factors for the current head and block

157 + k_decay_ptr = (K_decay + off_h * BLOCK + tl.arange(0, CBLOCK)[None, :])

158 +

159 + kv_index = tl.arange(0, CBLOCK)

160 +

161 + # Initialize the key-value outer product accumulator

162 + kv = tl.zeros([D_FBLOCK, E_FBLOCK], dtype=tl.float32)

163 +

164 + # Handle the last block which might be smaller than BLOCK

165 + if off_block == NUM_BLOCK - 1:

166 + split_n = n - (NUM_BLOCK - 1) * BLOCK

167 + else:

168 + split_n = BLOCK

169 + left_shift = tl.cdiv(split_n, CBLOCK) * CBLOCK - split_n

170 + num_blocks = min(tl.cdiv(split_n, CBLOCK), NUM_CBLOCK)

171 + k_decay_ptr += (NUM_CBLOCK - num_blocks) * CBLOCK

172 +

173 + # Process all sub-blocks in the current block

174 + for j in range(num_blocks):

175 + left_bound = (1 - j) * left_shift

176 + # Load key and value, handling boundary conditions

177 + k_trans = tl.load(K_trans_block_ptr - left_shift * d,

178 + mask=kv_index[None, :] >= left_bound,

179 + other=0.0)

180 + v = tl.load(V_block_ptr - left_shift * e,

181 + mask=kv_index[:, None] >= left_bound,

182 + other=0.0)

183 +

184 + # Load decay factor and compute weighted key-value outer product

185 + k_decay = tl.load(k_decay_ptr)

186 + kv += tl.dot(k_trans * k_decay, v)

187 +

188 + # Move to the next sub-block

189 + K_trans_block_ptr += CBLOCK * d

190 + V_block_ptr += CBLOCK * e

191 + k_decay_ptr += CBLOCK

192 +

193 + # Store the result

194 + tl.store(KV_block_ptr, kv.to(KV_block_ptr.dtype.element_ty))

195 +

196 +

197 +@triton.jit

198 +def _fwd_kv_reduce(S, KV, KV_HISTORY, b: tl.constexpr, h: tl.constexpr, n,

199 + d: tl.constexpr, e: tl.constexpr, BLOCK: tl.constexpr,

200 + NUM_BLOCK, D_FBLOCK: tl.constexpr, E_FBLOCK: tl.constexpr):

201 + # This kernel reduces the key-value outer products

202 + # across blocks and updates the KV history

203 + off_bh = tl.program_id(0) # batch-head index

204 + off_h = off_bh % h # head index

205 +

206 + kv_offset = off_bh * NUM_BLOCK * d * e

207 +

208 + # Calculate pointer to the key-value tensor

209 + KV_block_ptr = (KV + kv_offset + tl.arange(0, D_FBLOCK)[:, None] * e +

210 + tl.arange(0, E_FBLOCK)[None, :])

211 +

212 + # Load the decay rate for the current head

213 + s_ptrs = S + off_h

214 + s = tl.load(s_ptrs)

215 +

216 + # Calculate pointer to the key-value history tensor

217 + kv_history_offset = off_bh * d * e

218 + KV_HISTORY_block_ptr = (KV_HISTORY + kv_history_offset +

219 + tl.arange(0, D_FBLOCK)[:, None] * e +

220 + tl.arange(0, E_FBLOCK)[None, :])

221 +

222 + # Load the previous key-value history

223 + kv_pre = tl.load(KV_HISTORY_block_ptr).to(tl.float32)

224 +

225 + # Process all blocks in reverse order to compute the prefix sum

226 + for i in range(NUM_BLOCK):

227 + block_size = min(n - i * BLOCK, BLOCK)

228 + # Compute decay factor for the current block

229 + block_decay = tl.exp(-s.to(tl.float32) * block_size)

230 +

231 + # Load the current key-value outer product

232 + kv_cur = tl.load(KV_block_ptr).to(tl.float32)

233 + # Store the previous key-value history to the current block

234 + tl.store(KV_block_ptr, kv_pre.to(KV_block_ptr.dtype.element_ty))

235 +

236 + # Update the key-value history with the current block

237 + kv_pre = block_decay * kv_pre + kv_cur

238 + KV_block_ptr += d * e

239 +

240 + # Store the updated key-value history

241 + tl.store(KV_HISTORY_block_ptr, kv_pre)

242 +

243 +

244 +@triton.jit

245 +def _fwd_none_diag_kernel(

246 + Q,

247 + Out,

248 + S,

249 + KV,

250 + b: tl.constexpr,

251 + h: tl.constexpr,

252 + n,

253 + d: tl.constexpr,

254 + e: tl.constexpr,

255 + BLOCK: tl.constexpr,

256 + NUM_BLOCK,

257 + E_FBLOCK: tl.constexpr,

258 + CBLOCK: tl.constexpr,

259 + NUM_CBLOCK: tl.constexpr,

260 +):

261 + # This kernel computes the non-diagonal blocks of the attention matrix

262 + # Each non-diagonal block represents attention

263 + # where queries attend to keys in different blocks

264 + off_bh = tl.program_id(0) # batch-head index

265 + off_h = off_bh % h # head index

266 +

267 + off_nc = tl.program_id(1)

268 + off_n = off_nc // NUM_CBLOCK # block index

269 + off_c = off_nc % NUM_CBLOCK # sub-block index

270 + off_e = tl.program_id(2) # output feature block index

271 +

272 + n_offset = off_n * BLOCK

273 + c_offset = off_c * CBLOCK

274 + e_offset = off_e * E_FBLOCK

275 + block_offset = n_offset + c_offset

276 +

277 + # Calculate offsets for the current batch, head, and block

278 + q_offset = off_bh * n * d + (n_offset + c_offset) * d

279 + o_offset = off_bh * n * e + (n_offset + c_offset) * e + e_offset

280 + kv_offset = off_bh * NUM_BLOCK * d * e + off_n * d * e + e_offset

281 +

282 + # Calculate pointers to the query, output, and key-value tensors

283 + Q_block_ptr = (Q + q_offset + tl.arange(0, CBLOCK)[:, None] * d +

284 + tl.arange(0, d)[None, :])

285 + O_block_ptr = (Out + o_offset + tl.arange(0, CBLOCK)[:, None] * e +

286 + tl.arange(0, E_FBLOCK)[None, :])

287 + KV_block_ptr = (KV + kv_offset + tl.arange(0, d)[:, None] * e +

288 + tl.arange(0, E_FBLOCK)[None, :])

289 +

290 + # Load the decay rate for the current head

291 + S_block_ptr = S + off_h

292 + s = tl.load(S_block_ptr)

293 +

294 + c_array = tl.arange(0, CBLOCK)

295 +

296 + # Load the key-value outer product for the current block

297 + kv = tl.load(KV_block_ptr).to(tl.float32)

298 + q_index = block_offset + tl.arange(0, CBLOCK)

299 +

300 + # Load query values

301 + q = tl.load(Q_block_ptr, mask=q_index[:, None] < n,

302 + other=0.).to(tl.float32)

303 +

304 + # Compute decay factors for the current sub-block

305 + q_decay = tl.exp(-s.to(tl.float32) * (off_c * CBLOCK + c_array[:, None]))

306 +

307 + # Compute non-diagonal attention output

308 + qkv_none_diag = tl.dot(q, kv) * q_decay

309 +

310 + # Load diagonal attention output (computed by _fwd_diag_kernel)

311 + qkv_diag = tl.load(O_block_ptr, mask=q_index[:, None] < n,

312 + other=0.).to(tl.float32)

313 +

314 + # Combine diagonal and non-diagonal attention outputs

315 + qkv = qkv_diag + qkv_none_diag

316 +

317 + # Store the result

318 + tl.store(O_block_ptr,

319 + qkv.to(O_block_ptr.dtype.element_ty),

320 + mask=q_index[:, None] < n)

321 +

322 +

323 +class _attention(torch.autograd.Function):

324 +

325 + @staticmethod

326 + def forward(ctx, q, k, v, s, kv_history):

327 + # Forward pass of the lightning attention algorithm

328 + q = q.contiguous()

329 + k = k.contiguous()

330 + v = v.contiguous()

331 + s = s.contiguous()

332 +

333 + # Check CUDA compute capability

334 + capability = torch.cuda.get_device_capability()

335 + if capability[0] < 8:

336 + raise RuntimeError("Flash attention currently only supported",

337 + "for compute capability >= 80")

338 +

339 + # Get input dimensions

340 + b, h, n, d = q.shape

341 + e = v.shape[-1]

342 +

343 + # Initialize output tensor

344 + o = torch.empty((b, h, n, e), dtype=q.dtype, device=q.device)

345 +

346 + # Set block sizes

347 + BLOCK = 256

348 + NUM_BLOCK = triton.cdiv(n, BLOCK)

349 +

350 + CBLOCK = 32

351 + NUM_CBLOCK = BLOCK // CBLOCK

352 + assert BLOCK % CBLOCK == 0, "BLOCK must be a multiple of CBLOCK"

353 +

354 + # Compute decay factors for keys

355 + array = torch.arange(0, BLOCK, device=q.device) + 1

356 + k_decay = torch.exp(-s * (BLOCK - array.reshape(1, -1)))

357 +

358 + # Step 1: Compute diagonal blocks of attention

359 + grid = (b * h * NUM_BLOCK, NUM_CBLOCK)

360 + _fwd_diag_kernel[grid](q,

361 + k,

362 + v,

363 + o,

364 + s,

365 + b,

366 + h,

367 + n,

368 + d,

369 + e,

370 + BLOCK=BLOCK,

371 + NUM_BLOCK=NUM_BLOCK,

372 + CBLOCK=CBLOCK)

373 +

374 + # Set feature block sizes

375 + NUM_FBLOCK = 1

376 + D_FBLOCK = d // NUM_FBLOCK

377 + assert d % NUM_FBLOCK == 0

378 + E_FBLOCK = e // NUM_FBLOCK

379 + assert e % NUM_FBLOCK == 0

380 +

381 + CBLOCK = 64

382 + NUM_CBLOCK = BLOCK // CBLOCK

383 + assert BLOCK % CBLOCK == 0, "BLOCK must be a multiple of CBLOCK"

384 +

385 + # Step 2: Compute key-value outer products for each block in parallel

386 + kv = torch.empty((b, h, NUM_BLOCK, d, e),

387 + dtype=torch.float32,

388 + device=q.device)

389 + grid = (b * h, NUM_BLOCK)

390 + _fwd_kv_parallel[grid](

391 + k,

392 + v,

393 + k_decay,

394 + kv,

395 + b,

396 + h,

397 + n,

398 + d,

399 + e,

400 + BLOCK=BLOCK,

401 + NUM_BLOCK=NUM_BLOCK,

402 + D_FBLOCK=D_FBLOCK,

403 + E_FBLOCK=E_FBLOCK,

404 + NUM_FBLOCK=NUM_FBLOCK,

405 + CBLOCK=CBLOCK,

406 + NUM_CBLOCK=NUM_CBLOCK,

407 + )

408 +

409 + # Step 3: Reduce key-value outer products

410 + # across blocks and update KV history

411 + grid = (b * h, NUM_FBLOCK)

412 + _fwd_kv_reduce[grid](s,

413 + kv,

414 + kv_history,

415 + b,

416 + h,

417 + n,

418 + d,

419 + e,

420 + BLOCK=BLOCK,

421 + NUM_BLOCK=NUM_BLOCK,

422 + D_FBLOCK=D_FBLOCK,

423 + E_FBLOCK=E_FBLOCK)

424 +

425 + # Step 4: Compute non-diagonal blocks of attention

426 + grid = (b * h, NUM_BLOCK * NUM_CBLOCK)

427 + _fwd_none_diag_kernel[grid](

428 + q,

429 + o,

430 + s,

431 + kv,

432 + b,

433 + h,

434 + n,

435 + d,

436 + e,

437 + BLOCK=BLOCK,

438 + NUM_BLOCK=NUM_BLOCK,

439 + E_FBLOCK=E_FBLOCK,

440 + CBLOCK=CBLOCK,

441 + NUM_CBLOCK=NUM_CBLOCK,

442 + )

443 +

444 + # Save tensors for backward pass

445 + ctx.save_for_backward(q, k, v, s, kv)

446 + ctx.BLOCK = BLOCK

447 +

448 + return o, torch.cat([kv, kv_history.unsqueeze(2)], dim=2)

449 +

450 +

451 +# Apply the lightning attention function

452 +lightning_attention_ = _attention.apply

453 +

454 +

455 +def lightning_attention(q, k, v, ed, block_size=256, kv_history=None):

456 + """

457 + Apply lightning attention algorithm

458 + to compute attention efficiently.

459 +

460 + Args:

461 + q: Query tensor of shape [batch, heads, seq_len, dim]

462 + k: Key tensor of shape [batch, heads, seq_len, dim]

463 + v: Value tensor of shape [batch, heads, seq_len, dim_v]

464 + ed: Decay rate tensor of shape [heads]

465 + block_size: Size of blocks for block-sparse attention

466 + kv_history: Optional key-value history from previous computations

467 +

468 + Returns:

469 + output: Attention output

470 + kv: Updated key-value history

471 + """

472 + d = q.shape[-1]

473 + e = v.shape[-1]

474 +

475 + if ed.dim() == 1:

476 + ed = ed.view(1, -1, 1, 1)

477 +

478 + # Split the computation into chunks for better parallelism

479 + m = 128 if d >= 128 else 64

480 + assert d % m == 0, f"Dimension d ({d}) must be divisible by m ({m})"

481 + arr = [m * i for i in range(d // m + 1)]

482 + if arr[-1] != d:

483 + arr.append(d)

484 + n = len(arr)

485 + output = 0

486 +

487 + # Initialize or clone key-value history

488 + if kv_history is None:

489 + kv_history = torch.zeros((q.shape[0], q.shape[1], d, e),

490 + dtype=torch.float32,

491 + device=q.device)

492 + else:

493 + kv_history = kv_history.clone().contiguous()

494 +

495 + # Process each chunk and accumulate results

496 + for i in range(n - 1):

497 + s = arr[i]

498 + e = arr[i + 1]

499 + q1 = q[..., s:e]

500 + k1 = k[..., s:e]

501 + o, kv = lightning_attention_(q1, k1, v, ed, kv_history)

502 + output = output + o

503 + return output, kv

504 +

505 +

506 +@triton.jit

507 +def _linear_attn_decode_kernel(

508 + q_ptr,

509 + k_ptr,

510 + v_ptr,

511 + kv_cache_ptr,

512 + slope_rate,

513 + slot_idx,

514 + output_ptr,

515 + D: tl.constexpr,

516 + qkv_b_stride,

517 + qkv_h_stride,

518 + cache_b_stride,

519 + cache_h_stride,

520 + cache_d0_stride,

521 + cache_d1_stride,

522 + BLOCK_SIZE: tl.constexpr,

523 +):

524 + """

525 + Kernel for linear attention decoding with KV cache.

526 +

527 + This kernel computes attention for a single token using the KV cache.

528 + """

529 + pid_b = tl.program_id(0) # batch index

530 + pid_h = tl.program_id(1) # head index

531 + pid_d = tl.program_id(2) # dimension block index

532 +

533 + # Load slot index for the current batch

534 + slot_id = tl.load(slot_idx + pid_b)

535 +

536 + # Skip if slot_id is -1 (padding)

537 + if slot_id == -1:

538 + return

539 +

540 + batch_id = pid_b

541 + head_id = pid_h

542 +

543 + # Load decay rate for the current head

544 + ratio = tl.load(slope_rate + pid_h)

545 +

546 + # Calculate offsets for dimensions

547 + qk_d_offsets = tl.arange(0, D)

548 + v_d_offsets = tl.arange(0, BLOCK_SIZE) + pid_d * BLOCK_SIZE

549 + cache_d_offsets = qk_d_offsets[:, None] * cache_d0_stride + v_d_offsets[

550 + None, :] * cache_d1_stride

551 +

552 + # Calculate offsets for the current batch and head

553 + q_offset = batch_id * qkv_b_stride + head_id * qkv_h_stride

554 + k_offset = batch_id * qkv_b_stride + head_id * qkv_h_stride

555 + v_offset = batch_id * qkv_b_stride + head_id * qkv_h_stride

556 +

557 + cache_offset = slot_id * cache_b_stride + head_id * cache_h_stride

558 +

559 + # Create masks for loading tensors

560 + qk_mask = qk_d_offsets < D

561 + v_mask = v_d_offsets < D

562 +

563 + # Load query, key, and value tensors

564 + q = tl.load(q_ptr + q_offset + qk_d_offsets, mask=qk_mask, other=0.0)

565 + k = tl.load(k_ptr + k_offset + qk_d_offsets, mask=qk_mask, other=0.0)

566 + v = tl.load(v_ptr + v_offset + v_d_offsets, mask=v_mask, other=0.0)

567 +

568 + # Compute key-value outer product

569 + kv_outer = k[:, None] * v[None, :]

570 + kv_mask = qk_mask[:, None] & v_mask[None, :]

571 +

572 + # Apply decay to previous KV cache

573 + ratio = tl.exp(-ratio)

574 + kv_ptr = kv_cache_ptr + cache_offset + cache_d_offsets

575 + kv_cache_old = tl.load(kv_ptr, mask=kv_mask, other=0.0)

576 + kv_outer = kv_outer + ratio * kv_cache_old

577 +

578 + # Compute attention output

579 + output = q[:, None].to(tl.float32) * kv_outer

580 + output = tl.sum(output, axis=0)

581 +

582 + # Update KV cache and store output

583 + tl.store(kv_ptr, kv_outer, mask=kv_mask)

584 + tl.store(output_ptr + q_offset + v_d_offsets, output, mask=v_mask)

585 +

586 +

587 +def linear_decode_forward_triton(

588 + q: torch.Tensor,

589 + k: torch.Tensor,

590 + v: torch.Tensor,

591 + kv_caches: torch.Tensor,

592 + slope_rate: torch.Tensor,

593 + slot_idx: torch.Tensor,

594 + BLOCK_SIZE: int = 32,

595 +) -> torch.Tensor:

596 + """

597 + Perform linear attention decoding using Triton kernels.

598 +

599 + Args:

600 + q: Query tensor of shape [B, H, 1, D]

601 + k: Key tensor of shape [B, H, 1, D]

602 + v: Value tensor of shape [B, H, 1, D]

603 + kv_caches: Key-value cache tensor

604 + slope_rate: Decay rate tensor

605 + slot_idx: Slot indices for batches

606 + BLOCK_SIZE: Size of blocks for processing

607 +

608 + Returns:

609 + output: Attention output tensor

610 + """

611 + B, H, _, D = q.shape

612 + assert k.shape == (B, H, 1, D)

613 + assert v.shape == (B, H, 1, D)

614 +

615 + # Initialize output tensor

616 + output = torch.empty_like(q)

617 +

618 + # Set grid dimensions for the kernel

619 + grid = (B, H, D // BLOCK_SIZE)

620 +

621 + # Calculate strides for tensors

622 + qkv_b_stride = q.stride(0)

623 + qkv_h_stride = q.stride(1)

624 +

625 + cache_b_stride = kv_caches.stride(0)

626 + cache_h_stride = kv_caches.stride(1)

627 + cache_d0_stride = kv_caches.stride(2)

628 + cache_d1_stride = kv_caches.stride(3)

629 +

630 + # Launch the kernel

631 + _linear_attn_decode_kernel[grid](

632 + q,

633 + k,

634 + v,

635 + kv_caches,

636 + slope_rate,

637 + slot_idx,

638 + output,

639 + D,

640 + qkv_b_stride,

641 + qkv_h_stride,

642 + cache_b_stride,

643 + cache_h_stride,

644 + cache_d0_stride,

645 + cache_d1_stride,

646 + BLOCK_SIZE=BLOCK_SIZE,

647 + )

648 +

649 + # Reshape output and return

650 + output = rearrange(output, "b h n d -> b n (h d)")

651 + return output.squeeze(1).contiguous()

vllm/model_executor/models/constant_size_cache.py new

+136

@@ -0,0 +1,136 @@

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

2 +from abc import ABC, abstractmethod

3 +from typing import Any, Dict, List, Tuple

4 +

5 +import torch

6 +

7 +from vllm.attention.backends.utils import PAD_SLOT_ID

8 +

9 +

10 +class ConstantSizeCache(ABC):

11 + """

12 + Abstract base class for managing constant size caches

13 + like Mamba and Minimax.

14 + """

15 +

16 + def __init__(self, max_batch_size: int):

17 + # Maps between the request id and a dict that maps between the seq_id

18 + # and its index inside the cache

19 + self.cache_indices_mapping: Dict[str, Dict[int, int]] = {}

20 + self.free_cache_indices = list(range(max_batch_size))

21 +

22 + @property

23 + @abstractmethod

24 + def cache(self) -> Any:

25 + """Return the underlying cache tensor(s)"""

26 + pass

27 +

28 + @abstractmethod

29 + def _copy_cache(self, from_index: int, to_index: int):

30 + """Copy cache data from one index to another"""

31 + pass

32 +

33 + def current_run_tensors(self, **kwargs) -> Tuple:

34 + """

35 + Return the tensors for the current run's conv and ssm state.

36 + """

37 + if "seqlen_agnostic_capture_inputs" not in kwargs:

38 + # We get here only on Prefill/Eager mode runs

39 + request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"]

40 + finished_requests_ids = kwargs["finished_requests_ids"]

41 +

42 + self._release_finished_requests(finished_requests_ids)

43 + state_indices = self._prepare_current_run_cache(

44 + request_ids_to_seq_ids, finished_requests_ids)

45 +

46 + state_indices_tensor = torch.as_tensor(state_indices,

47 + dtype=torch.int32,

48 + device="cuda")

49 + cache_tensors = self.cache

50 + else:

51 + # CUDA graph capturing runs

52 + cache_tensors, state_indices_tensor = kwargs[

53 + "seqlen_agnostic_capture_inputs"]

54 +

55 + return (cache_tensors, state_indices_tensor)

56 +

57 + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):

58 + """

59 + Copy the relevant state_indices into the CUDA graph input buffer

60 + """

61 + assert all(

62 + key in kwargs

63 + for key in ["request_ids_to_seq_ids", "finished_requests_ids"])

64 + finished_requests_ids = kwargs["finished_requests_ids"]

65 + request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"]

66 + assert "seqlen_agnostic_capture_inputs" in input_buffers

67 + _, input_state_indices_buffer = input_buffers[

68 + "seqlen_agnostic_capture_inputs"]

69 +

70 + self._release_finished_requests(finished_requests_ids)

71 + state_indices = self._prepare_current_run_cache(

72 + request_ids_to_seq_ids, finished_requests_ids)

73 + cuda_graph_pad_len = input_state_indices_buffer.shape[0] - len(

74 + state_indices)

75 + state_indices.extend([PAD_SLOT_ID] * cuda_graph_pad_len)

76 +

77 + input_state_indices_buffer.copy_(

78 + torch.as_tensor(state_indices, dtype=torch.int32, device="cuda"))

79 +

80 + def get_seqlen_agnostic_capture_inputs(self, batch_size: int):

81 + """

82 + Provide the CUDA graph capture runs with a buffer in adjusted size.

83 + The buffer is used to maintain the Cache during the CUDA graph replay

84 + runs.

85 + """

86 + state_indices_tensor = torch.as_tensor([PAD_SLOT_ID] * batch_size,

87 + dtype=torch.int32,

88 + device="cuda")

89 + return (self.cache, state_indices_tensor)

90 +

91 + def _assign_seq_id_to_cache_index(self, cur_rid: str, seq_id: int,

92 + finished_requests_ids) -> int:

93 + """

94 + Assign (req_id,seq_id) pair to a `destination_index` index, if

95 + already occupied, move the occupying index to a free index.

96 + """

97 + if cur_rid in finished_requests_ids:

98 + # set as pad, do not allocate destination index

99 + return PAD_SLOT_ID

100 + elif cur_rid not in self.cache_indices_mapping:

101 + destination_index = self.free_cache_indices.pop()

102 + self.cache_indices_mapping[cur_rid] = {seq_id: destination_index}

103 + return destination_index

104 + elif seq_id not in (seq_ids2indices :=

105 + self.cache_indices_mapping[cur_rid]):

106 + # parallel sampling , where n > 1, assume prefill have

107 + # already happened, so we copy the

108 + # existing cache into the siblings seq_ids caches

109 + index_exists = next(iter(seq_ids2indices.values()))

110 + # case of decoding n>1, copy prefill cache to decoding indices

111 + destination_index = self.free_cache_indices.pop()

112 + self._copy_cache(from_index=index_exists,

113 + to_index=destination_index)

114 + self.cache_indices_mapping[cur_rid][seq_id] = destination_index

115 + return destination_index

116 + else:

117 + return self.cache_indices_mapping[cur_rid][seq_id]

118 +

119 + def _prepare_current_run_cache(

120 + self, request_ids_to_seq_ids: Dict[str, list[int]],

121 + finished_requests_ids: List[str]) -> List[int]:

122 + return [

123 + self._assign_seq_id_to_cache_index(req_id, seq_id,

124 + finished_requests_ids)

125 + for req_id, seq_ids in request_ids_to_seq_ids.items()

126 + for seq_id in seq_ids

127 + ]

128 +

129 + def _release_finished_requests(self,

130 + finished_seq_groups_req_ids: List[str]):

131 + for req_id in finished_seq_groups_req_ids:

132 + if req_id in self.cache_indices_mapping:

133 + for seq_id in self.cache_indices_mapping[req_id]:

134 + self.free_cache_indices.append(

135 + self.cache_indices_mapping[req_id][seq_id])

136 + self.cache_indices_mapping.pop(req_id)

vllm/model_executor/models/mamba_cache.py

+21 -111

@@ -1,12 +1,13 @@

1 # SPDX-License-Identifier: Apache-2.0

3 from dataclasses import dataclass

4 -from typing import Dict, List, Tuple

4 +from typing import Tuple

6 import torch

8 from vllm.attention.backends.utils import PAD_SLOT_ID

9 from vllm.config import VllmConfig

10 +from vllm.model_executor.models.constant_size_cache import ConstantSizeCache

13 @dataclass

@@ -21,7 +22,7 @@ class MambaCacheParams:

22 self.state_indices_tensor)

24 -class MambaCacheManager:

25 +class MambaCacheManager(ConstantSizeCache):

27 def __init__(self, vllm_config: VllmConfig, dtype: torch.dtype,

28 num_mamba_layers: int, conv_state_shape: Tuple[int, int],

@@ -32,6 +33,9 @@ class MambaCacheManager:

33 if not vllm_config.model_config.enforce_eager:

34 max_batch_size = vllm_config.pad_for_cudagraph(max_batch_size)

36 + # Initialize parent class

37 + super().__init__(max_batch_size)

38 +

39 conv_state = torch.empty(size=(num_mamba_layers, max_batch_size) +

40 conv_state_shape,

41 dtype=dtype,

@@ -41,126 +45,32 @@ class MambaCacheManager:

45 dtype=dtype,

46 device="cuda")

44 - self.mamba_cache = (conv_state, temporal_state)

48 + self._mamba_cache = (conv_state, temporal_state)

49 +

50 + @property

51 + def cache(self):

52 + return self._mamba_cache

46 - # Maps between the request id and a dict that maps between the seq_id

47 - # and its index inside the self.mamba_cache

48 - self.mamba_cache_indices_mapping: Dict[str, Dict[int, int]] = {}

49 - self.free_cache_indices = list(range(max_batch_size))

54 + def _copy_cache(self, from_index: int, to_index: int):

55 + for cache_t in self.cache:

56 + cache_t[:, to_index].copy_(cache_t[:, from_index],

57 + non_blocking=True)

59 def current_run_tensors(self, **kwargs) -> MambaCacheParams:

60 """

61 Return the tensors for the current run's conv and ssm state.

62 """

55 - if "seqlen_agnostic_capture_inputs" not in kwargs:

56 - # We get here only on Prefill/Eager mode runs

57 - request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"]

58 - finished_requests_ids = kwargs["finished_requests_ids"]

59 -

60 - self._release_finished_requests(finished_requests_ids)

61 - state_indices = self._prepare_current_run_mamba_cache(

62 - request_ids_to_seq_ids, finished_requests_ids)

63 -

64 - state_indices_tensor = torch.as_tensor(state_indices,

65 - dtype=torch.int32,

66 - device="cuda")

67 - mamba_cache_tensors = self.mamba_cache

68 -

69 - else:

70 - # CUDA graph capturing runs

71 - (mamba_cache_tensors,

72 - state_indices_tensor) = kwargs["seqlen_agnostic_capture_inputs"]

73 -

74 - return MambaCacheParams(mamba_cache_tensors[0], mamba_cache_tensors[1],

63 + cache_tensors, state_indices_tensor = super().current_run_tensors(

64 + **kwargs)

65 + return MambaCacheParams(cache_tensors[0], cache_tensors[1],

66 state_indices_tensor)

77 - def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):

78 - """

79 - Copy the relevant state_indices into the CUDA graph input buffer

80 - """

81 - assert all(

82 - key in kwargs

83 - for key in ["request_ids_to_seq_ids", "finished_requests_ids"])

84 - finished_requests_ids = kwargs["finished_requests_ids"]

85 - request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"]

86 - assert "seqlen_agnostic_capture_inputs" in input_buffers

87 - _, input_state_indices_buffer = input_buffers[

88 - "seqlen_agnostic_capture_inputs"]

89 -

90 - self._release_finished_requests(finished_requests_ids)

91 - state_indices = self._prepare_current_run_mamba_cache(

92 - request_ids_to_seq_ids, finished_requests_ids)

93 - cuda_graph_pad_len = input_state_indices_buffer.shape[0] - len(

94 - state_indices)

95 - state_indices.extend([PAD_SLOT_ID] * cuda_graph_pad_len)

96 -

97 - input_state_indices_buffer.copy_(

98 - torch.as_tensor(state_indices, dtype=torch.int32, device="cuda"))

99 -

68 def get_seqlen_agnostic_capture_inputs(self, batch_size: int):

69 """

70 Provide the CUDA graph capture runs with a buffer in adjusted size.

71 The buffer is used to maintain the Mamba Cache during the CUDA graph

72 replay runs.

73 """

106 - state_indices_tensor = torch.as_tensor([PAD_SLOT_ID] * batch_size,

107 - dtype=torch.int32,

108 - device="cuda")

109 - return (self.mamba_cache, state_indices_tensor)

110 -

111 - def _copy_mamba_cache(self, from_index: int, to_index: int):

112 - assert len(self.mamba_cache) > 0

113 - for cache_t in self.mamba_cache:

114 - cache_t[:, to_index].copy_(cache_t[:, from_index],

115 - non_blocking=True)

116 -

117 - def _assign_seq_id_to_cache_index(self, cur_rid: str, seq_id: int,

118 - finished_requests_ids) -> int:

119 - """

120 - Assign (req_id,seq_id) pair to a `destination_index` index, if

121 - already occupied, move the occupying index to a free index.

122 - """

123 - if cur_rid in finished_requests_ids:

124 - # set as pad, do not allocate destination index

125 - return PAD_SLOT_ID

126 - elif cur_rid not in self.mamba_cache_indices_mapping:

127 - destination_index = self.free_cache_indices.pop()

128 - self.mamba_cache_indices_mapping[cur_rid] = {

129 - seq_id: destination_index

130 - }

131 - return destination_index

132 - elif seq_id not in (seq_ids2indices :=

133 - self.mamba_cache_indices_mapping[cur_rid]):

134 - # parallel sampling , where n > 1, assume prefill have

135 - # already happened, so we copy the

136 - # existing cache into the siblings seq_ids caches

137 - index_exists = next(iter(seq_ids2indices.values()))

138 - # case of decoding n>1, copy prefill cache to decoding indices

139 - destination_index = self.free_cache_indices.pop()

140 - self._copy_mamba_cache(from_index=index_exists,

141 - to_index=destination_index)

142 - self.mamba_cache_indices_mapping[cur_rid][

143 - seq_id] = destination_index

144 - return destination_index

145 - else:

146 - # already exists

147 - return self.mamba_cache_indices_mapping[cur_rid][seq_id]

148 -

149 - def _prepare_current_run_mamba_cache(

150 - self, request_ids_to_seq_ids: Dict[str, list[int]],

151 - finished_requests_ids: List[str]) -> List[int]:

152 - return [

153 - self._assign_seq_id_to_cache_index(req_id, seq_id,

154 - finished_requests_ids)

155 - for req_id, seq_ids in request_ids_to_seq_ids.items()

156 - for seq_id in seq_ids

157 - ]

158 -

159 - def _release_finished_requests(self,

160 - finished_seq_groups_req_ids: List[str]):

161 - for req_id in finished_seq_groups_req_ids:

162 - if req_id in self.mamba_cache_indices_mapping:

163 - for seq_id in self.mamba_cache_indices_mapping[req_id]:

164 - self.free_cache_indices.append(

165 - self.mamba_cache_indices_mapping[req_id][seq_id])

166 - self.mamba_cache_indices_mapping.pop(req_id)

74 + return self._mamba_cache, torch.as_tensor([PAD_SLOT_ID] * batch_size,

75 + dtype=torch.int32,

76 + device="cuda")

vllm/model_executor/models/minimax_cache.py new

+35

@@ -0,0 +1,35 @@

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

2 +from dataclasses import dataclass

3 +

4 +import torch

5 +

6 +from vllm.model_executor.models.constant_size_cache import ConstantSizeCache

7 +

8 +

9 +@dataclass

10 +class MinimaxCacheParams:

11 + minimax_cache: torch.Tensor = torch.Tensor()

12 + state_indices_tensor: torch.Tensor = torch.Tensor()

13 +

14 + def at_layer_idx(self, layer_idx):

15 + return MinimaxCacheParams(self.minimax_cache[layer_idx, ...],

16 + self.state_indices_tensor)

17 +

18 +

19 +class MinimaxCacheManager(ConstantSizeCache):

20 +

21 + def __init__(self, dtype, cache_shape):

22 + super().__init__(cache_shape[1]) # max_batch_size is cache_shape[1]

23 + self._minimax_cache = torch.empty(size=cache_shape,

24 + dtype=dtype,

25 + device="cuda")

26 +

27 + @property

28 + def cache(self):

29 + return self._minimax_cache

30 +

31 + def _copy_cache(self, from_index: int, to_index: int):

32 + assert len(self.cache) > 0

33 + for cache_t in self.cache:

34 + cache_t[:, to_index].copy_(cache_t[:, from_index],

35 + non_blocking=True)

vllm/model_executor/models/minimax_text_01.py new

+1273

@@ -0,0 +1,1273 @@

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

2 +"""Inference-only MiniMaxText01 model."""

3 +import copy

4 +import math

5 +import re

6 +from typing import Dict, Iterable, List, Optional, Tuple, Union

7 +

8 +import torch

9 +import torch.distributed

10 +import torch.nn.functional as F

11 +from einops import rearrange

12 +from torch import nn

13 +from transformers.configuration_utils import PretrainedConfig

14 +

15 +from vllm.attention import Attention, AttentionMetadata

16 +from vllm.config import CacheConfig, VllmConfig

17 +from vllm.distributed.communication_op import tensor_model_parallel_all_reduce

18 +from vllm.distributed.parallel_state import (

19 + get_pp_group, get_tensor_model_parallel_rank,

20 + get_tensor_model_parallel_world_size)

21 +from vllm.forward_context import get_forward_context

22 +from vllm.model_executor.custom_op import CustomOp

23 +from vllm.model_executor.layers.activation import SiluAndMul

24 +from vllm.model_executor.layers.fused_moe import FusedMoE

25 +from vllm.model_executor.layers.layernorm import RMSNorm

26 +from vllm.model_executor.layers.lightning_attn import (

27 + lightning_attention, linear_decode_forward_triton)

28 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,

29 + MergedColumnParallelLinear,

30 + QKVParallelLinear,

31 + ReplicatedLinear,

32 + RowParallelLinear)

33 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

34 +from vllm.model_executor.layers.quantization.base_config import (

35 + QuantizationConfig)

36 +from vllm.model_executor.layers.sampler import Sampler

37 +from vllm.model_executor.layers.vocab_parallel_embedding import (

38 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

39 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

40 +from vllm.model_executor.models.utils import maybe_prefix

41 +from vllm.model_executor.sampling_metadata import SamplingMetadata

42 +from vllm.sequence import IntermediateTensors

43 +

44 +from .interfaces import HasInnerState, IsHybrid, SupportsV0Only

45 +from .minimax_cache import MinimaxCacheManager, MinimaxCacheParams

46 +from .utils import PPMissingLayer, is_pp_missing_parameter, make_layers

47 +

48 +

49 +def replace_weight_name(name: str,

50 + key: str = None,

51 + to: str = None,

52 + count: int = None,

53 + prefix: str = None) -> str:

54 + name = name.replace(key, to) if count is None else \

55 + name.replace(key, to, count)

56 + return name

57 +

58 +

59 +def weight_loader_with_alias(alias: str):

60 +

61 + def wrapper(func: callable):

62 +

63 + def inner_func(param: torch.Tensor,

64 + loaded_weight: torch.Tensor,

65 + *args,

66 + prefix: str = None,

67 + **kwargs):

68 + value = func(param, loaded_weight, *args, **kwargs)

69 + return value

70 +

71 + return inner_func

72 +

73 + return wrapper

74 +

75 +

76 +class MiniMaxText01RMSNormTP(CustomOp):

77 + name = "MiniMaxText01RMSNormTP"

78 +

79 + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:

80 + super().__init__()

81 + self.tp_world = get_tensor_model_parallel_world_size()

82 + self.tp_rank = get_tensor_model_parallel_rank()

83 + self.weight = nn.Parameter(torch.ones(int(hidden_size /

84 + self.tp_world)))

85 +

86 + self.weight.weight_loader = self.weight_loader

87 + self.variance_epsilon = eps

88 + return

89 +

90 + @staticmethod

91 + def weight_loader(

92 + param: nn.Parameter,

93 + loaded_weight: torch.Tensor,

94 + ) -> None:

95 + tp_world = get_tensor_model_parallel_world_size()

96 + tp_rank = get_tensor_model_parallel_rank()

97 +

98 + shard_size = loaded_weight.shape[0] // tp_world

99 + shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)

100 + param.data.copy_(loaded_weight[shard])

101 + return

102 +

103 + def _forward(

104 + self,

105 + x: torch.Tensor,

106 + ) -> torch.Tensor:

107 + orig_dtype = x.dtype

108 + x = x.to(torch.float32)

109 + variance = x.pow(2).mean(dim=-1, keepdim=True, dtype=torch.float32)

110 + if self.tp_world > 1:

111 + variance = tensor_model_parallel_all_reduce(

112 + variance) / self.tp_world

113 + x = x * torch.rsqrt(variance + self.variance_epsilon)

114 + x = x.to(orig_dtype) * self.weight

115 + return x

116 +

117 + def forward(

118 + self,

119 + x: torch.Tensor,

120 + residual: Optional[torch.Tensor] = None,

121 + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:

122 + assert residual is None, "RMSNorm does not support residual connection."

123 + return self._forward(x)

124 +

125 +

126 +class MiniMaxText01RotaryEmbedding(CustomOp):

127 + name = "MiniMaxText01RotaryEmbedding"

128 +

129 + def __init__(

130 + self,

131 + head_size: int,

132 + rotary_dim: int,

133 + max_position: int,

134 + base: int,

135 + is_neox_style: bool,

136 + cache_dtype: torch.dtype,

137 + ) -> None:

138 + super().__init__()

139 + self.head_size = head_size

140 + self.rotary_dim = rotary_dim

141 + self.max_position_embeddings = max_position

142 + self.base = base

143 + self.is_neox_style = is_neox_style

144 + self.cache_dtype = cache_dtype

145 + cache = self._compute_cos_sin_cache().to(cache_dtype)

146 + self.register_buffer("cos_sin_cache", cache, persistent=False)

147 +

148 + def _compute_inv_freq(

149 + self,

150 + base: Union[int, float],

151 + ) -> torch.Tensor:

152 + """Compute the inverse frequency."""

153 + inv_freq = 1.0 / (base**(torch.arange(

154 + 0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim))

155 + return inv_freq

156 +

157 + def _compute_cos_sin_cache(self) -> torch.Tensor:

158 + """Compute the cos and sin cache."""

159 + inv_freq = self._compute_inv_freq(self.base)

160 + t = torch.arange(self.max_position_embeddings, dtype=torch.float)

161 + freqs = torch.einsum("i,j -> ij", t, inv_freq)

162 + cos = freqs.cos()

163 + sin = freqs.sin()

164 + cache = torch.cat((cos, sin), dim=-1)

165 + return cache

166 +

167 + def forward(

168 + self,

169 + positions: torch.Tensor,

170 + query: torch.Tensor,

171 + key: torch.Tensor,

172 + ) -> Tuple[torch.Tensor, torch.Tensor]:

173 + from vllm import _custom_ops as ops

174 + self.cos_sin_cache = self.cos_sin_cache.to(positions.device)

175 + query_cast = query.to(self.cache_dtype)

176 + key_cast = key.to(self.cache_dtype)

177 + ops.rotary_embedding(positions, query_cast, key_cast, self.head_size,

178 + self.cos_sin_cache, self.is_neox_style)

179 + query = query_cast.to(query.dtype)

180 + key = key_cast.to(key.dtype)

181 + return query, key

182 +

183 +

184 +class MiniMaxText01MLP(nn.Module):

185 +

186 + def __init__(

187 + self,

188 + hidden_size: int,

189 + intermediate_size: int,

190 + quant_config: Optional[QuantizationConfig] = None,

191 + layer_idx: int = None,

192 + prefix: str = "mlp",

193 + ) -> None:

194 + super().__init__()

195 + self.layer_idx = layer_idx

196 +

197 + self.gate_up_proj = MergedColumnParallelLinear(

198 + hidden_size,

199 + [intermediate_size] * 2,

200 + bias=False,

201 + quant_config=quant_config,

202 + prefix=f"{prefix}.gate_up_proj",

203 + )

204 + self.down_proj = RowParallelLinear(

205 + intermediate_size,

206 + hidden_size,

207 + bias=False,

208 + quant_config=quant_config,

209 + prefix=f"{prefix}.down_proj",

210 + )

211 + self.act_fn = SiluAndMul()

212 + return

213 +

214 + def forward(self, x: torch.Tensor) -> torch.Tensor:

215 +

216 + gate_up, _ = self.gate_up_proj(x)

217 + x = self.act_fn(gate_up)

218 + x, _ = self.down_proj(x)

219 + return x

220 +

221 +

222 +class MiniMaxText01MoE(nn.Module):

223 +

224 + def __init__(

225 + self,

226 + num_experts: int,

227 + top_k: int,

228 + hidden_size: int,

229 + intermediate_size: int,

230 + params_dtype: Optional[torch.dtype] = None,

231 + layer_idx: int = None,

232 + quant_config: Optional[QuantizationConfig] = None,

233 + prefix: str = "moe",

234 + ) -> None:

235 + super().__init__()

236 +

237 + self.layer_idx = layer_idx

238 + self.tp_size = get_tensor_model_parallel_world_size()

239 + self.num_total_experts = num_experts

240 + self.top_k = top_k

241 + self.hidden_size = hidden_size

242 + self.intermediate_size = intermediate_size // self.tp_size

243 + self.quant_config = quant_config

244 +

245 + if params_dtype is None:

246 + params_dtype = torch.get_default_dtype()

247 + self.params_dtype = params_dtype

248 +

249 + self.gate = ReplicatedLinear(

250 + self.hidden_size,

251 + self.num_total_experts,

252 + bias=False,

253 + params_dtype=torch.float32,

254 + quant_config=None,

255 + prefix=f"{prefix}.gate",

256 + )

257 + self.gate.weight.weight_loader = MiniMaxText01MoE.gate_weight_loader

258 +

259 + self.experts = FusedMoE(

260 + num_experts=self.num_total_experts,

261 + top_k=self.top_k,

262 + hidden_size=self.hidden_size,

263 + intermediate_size=self.intermediate_size * self.tp_size,

264 + params_dtype=self.params_dtype,

265 + reduce_results=True,

266 + renormalize=True,

267 + quant_config=self.quant_config,

268 + tp_size=self.tp_size,

269 + prefix=f"{prefix}.experts",

270 + )

271 + return

272 +

273 + @staticmethod

274 + def gate_weight_loader(param: nn.Parameter,

275 + loaded_weight: torch.Tensor) -> None:

276 + assert param.size() == loaded_weight.size()

277 + param.data.copy_(loaded_weight.to(torch.float32))

278 + return

279 +

280 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

281 + num_tokens, hidden_size = hidden_states.shape

282 + hidden_states = hidden_states.view(-1, self.hidden_size)

283 + router_logits_fp32, _ = self.gate(hidden_states.to(torch.float32))

284 + final_hidden_states = self.experts(

285 + hidden_states, router_logits_fp32.to(hidden_states.dtype))

286 + final_hidden = final_hidden_states.view(num_tokens, hidden_size)

287 + return final_hidden

288 +

289 +

290 +class MiniMaxText01LinearKernel:

291 +

292 + @staticmethod

293 + def jit_linear_forward_prefix(q: torch.Tensor,

294 + k: torch.Tensor,

295 + v: torch.Tensor,

296 + kv_caches: torch.Tensor,

297 + slope_rate: torch.Tensor,

298 + block_size: int,

299 + layer_idx: int = None,

300 + **kwargs) -> torch.Tensor:

301 +

302 + slope_rate = slope_rate.to(torch.float32)

303 + should_pad_dim = q.dim() == 3

304 + if should_pad_dim:

305 + q = q.unsqueeze(0)

306 + k = k.unsqueeze(0)

307 + v = v.unsqueeze(0)

308 + b, h, n, d = q.shape

309 + e = d

310 + kv_history = kv_caches.reshape(1, h, d, e).contiguous()

311 + output, kv_history = lightning_attention(q,

312 + k,

313 + v,

314 + slope_rate,

315 + block_size=block_size,

316 + kv_history=kv_history)

317 + kv_caches.copy_(kv_history[:, :, -1, :, :].reshape(h, d, e))

318 + assert output.shape[0] == 1, "batch size must be 1"

319 + return rearrange(output.squeeze(0), "h n d -> n (h d)")

320 +

321 +

322 +class MiniMaxText01LinearAttention(nn.Module):

323 +

324 + def __init__(

325 + self,

326 + hidden_size: int,

327 + hidden_inner_size: int,

328 + num_heads: int,

329 + head_dim: int,

330 + max_position: int,

331 + block_size: int,

332 + num_hidden_layer: int,

333 + quant_config: Optional[QuantizationConfig] = None,

334 + layer_idx: int = 0,

335 + linear_layer_idx: int = 0,

336 + prefix: str = "linear_attn",

337 + ) -> None:

338 + super().__init__()

339 +

340 + self.layer_idx = layer_idx

341 + self.BLOCK = block_size

342 + self.hidden_size = hidden_size

343 + self.num_heads = num_heads

344 + self.head_dim = head_dim

345 + self.total_num_heads = num_heads

346 + self.hidden_inner_size = hidden_inner_size

347 + self.tp_size = get_tensor_model_parallel_world_size()

348 + self.tp_rank = get_tensor_model_parallel_rank()

349 +

350 + assert self.total_num_heads % self.tp_size == 0

351 + self.tp_heads = self.total_num_heads // self.tp_size

352 + self.qkv_size = self.num_heads * self.head_dim

353 + self.tp_hidden = self.head_dim * self.tp_heads

354 +

355 + self.qkv_proj = ColumnParallelLinear(

356 + hidden_size,

357 + self.hidden_inner_size * 3,

358 + bias=False,

359 + quant_config=quant_config,

360 + prefix=f"{prefix}.qkv_proj",

361 + )

362 + self.output_gate = ColumnParallelLinear(

363 + hidden_size,

364 + self.hidden_inner_size,

365 + bias=False,

366 + quant_config=quant_config,

367 + prefix=f"{prefix}.output_gate",

368 + )

369 + self.out_proj = RowParallelLinear(

370 + self.hidden_inner_size,

371 + hidden_size,

372 + bias=False,

373 + quant_config=quant_config,

374 + prefix=f"{prefix}.out_proj",

375 + )

376 + self.norm = MiniMaxText01RMSNormTP(

377 + self.hidden_inner_size,

378 + eps=1e-5,

379 + )

380 +

381 + slope_rate = MiniMaxText01LinearAttention._build_slope_tensor(

382 + self.num_heads)

383 + if num_hidden_layer <= 1:

384 + self.slope_rate = slope_rate * (1 + 1e-5)

385 + else:

386 + self.slope_rate = slope_rate * (1 - layer_idx /

387 + (num_hidden_layer - 1) + 1e-5)

388 + self.tp_slope = self.slope_rate[self.tp_rank *

389 + self.tp_heads:(self.tp_rank + 1) *

390 + self.tp_heads].contiguous()

391 +

392 + @staticmethod

393 + def weight_direct_load(param: torch.Tensor,

394 + loaded_weight: torch.Tensor) -> None:

395 + assert param.size() == loaded_weight.size()

396 + param.data.copy_(loaded_weight)

397 + return

398 +

399 + @staticmethod

400 + def _build_slope_tensor(n_attention_heads: int):

401 +

402 + def get_slopes(n):

403 +

404 + def get_slopes_power_of_2(n):

405 + start = 2**(-(2**-(math.log2(n) - 3)))

406 + ratio = start

407 + return [start * ratio**i for i in range(n)]

408 +

409 + if math.log2(n).is_integer():

410 + return get_slopes_power_of_2(n)

411 + else:

412 + closest_power_of_2 = 2**math.floor(math.log2(n))

413 + return (get_slopes_power_of_2(closest_power_of_2) + get_slopes(

414 + 2 * closest_power_of_2)[0::2][:n - closest_power_of_2])

415 +

416 + slopes = torch.tensor(get_slopes(n_attention_heads),

417 + dtype=torch.float32).reshape(

418 + n_attention_heads, 1, 1)

419 + return slopes

420 +

421 + def _prefill_and_mix_infer(self, q, k, v, kv_cache, state_indices_tensor,

422 + attn_metadata):

423 + hidden = []

424 + for _prefill_idx in range(getattr(attn_metadata, "num_prefills", 0)):

425 + _start = attn_metadata.query_start_loc[_prefill_idx]

426 + _end = attn_metadata.query_start_loc[_prefill_idx + 1]

427 + slot_id = state_indices_tensor[_prefill_idx]

428 + qs = q[_start:_end].transpose(0, 1).contiguous()

429 + ks = k[_start:_end].transpose(0, 1).contiguous()

430 + vs = v[_start:_end].transpose(0, 1).contiguous()

431 + slot_id = state_indices_tensor[_prefill_idx]

432 + slice_layer_cache = kv_cache[slot_id, ...]

433 +

434 + out_slice = MiniMaxText01LinearKernel.jit_linear_forward_prefix(

435 + qs,

436 + ks,

437 + vs,

438 + slice_layer_cache,

439 + self.tp_slope,

440 + self.BLOCK,

441 + layer_idx=self.layer_idx)

442 + hidden.append(out_slice.contiguous())

443 + if attn_metadata.num_decode_tokens > 0:

444 + hidden.append(

445 + self._decode_infer(q, k, v, kv_cache, state_indices_tensor,

446 + attn_metadata))

447 + hidden = torch.concat(hidden, dim=0).contiguous()

448 + return hidden

449 +

450 + def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor,

451 + attn_metadata):

452 + q = q[attn_metadata.num_prefill_tokens:].unsqueeze(2).contiguous()

453 + k = k[attn_metadata.num_prefill_tokens:].unsqueeze(2).contiguous()

454 + v = v[attn_metadata.num_prefill_tokens:].unsqueeze(2).contiguous()

455 + slot_id = state_indices_tensor[getattr(attn_metadata, "num_prefills", 0

456 + ):]

457 + hidden = linear_decode_forward_triton(q, k, v, kv_cache, self.tp_slope,

458 + slot_id, 32)

459 + return hidden

460 +

461 + def forward(self, hidden_states: torch.Tensor, positions: torch.Tensor,

462 + kv_caches: MinimaxCacheParams, **kwargs) -> torch.Tensor:

463 + qkv, _ = self.qkv_proj(hidden_states)

464 + qkv32 = qkv.to(torch.float32)

465 + qkvact = torch.nn.functional.silu(qkv32)

466 + qkvact = qkvact.view((qkv.shape[0], self.tp_heads, -1))

467 + q, k, v = torch.split(qkvact, [self.head_dim] * 3, dim=-1)

468 + forward_context = get_forward_context()

469 + attn_metadata = forward_context.attn_metadata

470 + kv_cache = kv_caches.minimax_cache

471 + state_indices_tensor = kv_caches.state_indices_tensor

472 +

473 + decode_only = getattr(attn_metadata, "num_prefills", 0) == 0

474 + if not decode_only:

475 + hidden = self._prefill_and_mix_infer(q, k, v, kv_cache,

476 + state_indices_tensor,

477 + attn_metadata)

478 + else:

479 + hidden = self._decode_infer(q, k, v, kv_cache,

480 + state_indices_tensor, attn_metadata)

481 +

482 + hidden = self.norm._forward(hidden)

483 + gate, _ = self.output_gate(hidden_states)

484 + hidden = F.sigmoid(gate) * hidden

485 + hidden = hidden.to(hidden_states.dtype)

486 + hidden, _ = self.out_proj(hidden)

487 + return hidden

488 +

489 +

490 +class MiniMaxText01Attention(nn.Module):

491 +

492 + def __init__(

493 + self,

494 + hidden_size: int,

495 + num_heads: int,

496 + head_dim: int,

497 + num_kv_heads: int,

498 + rotary_dim: int,

499 + max_position: int = 4096 * 32,

500 + rope_theta: float = 10000,

501 + sliding_window: Optional[int] = None,

502 + quant_config: Optional[QuantizationConfig] = None,

503 + layer_idx: int = None,

504 + cache_config: Optional[CacheConfig] = None,

505 + prefix: str = "mha",

506 + ) -> None:

507 + super().__init__()

508 + self.layer_idx = layer_idx

509 +

510 + self.hidden_size = hidden_size

511 + tp_size = get_tensor_model_parallel_world_size()

512 + self.total_num_heads = num_heads

513 + assert self.total_num_heads % tp_size == 0

514 + self.num_heads = self.total_num_heads // tp_size

515 + self.total_num_kv_heads = num_kv_heads

516 + if self.total_num_kv_heads >= tp_size:

517 + assert self.total_num_kv_heads % tp_size == 0

518 + else:

519 + assert tp_size % self.total_num_kv_heads == 0

520 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

521 + self.head_dim = head_dim

522 +

523 + self.q_size = self.num_heads * self.head_dim

524 + self.kv_size = self.num_kv_heads * self.head_dim

525 + self.scaling = self.head_dim**-0.5

526 + self.rope_theta = rope_theta

527 + self.sliding_window = sliding_window

528 +

529 + self.qkv_proj = QKVParallelLinear(

530 + hidden_size,

531 + self.head_dim,

532 + self.total_num_heads,

533 + self.total_num_kv_heads,

534 + bias=False,

535 + quant_config=quant_config,

536 + prefix=f"{prefix}.qkv_proj",

537 + )

538 + self.o_proj = RowParallelLinear(

539 + self.total_num_heads * self.head_dim,

540 + hidden_size,

541 + bias=False,

542 + quant_config=quant_config,

543 + prefix=f"{prefix}.o_proj",

544 + )

545 + self.attn = Attention(

546 + self.num_heads,

547 + self.head_dim,

548 + self.scaling,

549 + num_kv_heads=self.num_kv_heads,

550 + cache_config=cache_config,

551 + quant_config=quant_config,

552 + prefix=f"{prefix}.attn",

553 + )

554 + return

555 +

556 + def forward(self, hidden_states: torch.Tensor, positions: torch.Tensor,

557 + **kwargs) -> torch.Tensor:

558 + forward_context = get_forward_context()

559 + attn_metadata = forward_context.attn_metadata

560 + qkv, _ = self.qkv_proj(hidden_states)

561 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

562 + q, k = attn_metadata.rotary_emb(positions, q, k)

563 + attn_output = self.attn(q, k, v)

564 + output, _ = self.o_proj(attn_output)

565 + return output

566 +

567 +

568 +class MiniMaxText01DecoderLayer(nn.Module):

569 +

570 + def __init__(

571 + self,

572 + config: PretrainedConfig,

573 + cache_config: Optional[CacheConfig] = None,

574 + quant_config: Optional[QuantizationConfig] = None,

575 + expert_num: int = 1,

576 + layer_id: int = None,

577 + linear_layer_id: Optional[int] = None,

578 + prefix: str = "decoder",

579 + ) -> None:

580 + self._ilayer = layer_id

581 + self._irank = get_tensor_model_parallel_rank()

582 + super().__init__()

583 +

584 + self.hidden_size = config.hidden_size

585 + self.expert_num = expert_num

586 +

587 + rope_theta = getattr(config, "rope_theta", 10000)

588 +

589 + head_dim = getattr(config, "head_dim",

590 + config.hidden_size // config.num_attention_heads)

591 + if hasattr(config, "max_model_len") and isinstance(

592 + config.max_model_len, int):

593 + max_position_embeddings = min(config.max_position_embeddings,

594 + config.max_model_len)

595 + if config.attention_type == 0:

596 + use_headxdim = True

597 + hidden_inner = (head_dim * config.num_attention_heads

598 + if use_headxdim else config.hidden_size)

599 + self.self_attn = MiniMaxText01LinearAttention(

600 + hidden_size=self.hidden_size,

601 + hidden_inner_size=hidden_inner,

602 + num_heads=config.num_attention_heads,

603 + head_dim=head_dim,

604 + max_position=max_position_embeddings,

605 + block_size=config.block if hasattr(config, "block") else 256,

606 + num_hidden_layer=config.num_hidden_layers,

607 + quant_config=quant_config,

608 + layer_idx=self._ilayer,

609 + linear_layer_idx=linear_layer_id,

610 + prefix=prefix)

611 + elif config.attention_type == 1:

612 + self.self_attn = MiniMaxText01Attention(

613 + hidden_size=self.hidden_size,

614 + num_heads=config.num_attention_heads,

615 + head_dim=head_dim,

616 + rotary_dim=config.rotary_dim

617 + if hasattr(config, "rotary_dim") else head_dim,

618 + num_kv_heads=config.num_key_value_heads,

619 + max_position=max_position_embeddings,

620 + rope_theta=rope_theta,

621 + sliding_window=config.sliding_window,

622 + quant_config=quant_config,

623 + layer_idx=self._ilayer,

624 + cache_config=cache_config,

625 + prefix=prefix)

626 + else:

627 + raise ValueError(

628 + f"Unsupported attention type: {self.config.attention_type}")

629 +

630 + if expert_num == 1:

631 + self.mlp = MiniMaxText01MLP(

632 + hidden_size=self.hidden_size,

633 + intermediate_size=config.intermediate_size,

634 + quant_config=quant_config,

635 + layer_idx=self._ilayer,

636 + prefix=prefix)

637 + else:

638 + self.block_sparse_moe = MiniMaxText01MoE(

639 + num_experts=expert_num,

640 + top_k=config.num_experts_per_tok,

641 + hidden_size=config.hidden_size,

642 + intermediate_size=config.intermediate_size,

643 + layer_idx=self._ilayer,

644 + quant_config=quant_config,

645 + prefix=prefix)

646 +

647 + self.input_layernorm = RMSNorm(config.hidden_size,

648 + eps=config.rms_norm_eps)

649 + self.post_attention_layernorm = RMSNorm(config.hidden_size,

650 + eps=config.rms_norm_eps)

651 + if config.attention_type == 0:

652 + self.layernorm_attention_alpha = getattr(

653 + config, 'layernorm_linear_attention_alpha', 1)

654 + self.layernorm_attention_beta = getattr(

655 + config, 'layernorm_linear_attention_beta', 1)

656 + else:

657 + self.layernorm_attention_alpha = getattr(

658 + config, 'layernorm_full_attention_alpha', 1)

659 + self.layernorm_attention_beta = getattr(

660 + config, 'layernorm_full_attention_beta', 1)

661 + self.layernorm_mlp_alpha = getattr(config, 'layernorm_mlp_alpha', 1)

662 + self.layernorm_mlp_beta = getattr(config, 'layernorm_mlp_beta', 1)

663 + self.postnorm = getattr(config, 'postnorm', False)

664 + self.shared_moe = False

665 +

666 + shared_intermediate = getattr(config, 'shared_intermediate_size', 0)

667 + if shared_intermediate > 0:

668 + self.shared_moe = True

669 + self.shared_mlp = MiniMaxText01MLP(

670 + hidden_size=self.hidden_size,

671 + intermediate_size=shared_intermediate,

672 + quant_config=quant_config,

673 + layer_idx=self._ilayer,

674 + prefix=prefix)

675 + self.coefficient = ReplicatedLinear(

676 + self.hidden_size,

677 + 1,

678 + bias=False,

679 + quant_config=quant_config,

680 + params_dtype=torch.float32,

681 + )

682 + self.coefficient.weight.weight_loader = (

683 + self.shared_moe_coefficient_loader)

684 + self.shared_moe_mode = getattr(config, 'shared_moe_mode',

685 + 'softmax')

686 + return

687 +

688 + def forward(self,

689 + hidden_states: torch.Tensor,

690 + positions: torch.Tensor,

691 + kv_caches: Union[List[Dict], Optional[torch.Tensor]],

692 + attn_metadata: AttentionMetadata,

693 + residual: Optional[torch.Tensor],

694 + is_warmup: bool = False,

695 + **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:

696 +

697 + forward_context = get_forward_context()

698 + attn_metadata = forward_context.attn_metadata

699 + layernorm_input = hidden_states

700 + layernorm_output = self.input_layernorm(layernorm_input)

701 + residual = layernorm_output if self.postnorm else layernorm_input

702 + self_attention_output = self.self_attn(

703 + hidden_states=layernorm_output,

704 + positions=positions,

705 + kv_caches=kv_caches,

706 + attn_metadata=attn_metadata,

707 + )

708 +

709 + residual = residual * self.layernorm_attention_alpha

710 + self_attention_output = (self_attention_output *

711 + self.layernorm_attention_beta)

712 +

713 + layernorm_input = residual + self_attention_output

714 + layernorm_output = self.post_attention_layernorm(layernorm_input)

715 + residual = layernorm_output if self.postnorm else layernorm_input

716 +

717 + if self.expert_num == 1:

718 + hidden_states = self.mlp(layernorm_output)

719 + else:

720 + moe_hidden_states = self.block_sparse_moe(

721 + copy.deepcopy(layernorm_output))

722 + if self.shared_moe:

723 + before_moe_dtype = layernorm_output.dtype

724 + moe_hidden_fp32 = moe_hidden_states.to(torch.float32)

725 + output_mlp = self.shared_mlp(layernorm_output).to(

726 + torch.float32)

727 +

728 + coef, _ = self.coefficient(layernorm_output.to(torch.float32))

729 +

730 + if self.shared_moe_mode == 'softmax':

731 + coef = torch.nn.functional.softmax(coef, dim=-1)

732 + hidden_states = moe_hidden_fp32 * (

733 + 1 - coef) + output_mlp * coef

734 + elif self.shared_moe_mode == 'sigmoid':

735 + coef = torch.nn.functional.sigmoid(coef)

736 + hidden_states = moe_hidden_fp32 * (

737 + 1 - coef) + output_mlp * coef

738 +

739 + hidden_states = hidden_states.to(before_moe_dtype)

740 + else:

741 + hidden_states = moe_hidden_states

742 +

743 + residual = residual * self.layernorm_mlp_alpha

744 + hidden_states = hidden_states * self.layernorm_mlp_beta

745 +

746 + hidden_states = residual + hidden_states

747 +

748 + return hidden_states, None

749 +

750 + @staticmethod

751 + def shared_moe_coefficient_loader(param: torch.Tensor,

752 + loaded_weight: torch.Tensor) -> None:

753 + assert param.size() == loaded_weight.size()

754 +

755 + param.data.copy_(loaded_weight.to(torch.float32))

756 + return

757 +

758 +

759 +class MiniMaxText01Model(nn.Module):

760 +

761 + def __init__(

762 + self,

763 + config: PretrainedConfig,

764 + quant_config: Optional[QuantizationConfig] = None,

765 + cache_config: Optional[CacheConfig] = None,

766 + scheduler_config=None,

767 + prefix: str = "",

768 + ) -> None:

769 + super().__init__()

770 +

771 + self.padding_idx = config.pad_token_id

772 + self.vocab_size = config.vocab_size

773 +

774 + self.decoder_attention_types = getattr(

775 + config, "attn_type_list", False) or getattr(

776 + config, "decoder_attention_types", False)

777 + if not self.decoder_attention_types:

778 + self.decoder_attention_types = [1] * config.num_hidden_layers

779 + self.num_layers = config.num_hidden_layers

780 +

781 + self._layer_barrier = False

782 + if get_pp_group().is_first_rank:

783 + self.embed_tokens = VocabParallelEmbedding(

784 + self.vocab_size,

785 + config.hidden_size,

786 + org_num_embeddings=self.vocab_size,

787 + )

788 + else:

789 + self.embed_tokens = PPMissingLayer()

790 +

791 + def layer_fn(prefix):

792 + layer_idx = int(prefix.split('.')[-1])

793 + layer_config = config

794 + layer_config.attention_type = self.decoder_attention_types[

795 + layer_idx]

796 + layer_config.layer_idx = layer_idx

797 +

798 + decoder_kwargs = {

799 + "quant_config": quant_config,

800 + "layer_id": layer_idx,

801 + "cache_config": cache_config

802 + }

803 +

804 + if layer_config.attention_type == 0:

805 + decoder_kwargs["linear_layer_id"] = sum(

806 + 1 for i in range(layer_idx)

807 + if self.decoder_attention_types[i] == 0)

808 + else:

809 + decoder_kwargs["linear_layer_id"] = None

810 +

811 + if hasattr(config, "num_local_experts") and isinstance(

812 + config.num_local_experts, list):

813 + decoder_kwargs["expert_num"] = config.num_local_experts[

814 + layer_idx]

815 + elif hasattr(config, "num_local_experts") and isinstance(

816 + config.num_local_experts, int):

817 + decoder_kwargs["expert_num"] = config.num_local_experts

818 + else:

819 + decoder_kwargs["expert_num"] = 1

820 +

821 + return MiniMaxText01DecoderLayer(layer_config,

822 + **decoder_kwargs,

823 + prefix=prefix)

824 +

825 + self.start_layer, self.end_layer, self.layers = make_layers(

826 + config.num_hidden_layers, layer_fn, prefix=f"{prefix}.layers")

827 +

828 + linear_layer_nums = sum(1 for i in range(config.num_hidden_layers)

829 + if self.decoder_attention_types[i] == 0)

830 + max_slots_number = scheduler_config.max_num_seqs

831 + self.cache_shape = (linear_layer_nums, max_slots_number,

832 + config.num_attention_heads //

833 + get_tensor_model_parallel_world_size(),

834 + config.head_dim, config.head_dim)

835 + _dummy = torch.zeros(1)

836 + self._dtype = _dummy.dtype

837 + del _dummy

838 +

839 + self.minimax_cache = MinimaxCacheManager(dtype=self._dtype,

840 + cache_shape=self.cache_shape)

841 +

842 + rope_theta = getattr(config, "rope_theta", 10000)

843 + head_dim = getattr(config, "head_dim",

844 + config.hidden_size // config.num_attention_heads)

845 + if hasattr(config, "max_model_len") and isinstance(

846 + config.max_model_len, int):

847 + max_position_embeddings = min(config.max_position_embeddings,

848 + config.max_model_len)

849 + self.rotary_emb = MiniMaxText01RotaryEmbedding(

850 + head_dim,

851 + rotary_dim=config.rotary_dim

852 + if hasattr(config, "rotary_dim") else head_dim,

853 + max_position=max_position_embeddings,

854 + base=int(rope_theta),

855 + is_neox_style=True,

856 + cache_dtype=torch.float32,

857 + )

858 +

859 + norm_kwargs = {}

860 + if hasattr(config, "rms_norm_eps"):

861 + norm_kwargs["eps"] = config.rms_norm_eps

862 + if get_pp_group().is_last_rank:

863 + self.norm = RMSNorm(config.hidden_size, **norm_kwargs)

864 + else:

865 + self.norm = PPMissingLayer()

866 + self.embed_scale = 1.0

867 + return

868 +

869 + def _clear_prefill_cache(self, attn_metadata,

870 + minimax_cache_tensors: torch.Tensor, **kwargs):

871 + seq_to_slot_maps = {}

872 + seq_id_map = sum(list(kwargs["request_ids_to_seq_ids"].values()), [])

873 + for _, seq_to_slot_map in (

874 + self.minimax_cache.cache_indices_mapping.items()):

875 + seq_to_slot_maps.update(seq_to_slot_map)

876 +

877 + slots_to_clear = []

878 + for _prefill_id in range(getattr(attn_metadata, "num_prefills", 0)):

879 + seq_id = seq_id_map[_prefill_id]

880 + if attn_metadata.context_lens_tensor[

881 + _prefill_id] == 0 and seq_id in seq_to_slot_maps:

882 + slots_to_clear.append(seq_to_slot_maps[seq_id])

883 +

884 + if slots_to_clear:

885 + slots_tensor = torch.tensor(slots_to_clear,

886 + device=minimax_cache_tensors.device,

887 + dtype=torch.long)

888 + minimax_cache_tensors[:, slots_tensor, ...] = 0

889 +

890 + def forward(self,

891 + input_ids: Optional[torch.Tensor],

892 + positions: torch.Tensor,

893 + kv_caches: List[torch.Tensor],

894 + intermediate_tensors=None,

895 + inputs_embeds: Optional[torch.Tensor] = None,

896 + **kwargs) -> torch.Tensor:

897 + forward_context = get_forward_context()

898 + attn_metadata = forward_context.attn_metadata

899 + if attn_metadata is None:

900 + return None

901 + if "request_ids_to_seq_ids" not in kwargs:

902 + kwargs["request_ids_to_seq_ids"] = {}

903 + if "finished_requests_ids" not in kwargs:

904 + kwargs["finished_requests_ids"] = []

905 + (

906 + minimax_cache_tensors,

907 + state_indices_tensor,

908 + ) = self.minimax_cache.current_run_tensors(**kwargs)

909 + if getattr(attn_metadata, "num_prefills", 0) > 0:

910 + self._clear_prefill_cache(attn_metadata, minimax_cache_tensors,

911 + **kwargs)

912 +

913 + minimax_cache_params = MinimaxCacheParams(minimax_cache_tensors,

914 + state_indices_tensor)

915 + if get_pp_group().is_first_rank:

916 + if inputs_embeds is None:

917 + hidden_states = self.embed_scale * self.embed_tokens(input_ids)

918 + else:

919 + hidden_states = inputs_embeds

920 + residual = None

921 + else:

922 + assert intermediate_tensors is not None

923 + hidden_states = intermediate_tensors["hidden_states"]

924 + residual = intermediate_tensors["residual"]

925 +

926 + kv_cache_index = 0

927 + minimax_cache_index = 0

928 + attn_metadata.rotary_emb = self.rotary_emb

929 + for i in range(self.start_layer, self.end_layer):

930 + layer = self.layers[i]

931 + _caches = None

932 + if isinstance(layer.self_attn, MiniMaxText01Attention):

933 + _caches = kv_caches[kv_cache_index]

934 + kv_cache_index += 1

935 + if isinstance(layer.self_attn, MiniMaxText01LinearAttention):

936 + current_state_layer = minimax_cache_index

937 + _caches = minimax_cache_params.at_layer_idx(

938 + current_state_layer)

939 + minimax_cache_index += 1

940 + hidden_states, residual = layer(

941 + hidden_states=hidden_states,

942 + positions=positions,

943 + kv_caches=_caches,

944 + attn_metadata=attn_metadata,

945 + residual=residual,

946 + )

947 + if not get_pp_group().is_last_rank:

948 + return IntermediateTensors({

949 + "hidden_states": hidden_states,

950 + "residual": residual

951 + })

952 + if residual is not None:

953 + hidden_states, _ = self.norm(hidden_states, residual)

954 + else:

955 + hidden_states = self.norm(hidden_states)

956 +

957 + return hidden_states

958 +

959 +

960 +class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid,

961 + SupportsV0Only):

962 +

963 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

964 +

965 + super().__init__()

966 + config = vllm_config.model_config.hf_config

967 + quant_config = vllm_config.quant_config

968 + lora_config = vllm_config.lora_config

969 + self.config = config

970 + self.lora_config = lora_config

971 +

972 + if not hasattr(config, "sliding_window"):

973 + config.sliding_window = None

974 +

975 + self.CONCAT_FFN = True

976 +

977 + self.unpadded_vocab_size = self.config.vocab_size

978 + if hasattr(vllm_config.model_config, "max_model_len"):

979 + self.config.max_model_len = vllm_config.model_config.max_model_len

980 + self.model = MiniMaxText01Model(

981 + self.config,

982 + quant_config,

983 + cache_config=vllm_config.cache_config,

984 + scheduler_config=vllm_config.scheduler_config,

985 + prefix=maybe_prefix(prefix, "model"))

986 + if get_pp_group().is_last_rank:

987 + self.lm_head = ParallelLMHead(

988 + self.unpadded_vocab_size,

989 + self.config.hidden_size,

990 + org_num_embeddings=self.config.vocab_size,

991 + padding_size=DEFAULT_VOCAB_PADDING_SIZE,

992 + )

993 +

994 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

995 + self.config.vocab_size)

996 +

997 + self.sampler = Sampler()

998 + else:

999 + self.lm_head = PPMissingLayer()

1000 +

1001 + flash_layer_count = sum(1 for attn_type in self.config.attn_type_list

1002 + if attn_type == 1)

1003 + self.kv_cache = [torch.tensor([]) for _ in range(flash_layer_count)]

1004 + return

1005 +

1006 + def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):

1007 + return self.model.minimax_cache.copy_inputs_before_cuda_graphs(

1008 + input_buffers, **kwargs)

1009 +

1010 + def get_seqlen_agnostic_capture_inputs(self, batch_size: int):

1011 + return self.model.minimax_cache.get_seqlen_agnostic_capture_inputs(

1012 + batch_size)

1013 +

1014 + def forward(self,

1015 + input_ids: torch.Tensor,

1016 + positions: torch.Tensor,

1017 + intermediate_tensors: Optional[IntermediateTensors] = None,

1018 + inputs_embeds: Optional[torch.Tensor] = None,

1019 + **kwargs) -> torch.Tensor:

1020 + hidden_states = self.model(input_ids, positions, self.kv_cache,

1021 + intermediate_tensors, inputs_embeds,

1022 + **kwargs)

1023 +

1024 + return hidden_states

1025 +

1026 + def compute_logits(self, hidden_states: torch.Tensor,

1027 + sampling_metadata: SamplingMetadata) -> torch.Tensor:

1028 + logits = self.logits_processor(self.lm_head, hidden_states,

1029 + sampling_metadata)

1030 +

1031 + return logits

1032 +

1033 + def sample(

1034 + self,

1035 + logits: torch.Tensor,

1036 + sampling_metadata: SamplingMetadata,

1037 + ):

1038 +

1039 + next_tokens = self.sampler(logits, sampling_metadata)

1040 +

1041 + return next_tokens

1042 +

1043 + def make_empty_intermediate_tensors(

1044 + self, batch_size: int, dtype: torch.dtype,

1045 + device: torch.device) -> IntermediateTensors:

1046 + return IntermediateTensors({

1047 + "hidden_states":

1048 + torch.zeros((batch_size, self.config.hidden_size),

1049 + dtype=dtype,

1050 + device=device),

1051 + "residual":

1052 + torch.zeros((batch_size, self.config.hidden_size),

1053 + dtype=dtype,

1054 + device=device),

1055 + })

1056 +

1057 + def load_weights(self, weights: Iterable[Tuple[str,

1058 + torch.Tensor]]) -> None:

1059 + params_dict = dict(self.named_parameters())

1060 +

1061 + def which_layer(name: str) -> int:

1062 + if "layers" in name:

1063 + after_layer = name.split("layers")[-1]

1064 + return int(after_layer.split(".")[1])

1065 + return None

1066 +

1067 + def is_linear_attn_layer(layer_idx: int) -> bool:

1068 + if layer_idx is None or not hasattr(self.config, "attn_type_list"):

1069 + return False

1070 + return self.config.attn_type_list[layer_idx] == 0

1071 +

1072 + def is_moe_weight(name: str) -> bool:

1073 + return "block_sparse_moe" in name and not name.endswith(".bias")

1074 +

1075 + def get_expert_id(param_name):

1076 + pattern = r'model\.layers\.\d+\.block_sparse_moe\.experts\.(\d+)\.'

1077 + match = re.search(pattern, param_name)

1078 + if match:

1079 + return match.group(1)

1080 + return None

1081 +

1082 + def load_sparse_moe_weight(name: str, loaded_weight: torch.Tensor,

1083 + self) -> None:

1084 + if isinstance(self.config.num_local_experts, list):

1085 + expert_params_mapping = [

1086 + ("w13_weight"

1087 + if weight_name in ["w1", "w3"] else "w2_weight",

1088 + f"experts.{expert_id}.{weight_name}.weight", expert_id)

1089 + for expert_id in range(max(self.config.num_local_experts))

1090 + for weight_name in ["w1", "w2", "w3"]

1091 + ]

1092 + else:

1093 + expert_params_mapping = [

1094 + ("w13_scale" if weight_name in ["w1", "w3"] else

1095 + "w2_scale", f"{expert_id}.{weight_name}.weight_scale",

1096 + expert_id, weight_name)

1097 + for expert_id in range(self.config.num_local_experts)

1098 + for weight_name in ["w1", "w2", "w3"]

1099 + ] + [("w13_weight" if weight_name in ["w1", "w3"] else

1100 + "w2_weight", f"{expert_id}.{weight_name}.weight",

1101 + expert_id, weight_name)

1102 + for expert_id in range(self.config.num_local_experts)

1103 + for weight_name in ["w1", "w2", "w3"]]

1104 + for (param_name, weight_name, expert_id,

1105 + shard_id) in expert_params_mapping:

1106 + name_expert_id = get_expert_id(name)

1107 + if name_expert_id is not None and int(name_expert_id) != int(

1108 + expert_id):

1109 + continue

1110 + if weight_name not in name:

1111 + continue

1112 + name = name.replace(weight_name, param_name)

1113 + if is_pp_missing_parameter(name, self):

1114 + return

1115 + param = params_dict[name]

1116 + weight_loader = param.weight_loader

1117 + weight_loader = weight_loader_with_alias(name)(weight_loader)

1118 + weight_loader(param,

1119 + loaded_weight,

1120 + weight_name,

1121 + expert_id=expert_id,

1122 + shard_id=shard_id)

1123 + break

1124 + else:

1125 + if is_pp_missing_parameter(name, self):

1126 + return

1127 + param = params_dict[name]

1128 + weight_loader = getattr(param, "weight_loader",

1129 + default_weight_loader)

1130 + weight_loader = weight_loader_with_alias(name)(weight_loader)

1131 + weight_loader(param, loaded_weight)

1132 + return

1133 +

1134 + def is_shared_mlp_weight(name: str) -> bool:

1135 + return "shared_mlp" in name and not name.endswith(".bias")

1136 +

1137 + def load_shared_mlp_weight(name: str, loaded_weight: torch.Tensor,

1138 + self) -> None:

1139 + if not self.CONCAT_FFN:

1140 + if "gate_proj" in name:

1141 + name = name.replace("gate_proj", "w1", 1)

1142 + elif "up_proj" in name:

1143 + name = name.replace("up_proj", "w3", 1)

1144 + elif "down_proj" in name:

1145 + name = name.replace("down_proj", "w2", 1)

1146 + else:

1147 + if "gate_proj" in name:

1148 + name = name.replace("gate_proj", "gate_up_proj", 1)

1149 + loaded_shard_id = 0

1150 + elif "up_proj" in name:

1151 + name = name.replace("up_proj", "gate_up_proj", 1)

1152 + loaded_shard_id = 1

1153 + if is_pp_missing_parameter(name, self):

1154 + return

1155 + param = params_dict[name]

1156 + weight_loader = getattr(param, "weight_loader",

1157 + default_weight_loader)

1158 + weight_loader = weight_loader_with_alias(name)(weight_loader)

1159 + if not self.CONCAT_FFN:

1160 + weight_loader(param, loaded_weight)

1161 + else:

1162 + if "gate_up_proj" in name:

1163 + weight_loader(param, loaded_weight, loaded_shard_id)

1164 + elif "down_proj" in name:

1165 + weight_loader(param, loaded_weight)

1166 + else:

1167 + raise AssertionError(

1168 + "MLP weight not in [gate_up_proj, down_proj]")

1169 + return

1170 +

1171 + def is_mha_weight(name: str) -> bool:

1172 + return "self_attn" in name and not name.endswith(".bias")

1173 +

1174 + def load_linear_attn_weight(name: str, loaded_weight: torch.Tensor,

1175 + self) -> None:

1176 + if is_pp_missing_parameter(name, self):

1177 + return

1178 + param = params_dict[name]

1179 +

1180 + weight_loader = getattr(

1181 + param, "weight_loader",

1182 + MiniMaxText01LinearAttention.weight_direct_load)

1183 + weight_loader = weight_loader_with_alias(name)(weight_loader)

1184 + weight_loader(param, loaded_weight)

1185 + return

1186 +

1187 + def load_flash_attn_weight(name: str, loaded_weight: torch.Tensor,

1188 + self) -> None:

1189 +

1190 + flash_mha_params_mapping = [

1191 + ("qkv_proj", "q_proj", "q"),

1192 + ("qkv_proj", "k_proj", "k"),

1193 + ("qkv_proj", "v_proj", "v"),

1194 + ("gate_up_proj", "gate_proj", 0),

1195 + ("gate_up_proj", "up_proj", 1),

1196 + ]

1197 + for (param_name, weight_name,

1198 + shard_id) in flash_mha_params_mapping:

1199 + if weight_name not in name:

1200 + continue

1201 + name = name.replace(weight_name, param_name)

1202 + if is_pp_missing_parameter(name, self):

1203 + return

1204 + param = params_dict[name]

1205 + weight_loader = getattr(param, "weight_loader",

1206 + default_weight_loader)

1207 + weight_loader = weight_loader_with_alias(name)(weight_loader)

1208 + weight_loader(param, loaded_weight, shard_id)

1209 + break

1210 + else:

1211 + if is_pp_missing_parameter(name, self):

1212 + return

1213 + param = params_dict[name]

1214 +

1215 + weight_loader = getattr(param, "weight_loader",

1216 + default_weight_loader)

1217 + weight_loader = weight_loader_with_alias(name)(weight_loader)

1218 + weight_loader(param, loaded_weight)

1219 + return

1220 +

1221 + def is_layer_norm_weight(name: str) -> bool:

1222 + return "norm" in name and not name.endswith(

1223 + ".bias") and name in params_dict

1224 +

1225 + def load_layer_norm_weight(name: str, loaded_weight: torch.Tensor,

1226 + self) -> None:

1227 + if is_pp_missing_parameter(name, self):

1228 + return

1229 + param = params_dict[name]

1230 + weight_loader = getattr(param, "weight_loader",

1231 + default_weight_loader)

1232 + weight_loader = weight_loader_with_alias(name)(weight_loader)

1233 + weight_loader(param, loaded_weight)

1234 + return

1235 +

1236 + def load_basic_weight(name: str, loaded_weight: torch.Tensor,

1237 + self) -> None:

1238 + if is_pp_missing_parameter(name, self):

1239 + return

1240 + param = params_dict[name]

1241 + weight_loader = getattr(param, "weight_loader",

1242 + default_weight_loader)

1243 + weight_loader = weight_loader_with_alias(name)(weight_loader)

1244 + weight_loader(param, loaded_weight)

1245 + return

1246 +

1247 + for name, loaded_weight in weights:

1248 + weight_at_layer = which_layer(name)

1249 + if weight_at_layer and weight_at_layer >= len(

1250 + self.config.attn_type_list):

1251 + continue

1252 +

1253 + if is_layer_norm_weight(name):

1254 + load_layer_norm_weight(name, loaded_weight, self)

1255 + continue

1256 + if is_mha_weight(name):

1257 + if is_linear_attn_layer(weight_at_layer):

1258 + load_linear_attn_weight(name, loaded_weight, self)

1259 + else:

1260 + load_flash_attn_weight(name, loaded_weight, self)

1261 + continue

1262 + if is_moe_weight(name):

1263 + load_sparse_moe_weight(name, loaded_weight, self)

1264 + continue

1265 + if is_shared_mlp_weight(name):

1266 + load_shared_mlp_weight(name, loaded_weight, self)

1267 + continue

1268 +

1269 + if "rotary_emb.inv_freq" in name:

1270 + continue

1271 +

1272 + load_basic_weight(name, loaded_weight, self)

1273 + return

vllm/model_executor/models/registry.py

@@ -35,6 +35,7 @@ _TEXT_GENERATION_MODELS = {

35 "AquilaModel": ("llama", "LlamaForCausalLM"),

36 "AquilaForCausalLM": ("llama", "LlamaForCausalLM"), # AquilaChat2

37 "ArcticForCausalLM": ("arctic", "ArcticForCausalLM"),

38 + "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"),

39 # baichuan-7b, upper case 'C' in the class name

40 "BaiChuanForCausalLM": ("baichuan", "BaiChuanForCausalLM"),

41 # baichuan-13b, lower case 'c' in the class name