from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F from einops import repeat, rearrange from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined @dataclass class RoutingModuleOutput: # 路由模块的前向输出： # - boundary_prob: 每个位置为“分隔点/非分隔点”的二分类概率，形状(..., 2) # - boundary_mask: 基于最大概率得到的硬选择（True=分隔点），形状与输入序列相同的前两维 # - selected_probs: 对应硬选择类别的概率，形状(..., 1) boundary_prob: torch.Tensor boundary_mask: torch.Tensor selected_probs: torch.Tensor @dataclass class RoutingModuleState: """ 路由模块的推理状态（用于增量/流式step） 包含： - has_seen_tokens: (batch_size,) 是否已经见过任意token（用于首token强制边界） - last_hidden_state: (batch_size, d_model) 上一次的隐藏状态（用于与当前token做相邻相似度） """ has_seen_tokens: torch.Tensor # (batch_size,) last_hidden_state: torch.Tensor # (batch_size, d_model) @dataclass class DeChunkState: """ DeChunk 的推理状态（EMA的记忆值） 包含： - last_value: (batch_size, d_model) EMA反聚合的上一时刻值 """ last_value: torch.Tensor # (batch_size, d_model) def get_seq_idx(cu_seqlens, device=None): seq_idx = torch.zeros(cu_seqlens[-1], dtype=torch.long, device=device) seq_idx[cu_seqlens[:-1]] = 1 seq_idx = (torch.cumsum(seq_idx, dim=0) - 1).unsqueeze(0).int() return seq_idx class RoutingModule(nn.Module): """ 路由模块： 用相邻token的余弦相似度构造“成为分隔点”的概率： p_t = clamp((1 - cos(h_{t-1}, h_t)) / 2, 0, 1) 并强制首位置为边界（概率=1）。 支持： - 常规batch掩码 mask 模式 - packed序列 cu_seqlens 模式（把多序列打包成单条序列的拼接） - 流式推理 step()（维护状态） """ def __init__(self, d_model, device=None, dtype=None): self.d_model = d_model factory_kwargs = {"device": device, "dtype": dtype} super().__init__() # 相邻相似度计算前的线性投影（初始化为恒等） self.q_proj_layer = nn.Linear(d_model, d_model, bias=False, **factory_kwargs) self.k_proj_layer = nn.Linear(d_model, d_model, bias=False, **factory_kwargs) with torch.no_grad(): self.q_proj_layer.weight.copy_(torch.eye(d_model)) self.k_proj_layer.weight.copy_(torch.eye(d_model)) # 防止外部权重再初始化 self.q_proj_layer.weight._no_reinit = True self.k_proj_layer.weight._no_reinit = True def allocate_inference_cache(self, batch_size, max_seqlen, device, dtype=None): # 分配推理cache（用于step） return RoutingModuleState( has_seen_tokens=torch.zeros(batch_size, device=device, dtype=torch.bool), last_hidden_state=torch.zeros( batch_size, self.d_model, device=device, dtype=dtype ), ) def forward(self, hidden_states, cu_seqlens=None, mask=None, inference_params=None): """ hidden_states: - 若 cu_seqlens is None: (B, L, D) - 若 cu_seqlens 非 None: 期望 packed 模式 (T, D)，这里会临时扩维成 (1, T, D) cu_seqlens: packed模式下每条序列的前缀和下标，形如 [0, len1, len1+len2, ...] mask: (B, L) bool，True=有效（非packed） inference_params: RoutingModuleState，用于prefill时的校验与状态维护 """ assert (mask is not None) or ( cu_seqlens is not None ), "Either mask or cu_seqlens must be provided" if inference_params is not None: # prefill阶段必须提供mask，且不允许之前已经见过token assert ( mask is not None ), "Mask must be provided if inference_params is provided" assert ( ~inference_params.has_seen_tokens ).all(), "Cannot have seen tokens when inference_params is not provided" if cu_seqlens is not None: # packed 模式：把 (T, D) 临时变为 (1, T, D) hidden_states = hidden_states.unsqueeze(0) # 计算相邻余弦相似度 cos(h_{t-1}, h_t) cos_sim = torch.einsum( "b l d, b l d -> b l", F.normalize(self.q_proj_layer(hidden_states[:, :-1]), dim=-1), F.normalize(self.k_proj_layer(hidden_states[:, 1:]), dim=-1), ) # p = ((1 - cos) / 2) ∈ [0,1] boundary_prob = torch.clamp(((1 - cos_sim) / 2), min=0.0, max=1.0) # 强制首位置为边界：首位概率=1，补充后长度和输入序列长度想等 PAD_PROB = 1.0 boundary_prob = F.pad(boundary_prob, (1, 0), "constant", PAD_PROB) if cu_seqlens is not None: # packed 模式下，每段序列的第一个位置是 cu_seqlens[:-1] boundary_prob = boundary_prob.squeeze(0) boundary_prob[cu_seqlens[:-1]] = PAD_PROB # 组装为二分类概率 [非边界, 边界] boundary_prob = torch.stack(((1 - boundary_prob), boundary_prob), dim=-1) # 取最大概率类别 selected_idx = torch.argmax(boundary_prob, dim=-1) # 硬边界掩码 boundary_mask = selected_idx == 1 # 形状与 hidden_states 的前两维一致 if mask is not None: # 不允许选择到无效token boundary_mask = boundary_mask & mask if inference_params is not None: # 维护路由状态：是否见过token、最后一个有效token的隐藏状态 has_mask = mask.any(dim=-1) inference_params.has_seen_tokens.copy_( has_mask | inference_params.has_seen_tokens ) last_mask = torch.clamp(mask.sum(dim=-1) - 1, min=0) inference_params.last_hidden_state.copy_( torch.where( has_mask, hidden_states[ torch.arange( hidden_states.shape[0], device=hidden_states.device ), last_mask, ], inference_params.last_hidden_state, ) ) # 取硬选择对应的概率（便于可视化/正则） selected_probs = boundary_prob.gather( dim=-1, index=selected_idx.unsqueeze(-1) ) # (..., 1) return RoutingModuleOutput( boundary_prob=boundary_prob, # (..., 2) boundary_mask=boundary_mask, # (...) selected_probs=selected_probs, # (..., 1) ) def step(self, hidden_states, inference_params): """ 流式单步： hidden_states: (B, 1, D) 使用上一步缓存的 last_hidden_state 与当前token计算相邻相似度，得到当前步的边界概率 """ # (B, D) hidden_states = hidden_states.squeeze(1) cos_sim = torch.einsum( "b d, b d -> b", F.normalize(self.q_proj_layer(inference_params.last_hidden_state), dim=-1), F.normalize(self.k_proj_layer(hidden_states), dim=-1), ) boundary_prob = torch.clamp(((1 - cos_sim) / 2), min=0.0, max=1.0) # 更新最后隐藏状态 inference_params.last_hidden_state.copy_(hidden_states) # 首个token前，强制边界 boundary_prob = torch.where( inference_params.has_seen_tokens, boundary_prob, torch.ones_like(boundary_prob), ) boundary_prob = torch.stack(((1 - boundary_prob), boundary_prob), dim=-1) # 标记为已见token inference_params.has_seen_tokens.copy_( torch.ones_like(inference_params.has_seen_tokens) ) return RoutingModuleOutput( boundary_prob=boundary_prob, # (B, 2) boundary_mask=boundary_prob[..., 1] > 0.5, # (B,) selected_probs=boundary_prob.max(dim=-1).values.unsqueeze(-1), # (B, 1) ) class ChunkLayer(nn.Module): """ Chunk层：根据 boundary_mask 将被选中的“边界token”抽取出来，形成下一层序列。 支持两种模式： - packed（cu_seqlens 非 None）：直接在拼接后的序列上索引 - 非packed（mask 非 None）：通过排序 trick 把True位置排到前面，并生成 next_mask 返回： - next_hidden_states: 选中的token序列（packed: shape=(#selected, D)；非packed: (B, M, D)） - next_cu_seqlens: packed模式下新序列的cu_seqlens；否则None - next_max_seqlen: packed模式下选中的最大长度；非packed模式返回None - next_mask: 非packed模式下的右侧pad掩码；packed模式下None """ def forward(self, hidden_states, boundary_mask, cu_seqlens=None, mask=None): assert (mask is not None) or ( cu_seqlens is not None ), "Either mask or cu_seqlens must be provided" if cu_seqlens is not None: # packed：直接选择True的行，得到拼接后的 selected next_hidden_states = hidden_states[boundary_mask] # 新的cu_seqlens = 对每段最后一个位置(=cu_seqlens[1:]-1)累计True的计数，再前置0 next_cu_seqlens = F.pad( boundary_mask.cumsum(dim=0)[cu_seqlens[1:] - 1], (1, 0) ) # 新序列的最大段长（仅用于内核/优化） next_max_seqlen = int((next_cu_seqlens[1:] - next_cu_seqlens[:-1]).max()) next_mask = None else: # 非packed：对每个batch内，把True位置排到前面（False放到靠后） next_cu_seqlens = None num_tokens = boundary_mask.sum(dim=-1) # 每个样本被选中的数量 next_max_seqlen = int(num_tokens.max()) device = hidden_states.device L = hidden_states.shape[1] # trick：用 (~boundary_mask)*L 把False加大，从而 argsort 后 True 的下标排在前面 token_idx = ( torch.arange(L, device=device)[None, :] + (~boundary_mask).long() * L ) seq_sorted_indices = torch.argsort(token_idx, dim=1) # 收集前 next_max_seqlen 个（不足的样本右侧pad） next_hidden_states = torch.gather( hidden_states, dim=1, index=seq_sorted_indices[:, :next_max_seqlen, None].expand( -1, -1, hidden_states.shape[-1] ), ) # 下游的有效mask（右侧pad无效） next_mask = ( torch.arange(next_max_seqlen, device=device)[None, :] < num_tokens[:, None] ) # 非packed模式下，不再需要 max_seqlen（返回None） next_max_seqlen = None return next_hidden_states, next_cu_seqlens, next_max_seqlen, next_mask def step(self, hidden_states, boundary_mask): # 流式step：仅返回当前步被选中的token（用于下一层） return hidden_states[boundary_mask] class DeChunkLayer(nn.Module): """ DeChunk层：把“被选中的边界token序列”反聚合（EMA）回原始等长序列。 实现上复用 Mamba2 的 Triton 扫描核 mamba_chunk_scan_combined： - 将 d_model 切分为 nheads * headdim - 使用参数 A=-1, b=p, c=1 的一阶状态空间/EMA形式进行前向扫描 - 最终把扫描输出根据分段索引映射回原位置（plug back） 支持： - packed 模式（cu_seqlens） - 非packed（batch+右侧pad） - 流式 step（EMA递推） """ def __init__( self, d_model, dtype=torch.bfloat16, block_size=256, headdim=32, ): super().__init__() self.d_model = d_model # 仅为内核要求：使用 bfloat16，块大小与头维拆分 self.dtype = dtype self.block_size = block_size self.headdim = headdim assert d_model % self.headdim == 0 self.nheads = d_model // self.headdim def allocate_inference_cache(self, batch_size, max_seqlen, device, dtype=None): # 分配EMA的last_value缓存 return DeChunkState( last_value=torch.zeros( batch_size, self.d_model, device=device, dtype=dtype ), ) def forward( self, hidden_states, # 被选中的token序列（packed: (M, D)；非packed: (B, M, D)） boundary_mask, # 原序列上的边界掩码（(T,) 或 (B, L)） boundary_prob, # 原序列上的二分类概率（(..., 2)） cu_seqlens=None, inference_params=None, mask=None, ): """ 核心思路： 1) 从 boundary_prob 得到 p = P(boundary) ∈ (1e-4, 1-1e-4) 2) 构造 dt = log(1 / (1-p))，并对输入做缩放 x = h / dt 3) 用 mamba_chunk_scan_combined 扫描：A=-1, b=p, c=1，对 (B, M, H, P) 进行块扫描 4) 将结果根据 cumulative boundary index 回填到原序列位置 """ if inference_params is not None: # prefill时必须有mask，且首token必须是边界（保证EMA初始化） assert ( mask is not None ), "Mask must be provided if inference_params is provided" assert boundary_mask[ :, 0 ].all(), "First token must be a boundary if running prefill" # 取边界概率的“边界类”概率 p，并限制在(1e-4, 1-1e-4)内，避免数值不稳 p = torch.clamp(boundary_prob[..., -1].float(), min=1e-4, max=1 - (1e-4)) if cu_seqlens is not None: # packed：从原序列p中取出被选中的位置对应的概率，形状(B=1, M) p = p[boundary_mask].unsqueeze(0) # 为triton核准备packed序列的索引映射 seq_idx = get_seq_idx(cu_seqlens, device=hidden_states.device) else: B, L = boundary_mask.shape seq_idx = None # 与ChunkLayer一致的排序 trick，得到选中的顺序（True在前） token_idx = ( torch.arange(L, device=hidden_states.device)[None, :] + (~boundary_mask).long() * L ) seq_sorted_indices = torch.argsort(token_idx, dim=1) # 取出与 hidden_states 对应长度的 p（(B, M)） p = torch.gather( p, dim=1, index=seq_sorted_indices[:, : hidden_states.shape[1]] ) # (B, M) original_dtype = hidden_states.dtype # 构造 EMA 扫描所需变量 dt = torch.log(1 / (1 - p)).to(self.dtype) # (B, M) x = (hidden_states / dt[..., None]).to(self.dtype) # (B, M, D) / (B, M, 1) # A, b, c 分别对应一阶状态空间/EMA的参数 A = -torch.ones( (self.nheads,), device=hidden_states.device, dtype=torch.float32 ) b = p.to(self.dtype) c = torch.ones_like(b) # 调用triton核进行块扫描 out = mamba_chunk_scan_combined( rearrange(x, "b l (h p) -> b l h p", p=self.headdim), # (B, M, H, P) repeat(dt, "b l -> b l h", h=self.nheads), # (B, M, H) A, # (H,) rearrange(b, "b l -> b l 1 1"), # (B, M, 1, 1) rearrange(c, "b l -> b l 1 1"), # (B, M, 1, 1) chunk_size=self.block_size, seq_idx=seq_idx, # packed时提供 ) out = rearrange(out, "b l h p -> b l (h p)") # (B, M, D) # 将扫描结果回填（plug back）到原序列位置 if cu_seqlens is not None: out = out.squeeze(0) # (M, D) plug_back_idx = boundary_mask.cumsum(dim=0) - 1 # (T,) out = torch.gather( out, dim=0, index=plug_back_idx.unsqueeze(-1).expand(-1, self.d_model) ) # (T, D) else: plug_back_idx = torch.cumsum(boundary_mask, dim=1) - 1 # (B, L) out = torch.gather( out, dim=1, index=plug_back_idx.unsqueeze(-1).expand(-1, -1, self.d_model), ) # (B, L, D) # 更新流式缓存 if inference_params is not None: inference_params.last_value.copy_(out[:, -1]) return out.to(original_dtype) def step(self, hidden_states, boundary_mask, boundary_prob, inference_params): """ 流式单步 EMA 反聚合： hidden_states: (B', 1, D)，其中 B' = 当前步被选中的数量（boundary_mask.sum()） boundary_mask: (B,) 当前batch哪些位置被选中为边界 boundary_prob: (B, 2) 当前batch各位置的边界概率 输出：(B, 1, D)，对应对所有位置做了一步 EMA 更新后的值 """ B = boundary_mask.shape[0] D = hidden_states.shape[-1] # 构造当前步每个位置的 p（未被选中的位置 p=0） p = torch.zeros(B, device=hidden_states.device, dtype=hidden_states.dtype) p[boundary_mask] = boundary_prob[boundary_mask, -1].clamp( min=1e-4, max=1 - (1e-4) ) # 构造当前被选中的隐藏状态（未选中为0） current_hidden_states = torch.zeros( B, D, device=hidden_states.device, dtype=hidden_states.dtype ) current_hidden_states[boundary_mask] = hidden_states.squeeze(1) # EMA：result = p * x + (1 - p) * last result = p * current_hidden_states + (1 - p) * inference_params.last_value inference_params.last_value.copy_(result) return result.unsqueeze(1)