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
[FalconH1] Fix output dtype in RMSNorm fallback path for Falcon-H1 (e.g. 0.5B) (#18500)
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 22, 2025 at 06:23 UTC
20bd6f4d2ee10d6187d93e4bc6a64f78195243a2
2 files changed
+5
-4
vllm/model_executor/layers/mamba/mamba_mixer2.py
+5
-3
@@ -77,7 +77,7 @@ class Mixer2RMSNormGated(CustomOp):
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
80
+ return x.to(input_dtype)
81
82
if self.n_groups == 1:
83
if self.tp_size > 1:
@@ -117,9 +117,11 @@ class Mixer2RMSNormGated(CustomOp):
117
x: torch.Tensor,
118
gate: torch.Tensor,
119
) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
120
-
120
+ input_dtype = x.dtype
121
if not self.use_rms_norm:
122
- return x * nn.functional.silu(gate.to(torch.float32))
122
+ # Keep gate in float32 for numerical stability during silu
123
+ return x * nn.functional.silu(gate.to(
124
+ torch.float32)).to(input_dtype)
125
126
if self.tp_size > 1 or self.n_groups != 1:
127
return self.forward_native(x, gate)
vllm/model_executor/models/falcon_h1.py
-1
@@ -453,7 +453,6 @@ class FalconH1Model(nn.Module):
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,
456
attn_metadata=attn_metadata,
457
)
458
if get_pp_group().is_first_rank: