Skip to content

Fix Zamba2MambaMixer ignoring use_mamba_kernels=False #44853

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Cyrilvallez merged 7 commits into from
Apr 20, 2026
+102 −72
Merged

Fix Zamba2MambaMixer ignoring use_mamba_kernels=False #44853

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
4c2db4b
Fix NemotronH ignoring use_mamba_kernels=False
sergiopaniego Mar 19, 2026
5f0b11d
Move to zamba2
sergiopaniego Mar 19, 2026
8d7af3f
Merge branch 'main' into nemotron-mamba-kernel
sergiopaniego Mar 23, 2026
1a9d717
Merge branch 'main' into nemotron-mamba-kernel
sergiopaniego Apr 9, 2026
93840ed
moved to config
sergiopaniego Apr 9, 2026
fc48cae
Merge branch 'main' into nemotron-mamba-kernel
sergiopaniego Apr 9, 2026
ae83160
Merge branch 'main' into nemotron-mamba-kernel
sergiopaniego Apr 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 33 additions & 24 deletions src/transformers/models/nemotron_h/modeling_nemotron_h.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,34 +179,43 @@ def __init__(self, config: NemotronHConfig, layer_idx: int | None = None):
self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)

global causal_conv1d_update, causal_conv1d_fn
causal_conv1d = lazy_load_kernel("causal-conv1d")
causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", None)
causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", None)

global selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
mamba_ssm = lazy_load_kernel("mamba-ssm")
selective_state_update = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
)
mamba_chunk_scan_combined = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
)
mamba_split_conv1d_scan_combined = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
)

global is_fast_path_available
is_fast_path_available = all(
(
selective_state_update,
mamba_chunk_scan_combined,
mamba_split_conv1d_scan_combined,
causal_conv1d_fn,
causal_conv1d_update,

if config.use_mamba_kernels:
causal_conv1d = lazy_load_kernel("causal-conv1d")
causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", None)
causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", None)

mamba_ssm = lazy_load_kernel("mamba-ssm")
selective_state_update = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
)
mamba_chunk_scan_combined = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
)
mamba_split_conv1d_scan_combined = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
)
)

if not is_fast_path_available:
is_fast_path_available = all(
(
selective_state_update,
mamba_chunk_scan_combined,
mamba_split_conv1d_scan_combined,
causal_conv1d_fn,
causal_conv1d_update,
)
)
else:
causal_conv1d_update = None
causal_conv1d_fn = None
selective_state_update = None
mamba_chunk_scan_combined = None
mamba_split_conv1d_scan_combined = None
is_fast_path_available = False

if getattr(config, "use_mamba_kernels", True) and not is_fast_path_available:
logger.warning_once(
"The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
" is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
Expand Down
3 changes: 3 additions & 0 deletions src/transformers/models/zamba2/configuration_zamba2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,8 @@ class Zamba2Config(PreTrainedConfig):
Whether or not to use bias in the convolution layer of the mixer block.
chunk_size (`int`, *optional*, defaults to 256):
Size of the chunks that will comprise the sequence.
use_mamba_kernels (`bool`, *optional*, defaults to `True`):
Flag indicating whether or not to use the fast mamba kernels.
use_mem_eff_path (`bool`, *optional*, defaults to `False`):
Whether or not to use the fused conv1d and scan in mamba2 layers.
add_bias_linear (`bool`, *optional*, defaults to `False`):
Expand Down Expand Up @@ -83,6 +85,7 @@ class Zamba2Config(PreTrainedConfig):
time_step_floor: float = 1e-4
time_step_limit: list[float] | tuple[float, ...] | None = None
n_mamba_heads: int = 8
use_mamba_kernels: bool = True
use_conv_bias: bool = True
chunk_size: int = 256
use_mem_eff_path: bool = False
Expand Down
57 changes: 33 additions & 24 deletions src/transformers/models/zamba2/modeling_zamba2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -467,34 +467,43 @@ def __init__(self, config: Zamba2Config, layer_idx: int | None = None):
self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.add_bias_linear)

global causal_conv1d_update, causal_conv1d_fn
causal_conv1d = lazy_load_kernel("causal-conv1d")
causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", None)
causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", None)

global selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
mamba_ssm = lazy_load_kernel("mamba-ssm")
selective_state_update = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
)
mamba_chunk_scan_combined = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
)
mamba_split_conv1d_scan_combined = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
)

global is_fast_path_available
is_fast_path_available = all(
(
selective_state_update,
mamba_chunk_scan_combined,
mamba_split_conv1d_scan_combined,
causal_conv1d_fn,
causal_conv1d_update,

if config.use_mamba_kernels:
causal_conv1d = lazy_load_kernel("causal-conv1d")
causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", None)
causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", None)

mamba_ssm = lazy_load_kernel("mamba-ssm")
selective_state_update = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
)
mamba_chunk_scan_combined = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
)
mamba_split_conv1d_scan_combined = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
)
)

if not is_fast_path_available:
is_fast_path_available = all(
(
selective_state_update,
mamba_chunk_scan_combined,
mamba_split_conv1d_scan_combined,
causal_conv1d_fn,
causal_conv1d_update,
)
)
else:
causal_conv1d_update = None
causal_conv1d_fn = None
selective_state_update = None
mamba_chunk_scan_combined = None
mamba_split_conv1d_scan_combined = None
is_fast_path_available = False

if getattr(config, "use_mamba_kernels", True) and not is_fast_path_available:
logger.warning_once(
"The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
" is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
Expand Down
57 changes: 33 additions & 24 deletions src/transformers/models/zamba2/modular_zamba2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -255,34 +255,43 @@ def __init__(self, config: Zamba2Config, layer_idx: int | None = None):
self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.add_bias_linear)

global causal_conv1d_update, causal_conv1d_fn
causal_conv1d = lazy_load_kernel("causal-conv1d")
causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", None)
causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", None)

global selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
mamba_ssm = lazy_load_kernel("mamba-ssm")
selective_state_update = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
)
mamba_chunk_scan_combined = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
)
mamba_split_conv1d_scan_combined = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
)

global is_fast_path_available
is_fast_path_available = all(
(
selective_state_update,
mamba_chunk_scan_combined,
mamba_split_conv1d_scan_combined,
causal_conv1d_fn,
causal_conv1d_update,

if config.use_mamba_kernels:
causal_conv1d = lazy_load_kernel("causal-conv1d")
causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", None)
causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", None)

mamba_ssm = lazy_load_kernel("mamba-ssm")
selective_state_update = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
)
mamba_chunk_scan_combined = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
)
mamba_split_conv1d_scan_combined = resolve_internal_import(
mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
)
)

if not is_fast_path_available:
is_fast_path_available = all(
(
selective_state_update,
mamba_chunk_scan_combined,
mamba_split_conv1d_scan_combined,
causal_conv1d_fn,
causal_conv1d_update,
)
)
else:
causal_conv1d_update = None
causal_conv1d_fn = None
selective_state_update = None
mamba_chunk_scan_combined = None
mamba_split_conv1d_scan_combined = None
is_fast_path_available = False

if getattr(config, "use_mamba_kernels", True) and not is_fast_path_available:
logger.warning_once(
"The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
" is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
Expand Down
Loading