[XPU] Dispatch nn.LayerNorm to fused SYCL kernel via CustomOp (#57172)
Signed-off-by: Michal Ganczarenko <michal.ganczarenko@intel.com> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Michał Ganczarenko committed
Sep 30, 2026 at 14:19 UTC
df5668e5c935990cd9595bad2c3559b821092a5b
13 files changed
+237
-46
vllm/_xpu_ops.py
+66
@@ -135,6 +135,46 @@ def _fused_add_gemma_rms_norm_impl(
135
torch.ops._C.fused_add_gemma_rms_norm(input, residual, weight, epsilon)
136
137
138
+def _layer_norm_impl(
139
+ out: torch.Tensor,
140
+ input: torch.Tensor,
141
+ weight: torch.Tensor,
142
+ bias: torch.Tensor,
143
+ epsilon: float,
144
+) -> None:
145
+ torch.ops._C.layer_norm(out, input, weight, bias, epsilon)
146
+
147
+
148
+def _fused_add_layer_norm_impl(
149
+ input: torch.Tensor,
150
+ residual: torch.Tensor,
151
+ weight: torch.Tensor,
152
+ bias: torch.Tensor,
153
+ epsilon: float,
154
+) -> None:
155
+ torch.ops._C.fused_add_layer_norm(input, residual, weight, bias, epsilon)
156
+
157
+
158
+def _nemotron_layer_norm_impl(
159
+ out: torch.Tensor,
160
+ input: torch.Tensor,
161
+ weight: torch.Tensor,
162
+ bias: torch.Tensor,
163
+ epsilon: float,
164
+) -> None:
165
+ torch.ops._C.nemotron_layer_norm(out, input, weight, bias, epsilon)
166
+
167
+
168
+def _fused_add_nemotron_layer_norm_impl(
169
+ input: torch.Tensor,
170
+ residual: torch.Tensor,
171
+ weight: torch.Tensor,
172
+ bias: torch.Tensor,
173
+ epsilon: float,
174
+) -> None:
175
+ torch.ops._C.fused_add_nemotron_layer_norm(input, residual, weight, bias, epsilon)
176
+
177
+
178
def _gdn_attention_core_xpu_impl(
179
core_attn_out: torch.Tensor,
180
z: torch.Tensor,
@@ -1268,6 +1308,32 @@ class xpu_ops:
1308
mutates_args=["input", "residual"],
1309
)
1310
1311
+ if hasattr(torch.ops._C, "layer_norm"):
1312
+ direct_register_custom_op(
1313
+ op_name="xpu_layer_norm",
1314
+ op_func=_layer_norm_impl,
1315
+ mutates_args=["out"],
1316
+ )
1317
+
1318
+ direct_register_custom_op(
1319
+ op_name="xpu_fused_add_layer_norm",
1320
+ op_func=_fused_add_layer_norm_impl,
1321
+ mutates_args=["input", "residual"],
1322
+ )
1323
+
1324
+ if hasattr(torch.ops._C, "nemotron_layer_norm"):
1325
+ direct_register_custom_op(
1326
+ op_name="xpu_nemotron_layer_norm",
1327
+ op_func=_nemotron_layer_norm_impl,
1328
+ mutates_args=["out"],
1329
+ )
1330
+
1331
+ direct_register_custom_op(
1332
+ op_name="xpu_fused_add_nemotron_layer_norm",
1333
+ op_func=_fused_add_nemotron_layer_norm_impl,
1334
+ mutates_args=["input", "residual"],
1335
+ )
1336
+
1337
direct_register_custom_op(
1338
op_name="xpu_ops_deepseek_scaling_rope",
1339
op_func=_xpu_ops_deepseek_scaling_rope_impl,
vllm/model_executor/layers/layernorm.py
+73
-16
@@ -211,6 +211,79 @@ class GemmaRMSNorm(CustomOp):
211
return out
212
213
214
+@CustomOp.register("layer_norm")
215
+class LayerNorm(CustomOp):
216
+ """Standard (mean-centered) LayerNorm.
217
+
218
+ Drop-in for a bare `nn.LayerNorm` that dispatches to a fused XPU kernel
219
+ when one is available, and to the native implementation otherwise.
220
+
221
+ Normalization runs in the wider of the input and parameter dtypes and the
222
+ result is cast back to the input dtype, so `dtype=torch.float32` gives an
223
+ fp32 reduction for a lower-precision input.
224
+ """
225
+
226
+ def __init__(
227
+ self,
228
+ hidden_size: int,
229
+ eps: float = 1e-5,
230
+ elementwise_affine: bool = True,
231
+ bias: bool = True,
232
+ dtype: torch.dtype | None = None,
233
+ ) -> None:
234
+ super().__init__()
235
+ self.normalized_shape = (hidden_size,)
236
+ self.eps = eps
237
+ weight_dtype = dtype or torch.get_default_dtype()
238
+ self.weight: nn.Parameter | None = None
239
+ self.bias: nn.Parameter | None = None
240
+ if elementwise_affine:
241
+ self.weight = nn.Parameter(torch.ones(hidden_size, dtype=weight_dtype))
242
+ if bias:
243
+ self.bias = nn.Parameter(torch.zeros(hidden_size, dtype=weight_dtype))
244
+
245
+ def _compute_dtype(self, x: torch.Tensor) -> torch.dtype:
246
+ if self.weight is None:
247
+ return x.dtype
248
+ return torch.promote_types(x.dtype, self.weight.dtype)
249
+
250
+ def forward_native(self, x: torch.Tensor) -> torch.Tensor:
251
+ compute_dtype = self._compute_dtype(x)
252
+ if compute_dtype != x.dtype:
253
+ return F.layer_norm(
254
+ x.to(compute_dtype),
255
+ self.normalized_shape,
256
+ self.weight,
257
+ self.bias,
258
+ self.eps,
259
+ ).type_as(x)
260
+ return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
261
+
262
+ def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
263
+ return self.forward_native(x)
264
+
265
+ def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
266
+ import vllm._xpu_ops # noqa: F401 registers torch.ops.vllm.xpu_layer_norm
267
+
268
+ if (
269
+ self.weight is None
270
+ or self.bias is None
271
+ or not hasattr(torch.ops._C, "layer_norm")
272
+ # The kernel reduces in the input dtype, so a widened reduction
273
+ # takes the native path.
274
+ or self._compute_dtype(x) != x.dtype
275
+ ):
276
+ return self.forward_native(x)
277
+ # empty_like preserves x's strides, but the kernel requires a
278
+ # contiguous out (unlike x, which it can handle non-contiguous).
279
+ out = torch.empty(x.shape, device=x.device, dtype=x.dtype)
280
+ torch.ops.vllm.xpu_layer_norm(out, x, self.weight, self.bias, self.eps)
281
+ return out
282
+
283
+ def extra_repr(self) -> str:
284
+ return f"hidden_size={self.normalized_shape[0]}, eps={self.eps}"
285
+
286
+
287
# --8<-- [start:rms_norm_gated]
288
@CustomOp.register("rms_norm_gated")
289
class RMSNormGated(CustomOp):
@@ -350,19 +423,3 @@ class RMSNormGated(CustomOp):
423
self, x: torch.Tensor, z: torch.Tensor | None = None
424
) -> torch.Tensor:
425
return self.forward_cuda(x, z)
353
-
354
-
355
-class LayerNorm(nn.Module):
356
- """Layer Normalization."""
357
-
358
- def __init__(self, dim: int, eps: float = 1e-6):
359
- super().__init__()
360
- self.dim = dim
361
- self.eps = eps
362
- self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
363
- self.bias = nn.Parameter(torch.zeros(dim, dtype=torch.float32))
364
-
365
- def forward(self, x: torch.Tensor):
366
- return F.layer_norm(
367
- x.float(), (self.dim,), self.weight, self.bias, self.eps
368
- ).type_as(x)
vllm/model_executor/models/deepseek_v2.py
+1
-1
@@ -715,7 +715,7 @@ class Indexer(nn.Module):
715
disable_tp=True,
716
prefix=f"{prefix}.wk_weights_proj",
717
)
718
- self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
718
+ self.k_norm = LayerNorm(self.head_dim, eps=1e-6, dtype=torch.float32)
719
self.softmax_scale = self.head_dim**-0.5
720
721
self.scale_fmt = "ue8m0"
vllm/model_executor/models/gpt2.py
+4
-3
@@ -35,6 +35,7 @@ from vllm.distributed.parallel_state import (
35
)
36
from vllm.model_executor.layers.activation import get_act_fn
37
from vllm.model_executor.layers.attention import Attention
38
+from vllm.model_executor.layers.layernorm import LayerNorm
39
from vllm.model_executor.layers.linear import (
40
ColumnParallelLinear,
41
QKVParallelLinear,
@@ -156,11 +157,11 @@ class GPT2Block(nn.Module):
157
hidden_size = config.hidden_size
158
inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size
159
159
- self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
160
+ self.ln_1 = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
161
self.attn = GPT2Attention(
162
config, cache_config, quant_config, prefix=f"{prefix}.attn"
163
)
163
- self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
164
+ self.ln_2 = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
165
self.mlp = GPT2MLP(inner_dim, config, quant_config, prefix=f"{prefix}.mlp")
166
167
def forward(
@@ -212,7 +213,7 @@ class GPT2Model(nn.Module):
213
lambda prefix: GPT2Block(config, cache_config, quant_config, prefix=prefix),
214
prefix=f"{prefix}.h",
215
)
215
- self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
216
+ self.ln_f = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
217
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
218
["hidden_states"], config.n_embd
219
)
vllm/model_executor/models/gpt_j.py
+3
-2
@@ -31,6 +31,7 @@ from vllm.config import CacheConfig, VllmConfig
31
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
32
from vllm.model_executor.layers.activation import get_act_fn
33
from vllm.model_executor.layers.attention import Attention
34
+from vllm.model_executor.layers.layernorm import LayerNorm
35
from vllm.model_executor.layers.linear import (
36
ColumnParallelLinear,
37
QKVParallelLinear,
@@ -163,7 +164,7 @@ class GPTJBlock(nn.Module):
164
):
165
super().__init__()
166
inner_dim = 4 * config.n_embd if config.n_inner is None else config.n_inner
166
- self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
167
+ self.ln_1 = LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
168
self.attn = GPTJAttention(
169
config, cache_config, quant_config, prefix=f"{prefix}.attn"
170
)
@@ -206,7 +207,7 @@ class GPTJModel(nn.Module):
207
lambda prefix: GPTJBlock(config, cache_config, quant_config, prefix=prefix),
208
prefix=f"{prefix}.h",
209
)
209
- self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
210
+ self.ln_f = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
211
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
212
["hidden_states"], config.n_embd
213
)
vllm/model_executor/models/nemotron.py
+42
-13
@@ -33,6 +33,7 @@ from torch import nn
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36
+from vllm.model_executor.custom_op import CustomOp
37
from vllm.model_executor.layers.activation import get_act_fn
38
from vllm.model_executor.layers.attention import Attention
39
from vllm.model_executor.layers.linear import (
@@ -77,23 +78,22 @@ def _cast_if_autocast_enabled(device_type: str, *args):
78
)
79
80
80
-class NemotronLayerNorm1P(nn.LayerNorm):
81
- def __init__(
82
- self,
83
- normalized_shape: int | list[int] | torch.Size,
84
- eps: float = 1e-5,
85
- elementwise_affine: bool = True,
86
- bias: bool = True,
87
- device=None,
88
- dtype=None,
89
- ):
90
- super().__init__(normalized_shape, eps, elementwise_affine, bias, device, dtype)
81
+@CustomOp.register("nemotron_layer_norm")
82
+class NemotronLayerNorm1P(CustomOp):
83
+ """LayerNorm variant used by Nemotron: `x * (1 + w)` instead of `x * w`."""
84
92
- def forward(
85
+ def __init__(self, hidden_size: int, eps: float = 1e-5) -> None:
86
+ super().__init__()
87
+ self.normalized_shape = (hidden_size,)
88
+ self.eps = eps
89
+ self.weight = nn.Parameter(torch.zeros(hidden_size))
90
+ self.bias = nn.Parameter(torch.zeros(hidden_size))
91
+
92
+ def forward_native(
93
self,
94
x: torch.Tensor,
95
residual: torch.Tensor | None = None,
96
- ) -> torch.Tensor:
96
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
97
if residual is not None:
98
x = x + residual
99
residual = x
@@ -105,6 +105,35 @@ class NemotronLayerNorm1P(nn.LayerNorm):
105
x = torch.nn.functional.layer_norm(*args)
106
return x if residual is None else (x, residual)
107
108
+ def forward_cuda(
109
+ self,
110
+ x: torch.Tensor,
111
+ residual: torch.Tensor | None = None,
112
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
113
+ return self.forward_native(x, residual)
114
+
115
+ def forward_xpu(
116
+ self,
117
+ x: torch.Tensor,
118
+ residual: torch.Tensor | None = None,
119
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
120
+ import vllm._xpu_ops # noqa: F401 registers torch.ops.vllm.xpu_nemotron_layer_norm
121
+
122
+ if not hasattr(torch.ops._C, "nemotron_layer_norm"):
123
+ return self.forward_native(x, residual)
124
+ if residual is not None:
125
+ if not hasattr(torch.ops._C, "fused_add_nemotron_layer_norm"):
126
+ return self.forward_native(x, residual)
127
+ torch.ops.vllm.xpu_fused_add_nemotron_layer_norm(
128
+ x, residual, self.weight, self.bias, self.eps
129
+ )
130
+ return x, residual
131
+ # empty_like preserves x's strides, but the kernel requires a
132
+ # contiguous out (unlike x, which it can handle non-contiguous).
133
+ out = torch.empty(x.shape, device=x.device, dtype=x.dtype)
134
+ torch.ops.vllm.xpu_nemotron_layer_norm(out, x, self.weight, self.bias, self.eps)
135
+ return out
136
+
137
138
class NemotronMLP(nn.Module):
139
def __init__(
vllm/model_executor/models/opt.py
+5
-3
@@ -32,6 +32,7 @@ from vllm.config import CacheConfig, VllmConfig
32
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
33
from vllm.model_executor.layers.activation import get_act_fn
34
from vllm.model_executor.layers.attention import Attention
35
+from vllm.model_executor.layers.layernorm import LayerNorm
36
from vllm.model_executor.layers.linear import (
37
ColumnParallelLinear,
38
QKVParallelLinear,
@@ -143,7 +144,7 @@ class OPTDecoderLayer(nn.Module):
144
)
145
self.do_layer_norm_before = config.do_layer_norm_before
146
146
- self.self_attn_layer_norm = nn.LayerNorm(
147
+ self.self_attn_layer_norm = LayerNorm(
148
self.embed_dim, elementwise_affine=config.layer_norm_elementwise_affine
149
)
150
self.fc1 = ColumnParallelLinear(
@@ -161,7 +162,7 @@ class OPTDecoderLayer(nn.Module):
162
quant_config=quant_config,
163
prefix=f"{prefix}.fc2",
164
)
164
- self.final_layer_norm = nn.LayerNorm(
165
+ self.final_layer_norm = LayerNorm(
166
self.embed_dim, elementwise_affine=config.layer_norm_elementwise_affine
167
)
168
@@ -246,8 +247,9 @@ class OPTDecoder(nn.Module):
247
# keep backward compatibility with checkpoints that have been fine-tuned
248
# before transformers v4.20.1
249
# see https://github.com/facebookresearch/metaseq/pull/164
250
+ self.final_layer_norm: LayerNorm | None
251
if config.do_layer_norm_before and not config._remove_final_layer_norm:
250
- self.final_layer_norm = nn.LayerNorm(
252
+ self.final_layer_norm = LayerNorm(
253
config.hidden_size,
254
elementwise_affine=config.layer_norm_elementwise_affine,
255
)
vllm/model_executor/models/siglip.py
+7
-5
@@ -25,6 +25,7 @@ from vllm.model_executor.layers.attention import (
25
MMEncoderAttention,
26
)
27
from vllm.model_executor.layers.conv import Conv2dLayer
28
+from vllm.model_executor.layers.layernorm import LayerNorm
29
from vllm.model_executor.layers.linear import (
30
ColumnParallelLinear,
31
QKVParallelLinear,
@@ -467,13 +468,13 @@ class SiglipEncoderLayer(nn.Module):
468
prefix=f"{prefix}.self_attn",
469
attn_cls=attn_cls,
470
)
470
- self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
471
+ self.layer_norm1 = LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
472
self.mlp = SiglipMLP(
473
config,
474
quant_config=quant_config,
475
prefix=f"{prefix}.mlp",
476
)
476
- self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
477
+ self.layer_norm2 = LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
478
479
def forward(
480
self,
@@ -573,7 +574,7 @@ class SiglipTextTransformer(nn.Module):
574
attn_cls=EncoderOnlyAttention,
575
)
576
576
- self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
577
+ self.final_layer_norm = LayerNorm(embed_dim, eps=config.layer_norm_eps)
578
self.head = nn.Linear(embed_dim, config.projection_size)
579
580
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
@@ -616,7 +617,7 @@ class SiglipMultiheadAttentionPoolingHead(nn.Module):
617
self.attention = torch.nn.MultiheadAttention(
618
config.hidden_size, config.num_attention_heads, batch_first=True
619
)
619
- self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
620
+ self.layernorm = LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
621
self.mlp = SiglipMLP(
622
config=config,
623
quant_config=quant_config,
@@ -685,8 +686,9 @@ class SiglipVisionTransformer(nn.Module):
686
if require_post_norm is None:
687
require_post_norm = len(self.encoder.layers) == num_hidden_layers
688
689
+ self.post_layernorm: LayerNorm | None
690
if require_post_norm:
689
- self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
691
+ self.post_layernorm = LayerNorm(embed_dim, eps=config.layer_norm_eps)
692
else:
693
self.post_layernorm = None
694
vllm/model_executor/models/transformers/base.py
+4
@@ -71,6 +71,7 @@ from vllm.model_executor.models.transformers.utils import (
71
named_state,
72
replace_conv_class,
73
replace_embedding_class,
74
+ replace_layernorm_class,
75
replace_linear_class,
76
)
77
from vllm.model_executor.models.utils import (
@@ -466,6 +467,7 @@ class Base(
467
- `nn.Conv2d` / `nn.Conv3d` with vLLM's `Conv2d` / `Conv3d`
468
- Vocab `nn.Embedding`s with vLLM's `VocabParallelEmbedding`
469
- RMSNorm (detected from their dataflow) with vLLM's `RMSNorm`or `GemmaRMSNorm`
470
+ - `nn.LayerNorm` with vLLM's `LayerNorm`
471
"""
472
tp_plan = self.model.tp_plan or {}
473
@@ -556,6 +558,8 @@ class Base(
558
new_module = replace_embedding_class(
559
child_module, self.quant_config, prefix=qual_name
560
)
561
+ elif isinstance(child_module, nn.LayerNorm):
562
+ new_module = replace_layernorm_class(child_module)
563
elif child_module_fusers := fusers[child_module]:
564
for fuser in child_module_fusers:
565
register_fusion(fuser, qual_name, child_module)
vllm/model_executor/models/transformers/utils.py
+29
@@ -28,6 +28,7 @@ from torch import nn
28
29
from vllm.logger import init_logger
30
from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer
31
+from vllm.model_executor.layers.layernorm import LayerNorm
32
from vllm.model_executor.layers.linear import (
33
ColumnParallelLinear,
34
ReplicatedLinear,
@@ -187,6 +188,34 @@ def replace_conv_class(conv: TorchConv) -> VllmConv | TorchConv:
188
)
189
190
191
+def replace_layernorm_class(layernorm: nn.LayerNorm) -> nn.Module:
192
+ """Replace a standard (mean-centered) `nn.LayerNorm` with vLLM's
193
+ `LayerNorm`.
194
+
195
+ Args:
196
+ layernorm: `nn.LayerNorm` to be replaced.
197
+
198
+ Returns:
199
+ The new `LayerNorm`. If the layernorm is not supported (a subclass with
200
+ its own behavior, multi-dim `normalized_shape`, or not both
201
+ elementwise-affine and biased), returns the original module unchanged.
202
+
203
+ """
204
+ if (
205
+ type(layernorm) is not nn.LayerNorm
206
+ or len(layernorm.normalized_shape) != 1
207
+ or not layernorm.elementwise_affine
208
+ or layernorm.bias is None
209
+ ):
210
+ return layernorm
211
+
212
+ return LayerNorm(
213
+ layernorm.normalized_shape[0],
214
+ eps=layernorm.eps,
215
+ dtype=layernorm.weight.dtype,
216
+ )
217
+
218
+
219
def attrsetter(attr: str) -> Callable[[object, object], None]:
220
"""Set a possibly nested attribute, like the inverse of attrgetter."""
221
parent, _, name = attr.rpartition(".")
vllm/models/deepseek_v32/attention.py
+1
-1
@@ -84,7 +84,7 @@ class DeepseekV32Indexer(nn.Module):
84
disable_tp=True,
85
prefix=f"{prefix}.wk_weights_proj",
86
)
87
- self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
87
+ self.k_norm = LayerNorm(self.head_dim, eps=1e-6, dtype=torch.float32)
88
self.softmax_scale = self.head_dim**-0.5
89
90
self.scale_fmt = "ue8m0"
vllm/models/glm5next/common/attention.py
+1
-1
@@ -276,7 +276,7 @@ class Indexer(nn.Module):
276
disable_tp=True,
277
prefix=f"{prefix}.wk_weights_proj",
278
)
279
- self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
279
+ self.k_norm = LayerNorm(self.head_dim, eps=1e-6, dtype=torch.float32)
280
self.softmax_scale = self.head_dim**-0.5
281
282
# Hadamard-128 rotation of the indexer query is fused with the FP8
vllm/models/hy_v4/nvidia/attention.py
+1
-1
@@ -166,7 +166,7 @@ class Indexer(nn.Module):
166
disable_tp=True,
167
prefix=f"{prefix}.wk_weights_proj",
168
)
169
- self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
169
+ self.k_norm = LayerNorm(self.head_dim, eps=1e-6, dtype=torch.float32)
170
self.softmax_scale = self.head_dim**-0.5
171
172
self.scale_fmt = "ue8m0"