[None][feat] Add AD custom models for Jamba (Mamba v1 + Attention hybrid) by govind-ramnarayan · Pull Request #200 · nv-auto-deploy/TensorRT-LLM

govind-ramnarayan · 2026-03-11T07:30:46Z

Summary

Onboard ai21labs/AI21-Jamba-Reasoning-3B and ai21labs/AI21-Jamba-Mini-1.5 for AutoDeploy
New Mamba v1 selective scan custom op (torch_mamba_v1_selective_scan) and cached backend (torch_cached_mamba_v1) with AttentionDescriptor — Mamba v1 uses per-element state matrix A: [D_inner, d_state], fundamentally different from Mamba v2/SSD's scalar A per head
Self-contained prefill-only model (modeling_jamba.py) using torch_causal_conv1d + torch_mamba_v1_selective_scan custom ops
Bug fix: corrected mutates_args={} to declare actual cache mutations in torch_cached_mamba_v1, torch_cached_ssm, and torch_cached_causal_conv1d

Test plan

21 hierarchical equivalence tests (MLP, Attention, Mamba mixer, decoder layers, full model, export) — all pass
6 cached Mamba v1 op tests (prefill, decode, prefill→decode continuity) — all pass
AD e2e: Jamba-Reasoning-3B on 1x H100 — 9/10 coherent with chat template
AD e2e: Jamba-Mini-1.5 on 2x H100 (EP) — 10/10 coherent with chat template
AD e2e: Jamba-Mini-1.5 reduced layers (5/32) on 1x H100 — e2e flow verified

🤖 Generated with Claude Code

Onboard AI21-Jamba-Reasoning-3B and AI21-Jamba-Mini-1.5 for AutoDeploy. These hybrid models use Mamba v1 selective scan (per-element state matrix) plus GQA attention layers, with optional Sparse MoE. New files: - modeling_jamba.py: prefill-only custom model with torch_causal_conv1d and torch_mamba_v1_selective_scan custom ops - torch_mamba_v1.py: Mamba v1 selective scan reference op - torch_backend_mamba_v1.py: cached backend with AttentionDescriptor for autoregressive generation (prefill + decode paths) - test_jamba_modeling.py: 21 hierarchical equivalence tests - test_torch_mamba_v1_cached_op.py: 6 cached op tests including prefill-to-decode continuity validation - Config YAMLs for both models (Reasoning-3B single GPU, Mini-1.5 EP) Bug fix: corrected mutates_args={} to declare actual cache mutations in torch_cached_mamba_v1, torch_cached_ssm, and torch_cached_causal_conv1d custom ops. Signed-off-by: Gram Narayan <gramnarayan@nvidia.com> Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

govind-ramnarayan · 2026-03-12T01:16:03Z

[AGENT] Addressed review comments:

Config cleanup (remove defaults, minimize):

Removed skip_tokenizer_init: false and skip_loading_weights: false from all three YAML configs (both default False)
Removed max_num_tokens: 8192 from jamba_reasoning_3b.yaml and jamba_mini_1.5_reduced.yaml (equals LlmArgs default)
Removed world_size: 2 from jamba_mini_1.5.yaml — world size now comes from world_size_2.yaml in models.yaml per registry convention

models.yaml fixes:

AI21-Jamba-Reasoning-3B: corrected world_size_2.yaml → world_size_1.yaml (3B model fits on 1 GPU); added jamba_reasoning_3b.yaml reference (was missing — transforms were never applied)
Added ai21labs/AI21-Jamba-Mini-1.5 entry (was completely missing from registry)
Updated ai21labs/AI21-Jamba-1.5-Mini to also reference jamba_mini_1.5.yaml (same architecture, needs mamba v1 backend)

modeling_jamba.py:

Replaced manual JambaRMSNorm.forward with torch.ops.auto_deploy.torch_rmsnorm canonical op (consistent with other AD custom models)
Removed dead isinstance(ff_output, tuple) checks in both decoder layers (JambaSparseMoeBlock already returns a plain tensor)

All 27 unit tests pass after these changes:

pytest tests/unittest/auto_deploy/singlegpu/models/test_jamba_modeling.py \
       tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_torch_mamba_v1_cached_op.py
# 27 passed in 3.61s

- Strip all defaults from jamba_mini_1.5.yaml and jamba_reasoning_3b.yaml: trust_remote_code, enable_block_reuse (auto-disabled by AD runtime for non-paged SSM models), free_gpu_memory_fraction, max_batch_size, max_seq_len, max_num_tokens, compile_model cuda_graph settings, and chat-templated prompts - Delete jamba_mini_1.5_reduced.yaml (test-only, no longer needed) Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

govind-ramnarayan · 2026-03-12T02:25:44Z

  yaml_extra: ['dashboard_default.yaml', 'world_size_4.yaml', 'deepseek_v2_ep.yaml']
 - name: ai21labs/AI21-Jamba-1.5-Mini
-  yaml_extra: ['dashboard_default.yaml', 'world_size_2.yaml']
+  yaml_extra: ['dashboard_default.yaml', 'world_size_2.yaml', 'jamba_mini_1.5.yaml']


This seems duplicated. Can we get rid of it? I think it is not actually a different model...

govind-ramnarayan · 2026-03-12T02:26:02Z

 # Returns (seq_len, seq_start, slot_idx)
-@torch.library.custom_op("auto_deploy::torch_cached_causal_conv1d", mutates_args={})
+@torch.library.custom_op(
+    "auto_deploy::torch_cached_causal_conv1d", mutates_args={"conv_state_cache"}


I think this makes sense, but why was this working before?

AI21-Jamba-1.5-Mini and AI21-Jamba-Mini-1.5 are the same checkpoint (identical config.json). Keep only AI21-Jamba-Mini-1.5 (the canonical name per AI21's own README new_version pointer). Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

- Remove AI21-Jamba-Large-1.7 from registry (too large, no Mamba v1 backend config available yet) - torch_cached_mamba_v1: extract use_initial_states and raise ValueError for chunked prefill (consistent with torch_cached_ssm behavior) - modeling_jamba: add assert position_ids is not None in both forward methods (CA1 compliance) - test_jamba_modeling: add MoE equivalence tests for JambaSparseMoeBlock and Mamba+MoE decoder layer (num_experts=4 config) Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

govind-ramnarayan · 2026-03-13T02:13:12Z

+
+        # Infer dtype from hidden_states if available, else from A
+        hs_node = source_attn_node.args[0]
+        hs_dtype = hs_node.meta["val"].dtype if "val" in hs_node.meta else A_fake.dtype


I think this might always be there - we do not check this in the torch_backend_mamba implementation, we just access this, right?

govind-ramnarayan · 2026-03-13T02:14:18Z

+        hs_dtype = hs_node.meta["val"].dtype if "val" in hs_node.meta else A_fake.dtype
+
+        # Map to SSMResourceHandler: num_heads=D_inner, head_dim=1, d_state=d_state
+        ssm_state_dtype = cls.resolve_cache_dtype(cache_config.mamba_ssm_cache_dtype, hs_dtype)


Hmm, does this make sense? I thought that we did not have proper heads for Mamba 1- but I guess that Mamba 1 mathematically maps to Mamba2 with heads of size 1, is that correct? Since Mamba 2 has different values per head

govind-ramnarayan · 2026-03-13T02:16:21Z

+
+    @classmethod
+    def get_constants(cls, source_attn_node: Node) -> List[Constant]:
+        # No constants for Mamba v1 (unlike v2 which has time_step_limit, chunk_size)


Is the lack of chunk_size here because we are not supporting chunked prefill yet, or is it intrinsic to the op? If the former, can you clarify this comment?

govind-ramnarayan · 2026-03-13T02:21:22Z

+    )
+
+
+def _create_moe_config(num_experts=4):


Maybe you could make this cleaner by getting the jamba config with _create_small_config() and then editing certain parameters?

govind-ramnarayan · 2026-03-13T02:24:40Z

+    )
+
+
+def _create_moe_config(num_experts=4):


Check against other classes to see if this sort of config-per-test is necessary. I feel like this can just be one config with some potential overrides.

govind-ramnarayan · 2026-03-13T02:32:18Z

+def test_jamba_mamba_decoder_layer_equivalence(B, S, dtype):
+    """Test mamba decoder layer against HF."""
+    HFLayer = _get_hf_mamba_decoder_layer()
+    if HFLayer is None:


No need for this.

govind-ramnarayan · 2026-03-13T02:33:41Z

+def test_jamba_full_model_equivalence(B, S, dtype):
+    """Test full model logits against HF (slow_forward path) on CUDA."""
+    HFModel = _get_hf_jamba_model()
+    if HFModel is None:


Ditto for all the pytest skips. Let's let it error if it is missing.

govind-ramnarayan · 2026-03-13T02:34:51Z

+except ImportError:
+    from transformers import AutoConfig, PretrainedConfig
+
+    class JambaConfig(PretrainedConfig):


Is there a reason we have a fallback here? Is it possible that it is not in transformers? What am I missing?

Either way, in modeling code we don't want to create this config if we don't have to, if AutoConfig will pick it up. Let's see if we can get rid of this.

govind-ramnarayan · 2026-03-13T02:35:15Z

+                for i in range(self.num_hidden_layers)
+            ]
+
+    AutoConfig.register("jamba", JambaConfig, exist_ok=True)


Ditto, let's get rid of it.

govind-ramnarayan · 2026-03-13T02:45:47Z

+        hs, A, B, C, D, dt = _random_v1_params(device, dtype, batch, seq_len, d_inner, d_state)
+
+        # Compute reference final state using _mamba_v1_prefill directly
+        from tensorrt_llm._torch.auto_deploy.custom_ops.mamba.torch_backend_mamba_v1 import (


Put test imports on top of the file.

govind-ramnarayan · 2026-03-13T02:46:17Z

+    rtol = 1e-2
+    torch.manual_seed(42)
+    torch.cuda.empty_cache()
+    return {"device": device, "dtype": dtype, "atol": atol, "rtol": rtol}


nit: a bit overly long, just put those directly in here.

govind-ramnarayan · 2026-03-13T02:49:49Z

+    ResourceHandlerDict,
+    SSMResourceHandler,
+)
+from .torch_mamba_v1 import _torch_mamba_v1_selective_scan  # noqa: F401 — ensures op registered


This looks weird. What is going on here? We are not using the op in this file?

govind-ramnarayan · 2026-03-13T02:52:36Z



-@torch.library.custom_op("auto_deploy::torch_cached_ssm", mutates_args={})
+@torch.library.custom_op("auto_deploy::torch_cached_ssm", mutates_args={"ssm_state_cache"})


Ditto. Let us see if we can revert these mutates_args changes even if we think it might be a bug, because Mamba2 has been working well. Can you explain what this might effect?

govind-ramnarayan · 2026-03-13T02:55:17Z

+            ssm_cache,
+        )
+
+        torch.testing.assert_close(y_cached, y_ref, atol=env["atol"], rtol=env["rtol"])


Actually just use hard-coded tolerances here, it is clearer.

govind-ramnarayan · 2026-03-13T02:55:29Z

+    atol = 1e-2
+    rtol = 1e-2
+    torch.manual_seed(42)
+    torch.cuda.empty_cache()


can you explain this empty_cache() call?

govind-ramnarayan · 2026-03-13T02:56:22Z

+
+    def test_flattened_multi_sequence_prefill(self, env):
+        """Multiple sequences flattened into [1, total_len, D]: each segment matches reference."""
+        device, dtype = env["device"], env["dtype"]


Ditto, think it is just clearer to have these in the test. Maybe get rid of the env except to set random seed and similar things.

govind-ramnarayan · 2026-03-13T02:58:32Z

+            [sum(lens[:i]) for i in range(len(lens))], device=device, dtype=torch.int32
+        )
+        use_initial = torch.zeros(len(lens), device=device, dtype=torch.bool)
+        batch_info = torch.tensor([len(lens), total, 0], device=device, dtype=torch.int32)


Can you put a comment here as to what this means? There is an implicit assumption here on what the various entries of batch_info mean, and it will be helpful for later refactoring to comment what the interpretation here is.

govind-ramnarayan · 2026-03-13T03:00:59Z

+            B_seg = B[:, offset : offset + ln, :]
+            C_seg = C[:, offset : offset + ln, :]
+            dt_seg = dt[:, offset : offset + ln, :]
+            y_seg = torch.ops.auto_deploy.torch_mamba_v1_selective_scan(


I'm a little confused where the "state" is here. It almost looks like each part of y_ref is independently computed.

govind-ramnarayan · 2026-03-13T23:11:44Z

+        """Multi-batch prefill (flattened) then multi-batch decode."""
+        device, dtype = env["device"], env["dtype"]
+        d_inner, d_state = 16, 4
+        lens = [6, 4]


maybe rename to prefill_lens?

govind-ramnarayan · 2026-03-13T23:18:18Z

+
+    if s == 1:
+        # === Decode-only: single token per sequence ===
+        ssm_batch = cache_3d.index_select(0, slot_idx_t)  # [num_seq, D, N]


Can you comment here what this is doing?

govind-ramnarayan · 2026-03-13T23:18:56Z

+
+    seq_start = cu_seqlen[:num_seq]
+
+    for i in range(num_seq):


Why is this a for loop? It seems like for each batch we are doing the same thing, so can we vectorize this?

Can you see if it is a for loop in the mamba 2 cached op? If not then it seems like there should be no issue here. Also please check the PyTorch backend implementation for mamba 1 prefill as a reference.

github-actions Bot assigned govind-ramnarayan Mar 11, 2026

govind-ramnarayan force-pushed the gramnarayan/jamba branch from ff0d687 to f530ac8 Compare March 11, 2026 17:17

govind-ramnarayan added 2 commits March 11, 2026 17:25

address review

2b4a450

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

govind-ramnarayan force-pushed the gramnarayan/jamba branch from f530ac8 to 2b4a450 Compare March 12, 2026 01:14

govind-ramnarayan commented Mar 12, 2026

View reviewed changes

govind-ramnarayan added 2 commits March 11, 2026 19:50

govind-ramnarayan commented Mar 13, 2026

View reviewed changes



		@torch.library.custom_op("auto_deploy::torch_cached_ssm", mutates_args={})
		@torch.library.custom_op("auto_deploy::torch_cached_ssm", mutates_args={"ssm_state_cache"})

Uh oh!

Conversation

govind-ramnarayan commented Mar 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

govind-ramnarayan commented Mar 12, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

govind-ramnarayan Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

govind-ramnarayan commented Mar 11, 2026 •

edited

Loading

govind-ramnarayan Mar 13, 2026 •

edited

Loading