Hugging Face's logo

LiquidAI
/
d1-2B

Image-Text-to-Text
Transformers
Safetensors
lfm3_vl
feature-extraction
liquid
lfm3
edge
decision
classification
extraction
calibration
system-one
multimodal
decision-model
custom_code
conversational
Model card Files
xet
 Community
d1-2B / integrations_lfm3.py
paulpak58's picture
paulpak58
mlabonne's picture
mlabonne
Aurelien-Lac's picture
Aurelien-Lac
d1-2B
97c5704
Raw History Blame Contribute Delete
8.74 kB
# Vendored from transformers' LFM3 code (not in a release yet) at 69da1b47a11e362fbb15bfe5d5c8da0790c688c5, src/transformers/integrations/lfm3.py:
# its relative imports made absolute, nothing else changed.
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Inference kernels for LFM3's width-7 depthwise convolution."""
import torch
from transformers.utils import is_triton_available
if is_triton_available():
import triton
import triton.language as tl
# the length and strides are runtime arguments, or every prompt length would compile its own kernel
@triton.jit(do_not_specialize=["T", "X0", "X1", "X2", "S0", "S1", "ST"])
def _prefill(
X,
W,
BIAS,
SEQ,
Y,
T,
D: tl.constexpr,
X0,
X1,
X2,
S0,
S1,
ST,
HAS_BIAS: tl.constexpr,
PACKED: tl.constexpr,
BLOCK: tl.constexpr,
):
# int64 offsets: batch * channels * length can exceed 2**31 elements
b, d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
t = tl.program_id(2).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
acc = tl.full((BLOCK,), 0, tl.float32)
if PACKED:
doc = tl.load(SEQ + b * S0 + (t - T + ST) * S1, (t < T) & (t >= T - ST), other=-1)
for k in tl.static_range(7):
pos = t + k - 6
valid = (t < T) & (pos >= 0)
if PACKED:
prev = tl.load(SEQ + b * S0 + (pos - T + ST) * S1, valid & (pos >= T - ST), other=-1)
valid = valid & (doc == prev)
x = tl.load(X + b * X0 + d * X1 + pos * X2, valid, other=0).to(tl.float32)
w = tl.load(W + d * 7 + k).to(tl.float32)
acc += x * w
if HAS_BIAS:
acc += tl.load(BIAS + d).to(tl.float32)
tl.store(Y + (b * D + d) * T + t, acc, t < T)
@triton.jit
def _decode(
X,
W,
BIAS,
STATE,
Y,
D: tl.constexpr,
K: tl.constexpr,
X0: tl.constexpr,
X1: tl.constexpr,
S0: tl.constexpr,
S1: tl.constexpr,
S2: tl.constexpr,
HAS_BIAS: tl.constexpr,
BLOCK: tl.constexpr,
):
b = tl.program_id(0).to(tl.int64)
d = tl.program_id(1).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
k = tl.arange(0, 8)
# Load the complete state before writing the shifted window.
old = tl.load(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, (d[:, None] < D) & (k[None, :] < K), other=0)
x = tl.load(X + b * X0 + d * X1, d < D, other=0)
shifted = tl.gather(old, tl.broadcast_to(tl.minimum(k + 1, 7)[None, :], (BLOCK, 8)), 1)
shifted = tl.where(k[None, :] == K - 1, x[:, None], shifted)
# A cache may store either K=6 previous inputs or K=7 including the oldest discarded input.
positions = k + K - 6
values = tl.gather(old, tl.broadcast_to(tl.minimum(positions, 7)[None, :], (BLOCK, 8)), 1)
values = tl.where(k[None, :] == 6, x[:, None], values)
weights = tl.load(W + d[:, None] * 7 + k[None, :], (d[:, None] < D) & (k[None, :] < 7), other=0)
acc = tl.sum(values.to(tl.float32) * weights.to(tl.float32), 1)
if HAS_BIAS:
acc += tl.load(BIAS + d, d < D, other=0).to(tl.float32)
tl.store(STATE + b * S0 + d[:, None] * S1 + k[None, :] * S2, shifted, (d[:, None] < D) & (k[None, :] < K))
tl.store(Y + b * D + d, acc, d < D)
def can_use_lfm3_conv(hidden_states, weight):
return (
is_triton_available()
and hasattr(tl, "gather") # the decode kernel needs Triton >= 3.3
and hidden_states.is_cuda
and not torch.is_grad_enabled()
and hidden_states.dtype == weight.dtype
and hidden_states.dtype in (torch.float16, torch.bfloat16, torch.float32)
and weight.shape[-1] == 7
)
def lfm3_causal_conv1d(hidden_states, weight, bias=None, seq_idx=None):
batch, channels, length = hidden_states.shape
# the kernels index every tensor by the batch and channel of `hidden_states` and cannot bound the others
if weight.shape != (channels, 7) or (bias is not None and bias.shape != (channels,)):
raise ValueError(
f"`weight` must have shape ({channels}, 7) and `bias` ({channels},), got {tuple(weight.shape)}"
)
if seq_idx is not None and (
seq_idx.ndim != 2 or seq_idx.shape[0] not in (1, batch) or seq_idx.dtype not in (torch.int32, torch.int64)
):
raise ValueError(
f"`seq_idx` must be an int32 or int64 tensor of shape (1, length) or ({batch}, length), got "
f"{seq_idx.dtype} {tuple(seq_idx.shape)}"
)
if any(t is not None and t.device != hidden_states.device for t in (weight, bias, seq_idx)):
raise ValueError(
f"`weight`, `bias` and `seq_idx` must be on the device of `hidden_states` ({hidden_states.device})"
)
output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
# Triton launches on the current device; the model may live on another one (`device_map`, `.to("cuda:1")`).
with torch.cuda.device(hidden_states.device):
_prefill[(batch, channels, triton.cdiv(length, 128))](
hidden_states,
# the kernels index `bias` and `weight` densely
weight.contiguous(),
bias.contiguous() if bias is not None else None,
seq_idx,
output,
length,
channels,
*hidden_states.stride(),
*(
(0 if seq_idx.shape[0] == 1 else seq_idx.stride(0), seq_idx.stride(1))
if seq_idx is not None
else (0, 0)
),
seq_idx.shape[-1] if seq_idx is not None else 0,
bias is not None,
seq_idx is not None,
128,
enable_fp_fusion=False,
)
return output
def lfm3_causal_conv1d_update(hidden_states, conv_state, weight, bias=None):
batch, channels, length = hidden_states.shape
# the kernel addresses one position per row and bounds nothing by the other tensors' shapes
if length != 1:
raise ValueError(f"`hidden_states` must hold one position per row, got {tuple(hidden_states.shape)}")
# the 8-wide tile holds a 6- or 7-wide state (see `positions` in `_decode`)
if (
conv_state.ndim != 3
or conv_state.shape[:2] != (batch, channels)
or conv_state.shape[-1] not in (6, 7)
or weight.shape != (channels, 7)
):
raise ValueError(
f"`conv_state` must have shape ({batch}, {channels}, 6 or 7) and `weight` ({channels}, 7), got "
f"{tuple(conv_state.shape)} and {tuple(weight.shape)}"
)
if bias is not None and bias.shape != (channels,):
raise ValueError(f"`bias` must have shape ({channels},), got {tuple(bias.shape)}")
# the kernel writes the state in place: rows or channels sharing memory (an `expand`ed cache) would race
extent = 0
for stride, size in sorted(zip(conv_state.stride(), conv_state.shape)):
if size > 1 and stride <= extent:
raise ValueError(f"`conv_state` must not overlap itself, got strides {conv_state.stride()}")
extent += stride * (size - 1)
if any(t is not None and t.device != hidden_states.device for t in (conv_state, weight, bias)):
raise ValueError(
f"`conv_state`, `weight` and `bias` must be on the device of `hidden_states` ({hidden_states.device})"
)
output = torch.empty_like(hidden_states, memory_format=torch.contiguous_format)
with torch.cuda.device(hidden_states.device):
_decode[(batch, triton.cdiv(channels, 32))](
hidden_states,
weight.contiguous(),
bias.contiguous() if bias is not None else None,
conv_state,
output,
channels,
conv_state.shape[-1],
*hidden_states.stride()[:2],
*conv_state.stride(),
bias is not None,
32,
enable_fp_fusion=False,
)
return output