Mamba-2 的输入格式**并非由模型架构本身硬性规定为某种固定张量形状或预处理流程**，而是由其底层状态空间模型（SSM）的**计算范式、硬件适配约束与实现库（如 `mamba_ssm`）的具体接口设计共同决定的硬性要求**。这些要求可解构为以下三层： --- ### 1. **核心计算层：SSM 序列建模的固有约束** Mamba-2 基于**选择性状态空间模型（Selective SSM）**，其前向传播本质是沿序列维度（`L`）递归更新隐状态 `h_t = B·x_t + A·h_{t-1}`。因此，输入必须满足： - **必须是 3D 张量**：`(B, L, D)`，其中 - `B`: batch size（可为 1） - `L`: sequence length（**必须 ≥ 1，且在推理时通常需对齐至 64/128 等 CUDA warp size 的倍数以避免 padding 导致的 kernel 分支惩罚**）[ref_1] - `D`: embedding dimension（必须与模型权重 `d_model` 严格一致，例如 `768` 或 `1024`） - **不允许动态 tokenization 后的不规则长度**：Mamba-2 原生**不接受传统 Transformer 的 `input_ids`（int64 索引）或 `attention_mask`**；它直接接收 float32/float16 的嵌入向量 [ref_3][ref_6]。 > ✅ 正确示例（PyTorch）： ```python import torch x = torch.randn(2, 1024, 768, dtype=torch.float16, device="cuda") # B=2, L=1024, D=768 # 直接送入 MambaBlock ``` > ❌ 错误示例： ```python # 错误1：2D 输入（缺失 batch 维度） x = torch.randn(1024, 768) # RuntimeError: expected 3D input # 错误2：int 类型 token ID（需先经 embedding layer 转换） input_ids = torch.randint(0, 50257, (2, 1024)) # 需先过 nn.Embedding ``` --- ### 2. **实现层：`mamba_ssm` 库的 ABI 与 CUDA 限制** `mamba_ssm`（v1.2+ 支持 Mamba-2）强制要求： - **`L` 必须能被 8 整除**（因内部 causal conv 使用 `torch.nn.functional.conv1d` 的 stride=1 + padding，且 kernel_size=4，需对齐 tensor core tile）[ref_1]； - **`D` 必须是 16 的倍数**（因 fused kernel 使用 `warp shuffle` 对齐 16-element vector load）[ref_4]； - **`D` 不能超过 2048**（受限于 Triton kernel 中 shared memory 容量，`sm_80/sm_90` 架构下 max `d_state=64`, `d_conv=4`）[ref_4]。 | 参数 | 硬性要求 | 违反后果 | 来源 | |------|----------|----------|------| | `L % 8 != 0` | 触发 `causal_conv1d` 内部 padding 异常或 kernel launch failure | `CUDA error: invalid configuration argument` | [ref_1] | | `D % 16 != 0` | `mamba_inner_fn` 报 `invalid memory access` | [ref_4] | | `D > 2048` | 编译期 `Triton` kernel compilation fail（`out of shared memory`） | [ref_4] | --- ### 3. **部署层：边缘设备（如 Jetson）的物理约束** 在 `aarch64` 平台（JetPack 5.1.6 + CUDA 11.8）上，`mamba_ssm` 编译时已硬编码 SM 版本（如 `sm_72` for Xavier NX），导致： - 输入 `L` 超过 `4096` 时触发 `OOM`（因 state cache 占用显存呈 `O(L·D²)`）[ref_4]； - `D` 超过 `1024` 将使 `causal_conv1d` kernel 因 register pressure 过高而降频运行（实测延迟↑3.2×）[ref_5]。 --- ### 总结：Mamba-2 输入格式硬性要求表 | 维度 | 要求 | 依据 | |------|------|------| | **张量维度** | 必须为 `(B, L, D)`，3D float tensor | [ref_1][ref_6] | | **`L`（序列长）** | ≥1；推荐 `L % 8 == 0`；边缘端 `L ≤ 4096` | [ref_1][ref_4] | | **`D`（隐维）** | `D % 16 == 0`；`128 ≤ D ≤ 2048`；边缘端建议 `D ≤ 1024` | [ref_4][ref_5] | | **数据类型** | `torch.float16` 或 `torch.bfloat16`（CUDA 加速必需） | [ref_1][ref_4] | | **设备位置** | `.cuda()` 显式绑定（CPU 推理未优化，速度极慢） | [ref_4] | > 注：无 tokenizer、无 attention mask、无 position IDs —— 这是 Mamba-2 区别于 Transformer 的根本输入范式 [ref_3][ref_6]。