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

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, 09:45 UTC

[MODEL] FalconH1 (#18406)

Signed-off-by: dhia.rhaiem <dhia.rhaiem@tii.ae> Co-authored-by: younesbelkada <younesbelkada@gmail.com> Co-authored-by: Ilyas Chahed <ilyas.chahed@tii.ae> Co-authored-by: Jingwei Zuo <jingwei.zuo@tii.ae>

Dhia Eddine Rhaiem committed May 21, 2025 at 15:59 UTC eca18691d2fe29c4f6c1b466709eda9f123116ea

5 files changed +798 -59

docs/source/models/supported_models.md

@@ -392,6 +392,11 @@ Specified using `--task generate`.

392 * `tiiuae/falcon-mamba-7b`, `tiiuae/falcon-mamba-7b-instruct`, etc.

393 * ✅︎

394 * ✅︎

395 +- * `FalconH1ForCausalLM`

396 + * Falcon-H1

397 + * `tiiuae/Falcon-H1-34B-Base`, `tiiuae/Falcon-H1-34B-Instruct`, etc.

398 + * ✅︎

399 + * ✅︎

400 - * `GemmaForCausalLM`

401 * Gemma

402 * `google/gemma-2b`, `google/gemma-1.1-2b-it`, etc.

tests/models/registry.py

@@ -147,6 +147,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

147 "ExaoneForCausalLM": _HfExamplesInfo("LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct"), # noqa: E501

148 "Fairseq2LlamaForCausalLM": _HfExamplesInfo("mgleize/fairseq2-dummy-Llama-3.2-1B"), # noqa: E501

149 "FalconForCausalLM": _HfExamplesInfo("tiiuae/falcon-7b"),

150 + "FalconH1ForCausalLM":_HfExamplesInfo("tiiuae/Falcon-H1-1.5B-Instruct",

151 + is_available_online=False,

152 + min_transformers_version="4.52.2"),

153 "GemmaForCausalLM": _HfExamplesInfo("google/gemma-1.1-2b-it"),

154 "Gemma2ForCausalLM": _HfExamplesInfo("google/gemma-2-9b"),

155 "Gemma3ForCausalLM": _HfExamplesInfo("google/gemma-3-1b-it"),

vllm/model_executor/layers/mamba/mamba_mixer2.py

+104 -59

@@ -34,7 +34,11 @@ from vllm.model_executor.utils import set_weight_attrs

34 @CustomOp.register("mixer2_gated_rms_norm")

35 class Mixer2RMSNormGated(CustomOp):

37 - def __init__(self, full_hidden_size, full_n_groups, eps=1e-6):

37 + def __init__(self,

38 + full_hidden_size: int,

39 + full_n_groups: int,

40 + use_rms_norm: bool = True,

41 + eps: float = 1e-6):

42 super().__init__()

43 self.tp_size = get_tensor_model_parallel_world_size()

44 self.tp_rank = get_tensor_model_parallel_rank()

@@ -44,11 +48,17 @@ class Mixer2RMSNormGated(CustomOp):

48 self.n_groups = full_hidden_size // self.group_size

50 self.variance_epsilon = eps

47 - self.weight = nn.Parameter(torch.ones(self.per_rank_hidden_size))

48 - set_weight_attrs(self.weight,

49 - {"weight_loader": sharded_weight_loader(0)})

50 - assert self.full_hidden_size % self.tp_size== 0,\

51 - "Tensor parallel world size must divide hidden size."

51 + self.use_rms_norm = use_rms_norm

52 + if self.use_rms_norm:

53 + # Register norm weight only if we're actually applying RMSNorm

54 + self.weight = nn.Parameter(torch.ones(self.per_rank_hidden_size))

55 + set_weight_attrs(self.weight,

56 + {"weight_loader": sharded_weight_loader(0)})

57 + else:

58 + # Avoid checkpoint mismatch by skipping unused parameter

59 + self.register_parameter("weight", None)

60 + assert (self.full_hidden_size % self.tp_size == 0

61 + ), "Tensor parallel world size must divide hidden size."

63 def forward_native(

64 self,

@@ -66,6 +76,8 @@ class Mixer2RMSNormGated(CustomOp):

76 # the input and then redundantly compute the RMSNorm.

77 input_dtype = x.dtype

78 x = x * nn.functional.silu(gate.to(torch.float32))

79 + if not self.use_rms_norm:

80 + return x

82 if self.n_groups == 1:

83 if self.tp_size > 1:

@@ -74,7 +86,7 @@ class Mixer2RMSNormGated(CustomOp):

86 global_sums = tensor_model_parallel_all_reduce(local_sums)

87 # Calculate the variance

88 count = self.tp_size * x.shape[-1]

77 - variance = (global_sums / count)

89 + variance = global_sums / count

91 else:

92 variance = x.pow(2).mean(-1, keepdim=True)

@@ -106,6 +118,9 @@ class Mixer2RMSNormGated(CustomOp):

118 gate: torch.Tensor,

119 ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:

120

121 + if not self.use_rms_norm:

122 + return x * nn.functional.silu(gate.to(torch.float32))

123 +

124 if self.tp_size > 1 or self.n_groups != 1:

125 return self.forward_native(x, gate)

126

@@ -124,7 +139,7 @@ class Mixer2RMSNormGated(CustomOp):

139

140

141 def extra_groups_for_head_shards(ngroups: int, tp_size: int):

127 - """Compute the increase in group numbers to account for

142 + """Compute the increase in group numbers to account for

143 replication in order to accompany the head shards."""

144

145 # in the case ngoups % tp_size == 0, this will be zero

@@ -182,13 +197,15 @@ def mamba_v2_sharded_weight_loader(

197 # seem to handle slices well.

198 # https://github.com/python/mypy/issues/2410

199 param.data[

185 - boundary:(boundary + take), # type: ignore[misc]

186 - ...] = loaded_weight[loaded_start_idx:( # type: ignore[misc]

187 - loaded_start_idx + take)] # type: ignore[misc]

200 + boundary:(boundary + take),

201 + ... # type: ignore[misc]

202 + ] = loaded_weight[loaded_start_idx:(loaded_start_idx +

203 + take) # type: ignore[misc]

204 + ] # type: ignore[misc]

205

206 # move indexing boundaries

207 boundary += shard_size

191 - loaded_boundary += (full_dim - extra)

208 + loaded_boundary += full_dim - extra

209

210 return loader

211

@@ -206,19 +223,22 @@ class MambaMixer2(CustomOp):

223 **selective** state spaces)

224 """

225

209 - def __init__(self,

210 - hidden_size: int,

211 - ssm_state_size: int,

212 - conv_kernel_size: int,

213 - intermediate_size: int,

214 - use_conv_bias: bool,

215 - use_bias: bool,

216 - n_groups: int = 1,

217 - num_heads: int = 128,

218 - head_dim: int = 64,

219 - rms_norm_eps: float = 1e-5,

220 - activation="silu",

221 - quant_config: Optional[QuantizationConfig] = None):

226 + def __init__(

227 + self,

228 + hidden_size: int,

229 + ssm_state_size: int,

230 + conv_kernel_size: int,

231 + intermediate_size: int,

232 + use_conv_bias: bool,

233 + use_bias: bool,

234 + n_groups: int = 1,

235 + num_heads: int = 128,

236 + head_dim: int = 64,

237 + rms_norm_eps: float = 1e-5,

238 + activation: str = "silu",

239 + use_rms_norm: bool = True,

240 + quant_config: Optional[QuantizationConfig] = None,

241 + ):

242 super().__init__()

243

244 # For TP, the sharding plan is as follows:

@@ -238,17 +258,16 @@ class MambaMixer2(CustomOp):

258 self.tp_size = get_tensor_model_parallel_world_size()

259 tp_rank = get_tensor_model_parallel_rank()

260

241 - assert num_heads % self.tp_size == 0, \

242 - "Tensor parallel world size must divide num heads."

261 + assert (num_heads % self.tp_size == 0

262 + ), "Tensor parallel world size must divide num heads."

263

244 - assert (n_groups % self.tp_size) == 0 or n_groups == 1, \

245 - (

246 - "If tensor parallel world size does not divide num_heads, "

247 - "then num_groups must equal 1."

248 - )

264 + assert (n_groups % self.tp_size) == 0 or n_groups == 1, (

265 + "If tensor parallel world size does not divide num_heads, "

266 + "then num_groups must equal 1.")

267

250 - assert self.tp_size == 1 or quant_config is None, \

251 - "Tensor parallel currently not supported for quantized models."

268 + assert (

269 + self.tp_size == 1 or quant_config is None

270 + ), "Tensor parallel currently not supported for quantized models."

271

272 self.ssm_state_size = ssm_state_size

273 self.activation = activation

@@ -265,8 +284,7 @@ class MambaMixer2(CustomOp):

284 self.n_groups = n_groups + extra_groups_for_head_shards(

285 n_groups, self.tp_size)

286

268 - self.conv_dim = (intermediate_size +

269 - 2 * self.n_groups * ssm_state_size)

287 + self.conv_dim = intermediate_size + 2 * self.n_groups * ssm_state_size

288 self.conv1d = ColumnParallelLinear(

289 input_size=conv_kernel_size,

290 output_size=self.conv_dim,

@@ -279,11 +297,12 @@ class MambaMixer2(CustomOp):

297 # doesn't allow to override it

298 self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)

299

282 - self.in_proj = ColumnParallelLinear(input_size=hidden_size,

283 - output_size=intermediate_size +

284 - self.conv_dim + self.num_heads,

285 - bias=use_bias,

286 - quant_config=quant_config)

300 + self.in_proj = ColumnParallelLinear(

301 + input_size=hidden_size,

302 + output_size=intermediate_size + self.conv_dim + self.num_heads,

303 + bias=use_bias,

304 + quant_config=quant_config,

305 + )

306

307 # - because in_proj is a concatenation of 3 weights, we

308 # need to interleave them before sharding

@@ -305,7 +324,8 @@ class MambaMixer2(CustomOp):

324 # - ditto for the otther two weights below

325 delattr(self.conv1d.bias, "weight_loader")

326 set_weight_attrs(

308 - self.conv1d.bias, {

327 + self.conv1d.bias,

328 + {

329 "weight_loader":

330 mamba_v2_sharded_weight_loader(

331 [

@@ -316,18 +336,25 @@ class MambaMixer2(CustomOp):

336 self.tp_size,

337 tp_rank,

338 )

319 - })

339 + },

340 + )

341

342 delattr(self.conv1d.weight, "weight_loader")

343 set_weight_attrs(

323 - self.conv1d.weight, {

344 + self.conv1d.weight,

345 + {

346 "weight_loader":

325 - mamba_v2_sharded_weight_loader([

326 - intermediate_settings,

327 - group_shard_settings,

328 - group_shard_settings,

329 - ], self.tp_size, tp_rank)

330 - })

347 + mamba_v2_sharded_weight_loader(

348 + [

349 + intermediate_settings,

350 + group_shard_settings,

351 + group_shard_settings,

352 + ],

353 + self.tp_size,

354 + tp_rank,

355 + )

356 + },

357 + )

358

359 if quant_config is None:

360 # - quant layers do not have a weight loader

@@ -345,8 +372,10 @@ class MambaMixer2(CustomOp):

372 head_setings, # for dt

373 ],

374 self.tp_size,

348 - tp_rank)

349 - })

375 + tp_rank,

376 + )

377 + },

378 + )

379

380 # - these are TPed by heads to reduce the size of the

381 # temporal shape

@@ -357,6 +386,7 @@ class MambaMixer2(CustomOp):

386 ))

387 self.D = nn.Parameter(torch.ones(num_heads // self.tp_size))

388 self.dt_bias = nn.Parameter(torch.ones(num_heads // self.tp_size))

389 + self.use_rms_norm = use_rms_norm

390

391 set_weight_attrs(self.D, {"weight_loader": sharded_weight_loader(0)})

392 a_weight_loader = composed_weight_loader(

@@ -365,18 +395,25 @@ class MambaMixer2(CustomOp):

395 set_weight_attrs(self.dt_bias,

396 {"weight_loader": sharded_weight_loader(0)})

397

368 - self.out_proj = RowParallelLinear(intermediate_size,

369 - hidden_size,

370 - bias=use_bias,

371 - input_is_parallel=True,

372 - quant_config=quant_config)

398 + self.out_proj = RowParallelLinear(

399 + intermediate_size,

400 + hidden_size,

401 + bias=use_bias,

402 + input_is_parallel=True,

403 + quant_config=quant_config,

404 + )

405

406 self.norm = Mixer2RMSNormGated(intermediate_size,

407 n_groups,

408 + self.use_rms_norm,

409 eps=rms_norm_eps)

410

378 - def forward_native(self, hidden_states: torch.Tensor,

379 - conv_state: torch.Tensor, ssm_state: torch.Tensor):

411 + def forward_native(

412 + self,

413 + hidden_states: torch.Tensor,

414 + conv_state: torch.Tensor,

415 + ssm_state: torch.Tensor,

416 + ):

417 pass

418

419 def forward_cuda(

@@ -384,6 +421,7 @@ class MambaMixer2(CustomOp):

421 hidden_states: torch.Tensor,

422 mamba_cache_params: MambaCacheParams,

423 mamba2_metadata: Mamba2Metadata,

424 + mup_vector: Optional[torch.Tensor] = None,

425 ):

426 # mamba2_metadata contains metadata necessary for the mamba2 triton

427 # kernels to operate in continuous batching and in chunked prefill

@@ -401,6 +439,10 @@ class MambaMixer2(CustomOp):

439

440 # 1. Gated MLP's linear projection

441 projected_states, _ = self.in_proj(hidden_states)

442 +

443 + if mup_vector is not None:

444 + projected_states = projected_states * mup_vector

445 +

446 gate, hidden_states_B_C, dt = torch.split(

447 projected_states,

448 [

@@ -561,6 +603,9 @@ class MambaMixer2(CustomOp):

603 hidden_states = torch.vstack(ssd_output_list)

604

605 # 4. gated MLP

606 + # GatedRMSNorm internally applying SiLU to the gate

607 + # SiLU is applied internally before normalization, unlike standard

608 + # norm usage

609 hidden_states = self.norm(hidden_states, gate)

610

611 # 5. Final linear projection

vllm/model_executor/models/falcon_h1.py new

+685

@@ -0,0 +1,685 @@

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

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

3 +from collections.abc import Iterable

4 +from typing import Optional

5 +

6 +import torch

7 +from torch import nn

8 +from transformers import FalconH1Config

9 +

10 +from vllm.attention.layer import Attention

11 +from vllm.config import CacheConfig, VllmConfig

12 +from vllm.distributed import divide, get_tensor_model_parallel_world_size

13 +from vllm.distributed.parallel_state import get_pp_group

14 +from vllm.forward_context import get_forward_context

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

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

17 +from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,

18 + QKVParallelLinear,

19 + RowParallelLinear)

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

21 +from vllm.model_executor.layers.mamba.mamba2_metadata import (

22 + Mamba2Metadata, prepare_mamba2_metadata)

23 +from vllm.model_executor.layers.mamba.mamba_mixer2 import (

24 + MambaMixer2, extra_groups_for_head_shards)

25 +from vllm.model_executor.layers.quantization import QuantizationConfig

26 +from vllm.model_executor.layers.rotary_embedding import get_rope

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

28 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

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

30 +from vllm.model_executor.models.mamba_cache import (MambaCacheManager,

31 + MambaCacheParams)

32 +from vllm.model_executor.sampling_metadata import SamplingMetadata

33 +from vllm.sequence import IntermediateTensors

34 +

35 +from .interfaces import (HasInnerState, IsHybrid, SupportsLoRA, SupportsPP,

36 + SupportsV0Only)

37 +from .utils import (PPMissingLayer, is_pp_missing_parameter,

38 + make_empty_intermediate_tensors_factory, make_layers,

39 + maybe_prefix)

40 +

41 +

42 +class FalconH1MLP(nn.Module):

43 +

44 + def __init__(

45 + self,

46 + config: FalconH1Config,

47 + quant_config: Optional[QuantizationConfig] = None,

48 + bias: bool = False,

49 + ) -> None:

50 + super().__init__()

51 + self.gate_up_proj = MergedColumnParallelLinear(

52 + input_size=config.hidden_size,

53 + output_sizes=[config.intermediate_size] * 2,

54 + bias=bias,

55 + quant_config=quant_config,

56 + )

57 + self.down_proj = RowParallelLinear(

58 + input_size=config.intermediate_size,

59 + output_size=config.hidden_size,

60 + bias=bias,

61 + quant_config=quant_config,

62 + )

63 + self.tp_size = get_tensor_model_parallel_world_size()

64 + self.intermediate_size = config.intermediate_size

65 + self.gate_multiplier, self.down_multiplier = config.mlp_multipliers

66 + if config.hidden_act != "silu":

67 + raise ValueError(f"Unsupported activation: {config.hidden_act}. "

68 + "Only silu is supported for now.")

69 + self.act_fn = SiluAndMul()

70 +

71 + def forward(self, x):

72 + x, _ = self.gate_up_proj(x)

73 + x[:, :self.intermediate_size // self.tp_size] *= self.gate_multiplier

74 + x = self.act_fn(x)

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

76 + x = x * self.down_multiplier

77 + return x

78 +

79 +

80 +class FalconH1SSMDecoderLayer(nn.Module):

81 +

82 + def __init__(

83 + self,

84 + config: FalconH1Config,

85 + cache_config: Optional[CacheConfig] = None,

86 + quant_config: Optional[QuantizationConfig] = None,

87 + ) -> None:

88 + super().__init__()

89 + self.config = config

90 + self.tp_size = get_tensor_model_parallel_world_size()

91 +

92 + self.d_ssm = (int(config.mamba_expand * config.hidden_size)

93 + if config.mamba_d_ssm is None else config.mamba_d_ssm)

94 +

95 + self.mamba = MambaMixer2(

96 + hidden_size=config.hidden_size,

97 + ssm_state_size=config.mamba_d_state,

98 + conv_kernel_size=config.mamba_d_conv,

99 + intermediate_size=self.d_ssm,

100 + use_conv_bias=config.mamba_conv_bias,

101 + use_bias=config.mamba_proj_bias,

102 + n_groups=config.mamba_n_groups,

103 + num_heads=config.mamba_n_heads,

104 + head_dim=config.mamba_d_head,

105 + rms_norm_eps=config.rms_norm_eps,

106 + activation=config.hidden_act,

107 + quant_config=quant_config,

108 + use_rms_norm=config.mamba_rms_norm,

109 + )

110 + # n_groups is overridden later by `MambaMixer2`

111 + self.groups_time_state_size = self.mamba.n_groups * config.mamba_d_state

112 + self.zxbcdt_multipliers = config.ssm_multipliers

113 + self._init_mup_vector()

114 +

115 + def _init_mup_vector(self):

116 + """

117 + Non learnable per-block scaling vector composed of element-wise

118 + multipliersapplied to each separate contiguous block of the output

119 + of the linear projection (in_proj) before further processing

120 + (gating, convolution, SSM):

121 +

122 + - Z block: [0 : d_ssm] → zxbcdt_multipliers[0]

123 + - X block: [d_ssm : 2 * d_ssm] → zxbcdt_multipliers[1]

124 + - B block: [2 * d_ssm : 2 * d_ssm + G * S] → zxbcdt_multipliers[2]

125 + - C block: [2 * d_ssm + G * S : 2 * d_ssm + 2 * G * S]

126 + → zxbcdt_multipliers[3]

127 + - dt block: [2 * d_ssm + 2 * G * S : end] → zxbcdt_multipliers[4]

128 +

129 + where:

130 + - d_ssm: Dimension of state-space model latent

131 + - G: Number of groups (n_groups)

132 + - S: SSM state size per group

133 + - All indices are divided by tp_size to support tensor parallelism

134 + """

135 + vector_shape = (2 * self.d_ssm + 2 * self.groups_time_state_size +

136 + self.config.mamba_n_heads) // self.tp_size

137 + mup_vector = torch.ones(1, vector_shape)

138 + # Z vector 0 -> d_ssm

139 + mup_vector[:, :self.d_ssm //

140 + self.tp_size] *= self.zxbcdt_multipliers[0]

141 + # X vector d_ssm -> 2 * d_ssm

142 + mup_vector[:,

143 + (self.d_ssm //

144 + self.tp_size):(2 * self.d_ssm //

145 + self.tp_size)] *= self.zxbcdt_multipliers[1]

146 + # B vector 2 * d_ssm -> 2 * d_ssm + (n_group * d_state)

147 + mup_vector[

148 + :,

149 + (2 * self.d_ssm) //

150 + self.tp_size:(2 * self.d_ssm + self.groups_time_state_size) //

151 + self.tp_size,

152 + ] *= self.zxbcdt_multipliers[2]

153 + # C vector 2 * d_ssm + (n_group * d_state)

154 + # -> 2 * d_ssm + 2 * (n_group * d_state)

155 + mup_vector[

156 + :,

157 + (2 * self.d_ssm + self.groups_time_state_size) //

158 + self.tp_size:(2 * self.d_ssm + 2 * self.groups_time_state_size) //

159 + self.tp_size,

160 + ] *= self.zxbcdt_multipliers[3]

161 + # dt vector 2 * d_ssm + 2 * (n_group * d_state)

162 + # -> 2 * d_ssm + 2 * (n_group * d_state) + n_heads

163 + mup_vector[

164 + :,

165 + (2 * self.d_ssm + 2 * self.groups_time_state_size) //

166 + self.tp_size:,

167 + ] *= self.zxbcdt_multipliers[4]

168 +

169 + self.register_buffer("mup_vector", mup_vector, persistent=False)

170 +

171 + def forward(

172 + self,

173 + hidden_states: torch.Tensor,

174 + residual: Optional[torch.Tensor],

175 + mamba_cache_params: MambaCacheParams,

176 + mamba2_metadata: Mamba2Metadata,

177 + **kwargs,

178 + ):

179 + hidden_states = self.mamba(

180 + hidden_states,

181 + mamba_cache_params,

182 + mamba2_metadata=mamba2_metadata,

183 + mup_vector=self.mup_vector,

184 + )

185 + return hidden_states, residual

186 +

187 +

188 +class FalconH1AttentionDecoderLayer(nn.Module):

189 +

190 + def __init__(

191 + self,

192 + config: FalconH1Config,

193 + cache_config: Optional[CacheConfig] = None,

194 + quant_config: Optional[QuantizationConfig] = None,

195 + prefix: str = "",

196 + ) -> None:

197 + super().__init__()

198 + rope_theta = getattr(config, "rope_theta", 1e11)

199 + rope_scaling = getattr(config, "rope_scaling", None)

200 + max_position_embeddings = getattr(config, "max_position_embeddings",

201 + 8192)

202 + self.hidden_size = config.hidden_size

203 + tp_size = get_tensor_model_parallel_world_size()

204 + self.total_num_heads = config.num_attention_heads

205 + assert self.total_num_heads % tp_size == 0

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

207 + self.total_num_kv_heads = config.num_key_value_heads

208 + if self.total_num_kv_heads >= tp_size:

209 + # Number of KV heads is greater than TP size, so we partition

210 + # the KV heads across multiple tensor parallel GPUs.

211 + assert self.total_num_kv_heads % tp_size == 0

212 + else:

213 + # Number of KV heads is less than TP size, so we replicate

214 + # the KV heads across multiple tensor parallel GPUs.

215 + assert tp_size % self.total_num_kv_heads == 0

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

217 + self.head_dim = (config.hidden_size // self.total_num_heads if getattr(

218 + config, "head_dim", None) is None else config.head_dim)

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

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

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

222 + self.rope_theta = rope_theta

223 + self.max_position_embeddings = max_position_embeddings

224 +

225 + if hasattr(config, "partial_rotary_factor"):

226 + rotary_dim = self.head_dim * config.partial_rotary_factor

227 + elif hasattr(config, "attn_rotary_emb"):

228 + rotary_dim = config.attn_rotary_emb # for backward compatibility

229 + else:

230 + rotary_dim = self.head_dim # default

231 +

232 + self.rotary_emb = get_rope(

233 + head_size=self.head_dim,

234 + rotary_dim=rotary_dim,

235 + max_position=max_position_embeddings,

236 + rope_scaling=rope_scaling,

237 + base=rope_theta,

238 + is_neox_style=True,

239 + dtype=None, # see impl of get_rope

240 + )

241 +

242 + self.qkv_proj = QKVParallelLinear(

243 + config.hidden_size,

244 + self.head_dim,

245 + self.total_num_heads,

246 + self.total_num_kv_heads,

247 + bias=False,

248 + quant_config=quant_config,

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

250 + )

251 + self.o_proj = RowParallelLinear(

252 + self.total_num_heads * self.head_dim,

253 + config.hidden_size,

254 + bias=False,

255 + quant_config=quant_config,

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

257 + )

258 +

259 + self.attn = Attention(

260 + self.num_heads,

261 + self.head_dim,

262 + self.scaling,

263 + num_kv_heads=self.num_kv_heads,

264 + cache_config=cache_config,

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

266 + )

267 + self.key_multiplier = config.key_multiplier

268 +

269 + def self_attention(

270 + self,

271 + positions: torch.Tensor,

272 + hidden_states: torch.Tensor,

273 + **kwargs,

274 + ) -> torch.Tensor:

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

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

277 + k = k * self.key_multiplier

278 +

279 + q, k = self.rotary_emb(positions, q, k)

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

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

282 + return output

283 +

284 + def forward(

285 + self,

286 + positions: torch.Tensor,

287 + hidden_states: torch.Tensor,

288 + residual: Optional[torch.Tensor],

289 + **kwargs,

290 + ):

291 + hidden_states = self.self_attention(

292 + positions=positions,

293 + hidden_states=hidden_states,

294 + )

295 + return hidden_states, residual

296 +

297 +

298 +class FalconH1ParallelHybrid(nn.Module):

299 + """

300 + A hybrid decoder layer for FalconH1 where the input is processed

301 + in parallel through both the self-attention branch and the SSM (Mamba)

302 + branch. Their outputs are then summed to produce the final hidden state.

303 +

304 + This layer uses:

305 + - FalconH1AttentionDecoderLayer for the multi-head self-attention branch.

306 + - FalconH1SSMDecoderLayer for the state-space (Mamba) branch.

307 + """

308 +

309 + def __init__(

310 + self,

311 + config: FalconH1Config,

312 + layer_idx: int,

313 + cache_config: Optional[CacheConfig] = None,

314 + quant_config: Optional[QuantizationConfig] = None,

315 + prefix: str = "",

316 + ) -> None:

317 + super().__init__()

318 + # Instantiate the attention branch

319 + self.self_attn = FalconH1AttentionDecoderLayer(

320 + config=config,

321 + cache_config=cache_config,

322 + quant_config=quant_config,

323 + prefix=prefix,

324 + )

325 + # Instantiate the SSM branch

326 + self.mamba = FalconH1SSMDecoderLayer(

327 + config=config,

328 + cache_config=cache_config,

329 + quant_config=quant_config,

330 + )

331 + self.ssm_out_multiplier = config.ssm_out_multiplier

332 + self.ssm_in_multiplier = config.ssm_in_multiplier

333 +

334 + self.attention_in_multiplier = config.attention_in_multiplier

335 + self.attn_out_multiplier = config.attention_out_multiplier

336 +

337 + self.feed_forward = FalconH1MLP(config)

338 +

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

340 + eps=config.rms_norm_eps)

341 + self.pre_ff_layernorm = RMSNorm(config.hidden_size,

342 + eps=config.rms_norm_eps)

343 +

344 + def forward(

345 + self,

346 + positions: torch.Tensor,

347 + hidden_states: torch.Tensor,

348 + mamba_cache_params: MambaCacheParams,

349 + mamba2_metadata: Mamba2Metadata,

350 + **kwargs,

351 + ):

352 + residual = hidden_states

353 + hidden_states = self.input_layernorm(hidden_states)

354 + # Process input through the attention branch.

355 + # FalconH1AttentionDecoderLayer expects positions, hidden_states,

356 + # kv_cache, attn_metadata, and residual.

357 + attn_hidden, _ = self.self_attn(

358 + positions=positions,

359 + hidden_states=hidden_states * self.attention_in_multiplier,

360 + residual=residual,

361 + **kwargs,

362 + )

363 +

364 + # Process input through the SSM branch.

365 + # FalconH1SSMDecoderLayer expects hidden_states, attn_metadata,

366 + # residual, mamba_cache_params, and sequence_idx.

367 + ssm_hidden, _ = self.mamba(

368 + hidden_states=hidden_states * self.ssm_in_multiplier,

369 + residual=residual,

370 + mamba_cache_params=mamba_cache_params,

371 + mamba2_metadata=mamba2_metadata,

372 + **kwargs,

373 + )

374 + # Sum the outputs from both branches.

375 + # We assume both branches produce outputs of the same

376 + # dimensionality (config.hidden_size).

377 + hidden_states = (attn_hidden * self.attn_out_multiplier) + (

378 + ssm_hidden * self.ssm_out_multiplier)

379 + hidden_states = hidden_states + residual

380 +

381 + # feed-forward

382 + residual = hidden_states

383 + hidden_states = self.pre_ff_layernorm(hidden_states)

384 + hidden_states = self.feed_forward(hidden_states)

385 + hidden_states = residual + hidden_states

386 +

387 + return hidden_states

388 +

389 +

390 +class FalconH1Model(nn.Module):

391 +

392 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

393 + super().__init__()

394 + config: FalconH1Config = vllm_config.model_config.hf_config

395 + cache_config = vllm_config.cache_config

396 + quant_config = vllm_config.quant_config

397 + lora_config = vllm_config.lora_config

398 +

399 + self.config = config

400 + lora_vocab = ((lora_config.lora_extra_vocab_size *

401 + (lora_config.max_loras or 1)) if lora_config else 0)

402 + self.vocab_size = config.vocab_size + lora_vocab

403 + self.org_vocab_size = config.vocab_size

404 + if get_pp_group().is_first_rank:

405 +

406 + self.embed_tokens = VocabParallelEmbedding(

407 + self.vocab_size,

408 + config.hidden_size,

409 + org_num_embeddings=config.vocab_size,

410 + )

411 + self.embedding_multiplier = config.embedding_multiplier

412 + else:

413 + self.embed_tokens = PPMissingLayer()

414 + self.embedding_multiplier = 1.0

415 +

416 + def get_layer(prefix: str):

417 + layer_idx = int(prefix.rsplit(".", 1)[1])

418 + layer_class = FalconH1ParallelHybrid

419 + return layer_class(

420 + config,

421 + layer_idx,

422 + cache_config,

423 + quant_config=quant_config,

424 + prefix=prefix,

425 + )

426 +

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

428 + config.num_hidden_layers, get_layer, prefix=f"{prefix}.layers")

429 + self.make_empty_intermediate_tensors = (

430 + make_empty_intermediate_tensors_factory(

431 + ["hidden_states", "residual"], config.hidden_size))

432 + if get_pp_group().is_last_rank:

433 + self.final_layernorm = RMSNorm(config.hidden_size,

434 + eps=config.rms_norm_eps)

435 + else:

436 + self.final_layernorm = PPMissingLayer()

437 +

438 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

439 + return self.embed_tokens(input_ids)

440 +

441 + def forward(

442 + self,

443 + input_ids: torch.Tensor,

444 + positions: torch.Tensor,

445 + mamba_cache_params: MambaCacheParams,

446 + intermediate_tensors: Optional[IntermediateTensors] = None,

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

448 + ) -> torch.Tensor:

449 +

450 + # pass a sequence index tensor, that is required for

451 + # proper continuous batching computation including

452 + # chunked prefill

453 + attn_metadata = get_forward_context().attn_metadata

454 + mamba2_metadata = prepare_mamba2_metadata(

455 + chunk_size=self.config.mamba_chunk_size,

456 + input_ids=input_ids,

457 + attn_metadata=attn_metadata,

458 + )

459 + if get_pp_group().is_first_rank:

460 + if inputs_embeds is not None:

461 + hidden_states = inputs_embeds * self.embedding_multiplier

462 + else:

463 + hidden_states = (self.get_input_embeddings(input_ids) *

464 + self.embedding_multiplier)

465 + else:

466 + assert intermediate_tensors is not None

467 + hidden_states = intermediate_tensors["hidden_states"]

468 +

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

470 + layer = self.layers[i]

471 + layer_mamba_cache_params = mamba_cache_params.at_layer_idx(i)

472 + hidden_states = layer(

473 + positions=positions,

474 + hidden_states=hidden_states,

475 + mamba_cache_params=layer_mamba_cache_params,

476 + mamba2_metadata=mamba2_metadata,

477 + )

478 + if not get_pp_group().is_last_rank:

479 + return IntermediateTensors({

480 + "hidden_states": hidden_states,

481 + })

482 + hidden_states = self.final_layernorm(hidden_states)

483 + return hidden_states

484 +

485 +

486 +class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,

487 + IsHybrid, SupportsV0Only):

488 + packed_modules_mapping = {

489 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

490 + "gate_up_proj": ["gate_proj", "up_proj"],

491 + }

492 +

493 + embedding_modules = {

494 + "embed_tokens": "input_embeddings",

495 + "lm_head": "output_embeddings",

496 + }

497 + embedding_padding_modules = ["lm_head"]

498 +

499 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

500 + config = vllm_config.model_config.hf_config

501 + self.vllm_config = vllm_config

502 + self.model_config = vllm_config.model_config

503 + cache_config = vllm_config.cache_config

504 + lora_config = vllm_config.lora_config

505 + scheduler_config = vllm_config.scheduler_config

506 + assert (not cache_config.enable_prefix_caching

507 + ), "FalconH1 currently does not support prefix caching"

508 +

509 + self.quant_config = vllm_config.quant_config

510 +

511 + super().__init__()

512 + self.config = config

513 + self.scheduler_config = scheduler_config

514 + self.model = FalconH1Model(vllm_config=vllm_config,

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

516 + self.tie_word_embeddings = config.tie_word_embeddings

517 + self.unpadded_vocab_size = config.vocab_size

518 + self.mamba_cache: Optional[MambaCacheManager] = None

519 + if lora_config:

520 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size

521 + if get_pp_group().is_last_rank:

522 + self.lm_head = ParallelLMHead(

523 + self.unpadded_vocab_size,

524 + config.hidden_size,

525 + org_num_embeddings=config.vocab_size,

526 + padding_size=(

527 + DEFAULT_VOCAB_PADDING_SIZE

528 + # We need bigger padding if using lora for kernel

529 + # compatibility

530 + if not lora_config else

531 + lora_config.lora_vocab_padding_size),

532 + )

533 + self.lm_head_multiplier = config.lm_head_multiplier

534 + if self.tie_word_embeddings:

535 + self.lm_head = self.lm_head.tie_weights(

536 + self.model.embed_tokens)

537 + # Used to track and store by the Mamba cache between steps.

538 +

539 + self.logits_processor = LogitsProcessor(

540 + self.unpadded_vocab_size,

541 + config.vocab_size,

542 + scale=config.lm_head_multiplier,

543 + )

544 + else:

545 + self.lm_head = PPMissingLayer()

546 +

547 + self.make_empty_intermediate_tensors = (

548 + self.model.make_empty_intermediate_tensors)

549 +

550 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

551 + return self.model.get_input_embeddings(input_ids)

552 +

553 + def forward(

554 + self,

555 + input_ids: torch.Tensor,

556 + positions: torch.Tensor,

557 + intermediate_tensors: Optional[IntermediateTensors] = None,

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

559 + **kwargs,

560 + ):

561 + if self.mamba_cache is None:

562 + self.mamba_cache = MambaCacheManager(

563 + self.vllm_config,

564 + self.lm_head.weight.dtype

565 + if hasattr(self.lm_head, 'weight') else torch.bfloat16,

566 + self.config.num_hidden_layers,

567 + *self._get_mamba_cache_shape(),

568 + )

569 + mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)

570 + hidden_states = self.model(

571 + input_ids,

572 + positions,

573 + mamba_cache_params,

574 + intermediate_tensors,

575 + inputs_embeds,

576 + )

577 +

578 + return hidden_states

579 +

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

581 + return self.mamba_cache.copy_inputs_before_cuda_graphs(

582 + input_buffers, **kwargs)

583 +

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

585 + return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)

586 +

587 + def _get_mamba_cache_shape(

588 + self) -> tuple[tuple[int, int], tuple[int, int]]:

589 + world_size = get_tensor_model_parallel_world_size()

590 + hidden_size = self.config.hidden_size

591 +

592 + conv_state_shape, temporal_state_shape = None, None

593 +

594 + intermediate_size = (int(self.config.mamba_expand *

595 + hidden_size) if self.config.mamba_d_ssm

596 + is None else self.config.mamba_d_ssm)

597 +

598 + # if n_groups is not divisible by world_size, need to extend the shards

599 + # to ensure all groups needed by a head is sharded along with it

600 + n_groups = self.config.mamba_n_groups + extra_groups_for_head_shards(

601 + self.config.mamba_n_groups, world_size)

602 +

603 + # - heads and n_groups are TP-ed

604 + conv_dim = intermediate_size + 2 * n_groups * self.config.mamba_d_state

605 + conv_state_shape = (

606 + divide(conv_dim, world_size),

607 + self.config.mamba_d_conv - 1,

608 + )

609 +

610 + # These are not TP-ed as they depend on A, dt_bias, D

611 + # - they are typically small

612 + # e.g., (h_heads, d_head, d_state) = (128, 64, 128)

613 + temporal_state_shape = (

614 + divide(self.config.mamba_n_heads, world_size),

615 + self.config.mamba_d_head,

616 + self.config.mamba_d_state,

617 + )

618 + return conv_state_shape, temporal_state_shape

619 +

620 + def compute_logits(

621 + self,

622 + hidden_states: torch.Tensor,

623 + sampling_metadata: SamplingMetadata,

624 + ) -> Optional[torch.Tensor]:

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

626 + sampling_metadata)

627 +

628 + return logits

629 +

630 + def load_weights(self, weights: Iterable[tuple[str,

631 + torch.Tensor]]) -> set[str]:

632 + stacked_params_mapping = [

633 + # (param_name, shard_name, shard_id)

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

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

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

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

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

639 + ]

640 +

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

642 + loaded_params: set[str] = set()

643 + for name, loaded_weight in weights:

644 + if "rotary_emb.inv_freq" in name:

645 + continue

646 +

647 + if "A_log" in name:

648 + name = name.replace("A_log", "A")

649 +

650 + if "mamba" in name:

651 + name = name.replace("mamba", "mamba.mamba")

652 +

653 + for param_name, weight_name, shard_id in stacked_params_mapping:

654 + if weight_name not in name:

655 + continue

656 +

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

658 + # Skip loading extra bias for GPTQ models.

659 + if name.endswith(".bias") and name not in params_dict:

660 + continue

661 + # Skip layers on other devices.

662 + if is_pp_missing_parameter(name, self):

663 + continue

664 + param = params_dict[name]

665 + weight_loader = param.weight_loader

666 + weight_loader(param, loaded_weight, shard_id)

667 + break

668 + else:

669 + # Skip loading extra bias for GPTQ models.

670 + if name.endswith(".bias") and name not in params_dict:

671 + continue

672 + if is_pp_missing_parameter(name, self):

673 + continue

674 + if self.tie_word_embeddings and "lm_head" in name:

675 + continue

676 +

677 + param = params_dict[name]

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

679 + default_weight_loader)

680 + weight_loader(param, loaded_weight)

681 + loaded_params.add(name)

682 +

683 + if self.tie_word_embeddings:

684 + loaded_params.add("lm_head.weight")

685 + return loaded_params

vllm/model_executor/models/registry.py

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

79 "LLaMAForCausalLM": ("llama", "LlamaForCausalLM"),

80 "MambaForCausalLM": ("mamba", "MambaForCausalLM"),

81 "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"),

82 + "FalconH1ForCausalLM":("falcon_h1", "FalconH1ForCausalLM"),

83 "Mamba2ForCausalLM": ("mamba2", "Mamba2ForCausalLM"),

84 "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"),

85 "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),