import os import torch import torch.nn as nn import torch.nn.functional as F from typing import Optional, Tuple from transformers.models.qwen3.modeling_qwen3 import ( Qwen3Attention, Qwen3Config, apply_rotary_pos_emb, repeat_kv, ) try: from flash_attn import flash_attn_varlen_func except ImportError: print("请安装flash-attn库: pip install flash-attn --no-build-isolation") flash_attn_varlen_func = None class MemorySparseAttention(Qwen3Attention): def __init__(self, config: Qwen3Config, layer_idx: int): super().__init__(config=config, layer_idx=layer_idx) if flash_attn_varlen_func is None: raise ImportError("flash_attn is required. Please install it via 'pip install flash-attn --no-build-isolation'") self.layer_idx = layer_idx self.top_k_docs = config.msa_config.top_k_docs self.pooling_kernel_size = config.msa_config.pooling_kernel_size self.router_layer_idx = config.msa_config.router_layer_idx if self.router_layer_idx == "all": self.router_layer_idx = list(range(config.num_hidden_layers)) else: self.router_layer_idx = [int(i) for i in self.router_layer_idx.split(",")] self.is_router_layer = self.layer_idx in self.router_layer_idx self.head_reduce_method = config.msa_config.head_reduce_method self.query_reduce_method = config.msa_config.query_reduce_method self.chunk_reduce_method = config.msa_config.chunk_reduce_method self.decouple_pooling_mode = config.msa_config.decouple_pooling_mode self.aux_loss_method = config.msa_config.aux_loss_method self.decouple_router = config.msa_config.decouple_router if self.is_router_layer and self.decouple_router: self.router_k_proj = nn.Sequential( nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False), # nn.GELU(), # nn.Linear(config.num_key_value_heads * self.head_dim, config.num_key_value_heads * self.head_dim, bias=False) ) self.router_q_proj = nn.Sequential( nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False), # nn.GELU(), # nn.Linear(config.num_attention_heads * self.head_dim, config.num_attention_heads * self.head_dim, bias=False) ) self.num_kv_heads = config.num_key_value_heads self.sliding_window = None self.selected_docs_indices = None self.max_doc_id = None self.num_split_for_kv = 8 self.template_prefix_kcache = None self.template_prefix_vcache = None self.memory_client = None def set_memory_client(self, memory_client): self.memory_client = memory_client def forward( self, hidden_states: torch.Tensor, doc_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: if self.training: return self._forward( hidden_states, doc_ids, attention_mask, position_embeddings, past_key_value, **kwargs, ) elif past_key_value is not None: return self.forward_with_kvcache_for_batch_parrallel( hidden_states, doc_ids, attention_mask, position_embeddings, past_key_value, **kwargs, ) else: raise Exception("error!") @staticmethod def map_tensor_to_group_ids(a: torch.Tensor) -> torch.Tensor: if a.ndim != 1: raise ValueError("输入 Tensor a 必须是一维的。") diff_mask = torch.diff(a) != 0 # [L-1] id_increments = diff_mask.int() # [L-1] group_indices_offset = torch.cumsum(id_increments, dim=0) # [L-1] b = torch.cat(( torch.tensor([0], device=a.device, dtype=a.dtype), group_indices_offset )) + 1 return b def forward_with_kvcache_for_batch_parrallel( self, hidden_states: torch.Tensor, doc_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: bsz, q_len, _ = hidden_states.shape device, dtype = hidden_states.device, hidden_states.dtype hidden_shape = (bsz, q_len, -1, self.head_dim) query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) stage = past_key_value.cache_kwargs[self.layer_idx]["stage"] if stage == "prefill_stage1": max_doc_id = int(doc_ids.max().item()) doc_token_mask = (doc_ids > 0) & (attention_mask == 1) doc_indices = torch.nonzero(doc_token_mask, as_tuple=False) original_doc_ids = doc_ids[doc_token_mask] original_doc_batch_indices = doc_indices[:, 0] global_doc_ids = original_doc_batch_indices * (max_doc_id + 1) + original_doc_ids if self.is_router_layer: _, counts = torch.unique_consecutive(global_doc_ids, return_counts=True) total_doc_tokens = global_doc_ids.shape[0] cu_seqlens = counts.cumsum(0) offsets = torch.zeros(counts.shape[0] + 1, dtype=counts.dtype, device=device) offsets[1:] = cu_seqlens offsets = offsets[:-1] expanded_offsets = torch.repeat_interleave(offsets, counts) original_order_ranks = torch.arange(total_doc_tokens, device=device) - expanded_offsets chunk_indices = original_order_ranks // self.pooling_kernel_size max_chunks_per_doc = (q_len // self.pooling_kernel_size) + 1 global_chunk_ids = global_doc_ids * max_chunks_per_doc + chunk_indices unique_global_chunk_ids, chunk_token_counts = torch.unique_consecutive(global_chunk_ids, return_counts=True) pooled_doc_ids = unique_global_chunk_ids // max_chunks_per_doc % (max_doc_id + 1) pooled_k_chunks, pooled_v_chunks = self.sequence_pooling_kv( key_states, value_states, doc_indices, global_chunk_ids, ) pooled_k_chunks = pooled_k_chunks.transpose(0, 1).unsqueeze(0) pooled_v_chunks = pooled_v_chunks.transpose(0, 1).unsqueeze(0) pooled_router_k = None if self.decouple_router: r_k_raw = self.router_k_proj(hidden_states).view(hidden_shape).transpose(1, 2) r_k_docs = r_k_raw[doc_indices[:, 0], :, doc_indices[:, 1]] _, chunk_lengths = torch.unique_consecutive(global_chunk_ids, return_counts=True) chunk_counts_view = chunk_lengths.view(-1, 1, 1).to(dtype=torch.float32) b_k, h_k, d_k = r_k_docs.shape k_flat = r_k_docs.reshape(b_k, -1).to(dtype=torch.float32) k_cumsum = F.pad(torch.cumsum(k_flat, dim=0), (0, 0, 1, 0)) chunk_cu_seqlens = F.pad(torch.cumsum(chunk_lengths, 0), (1, 0)) k_sums_flat = k_cumsum[chunk_cu_seqlens[1:]] - k_cumsum[chunk_cu_seqlens[:-1]] pooled_router_k = (k_sums_flat.view(unique_global_chunk_ids.shape[0], h_k, d_k) / chunk_counts_view).to(dtype=r_k_docs.dtype) pooled_router_k = pooled_router_k.transpose(0, 1).unsqueeze(0) if self.aux_loss_method == "INFONCE": router_k = pooled_router_k if pooled_router_k is not None else pooled_k_chunks pooled_router_k = F.normalize(router_k, p=2, dim=-1) if past_key_value is not None: num_template_mask_prefix = (doc_ids == -2).sum() template_prefix_kcache = key_states[:, :, :num_template_mask_prefix] template_prefix_vcache = value_states[:, :, :num_template_mask_prefix] kwargs = { "template_prefix_kcache": template_prefix_kcache, "template_prefix_vcache": template_prefix_vcache, } if self.is_router_layer: pooled_k_chunks, pooled_v_chunks = past_key_value.update(pooled_k_chunks, pooled_v_chunks, self.layer_idx) kwargs2 = { "doc_id_bias": doc_ids.shape[1], "pooled_doc_ids": pooled_doc_ids, "prefill_stage1_kvcache_size": pooled_k_chunks.shape[2], } if pooled_router_k is not None: past_key_value.update_router_kcache(pooled_router_k, self.layer_idx) kwargs.update(kwargs2) past_key_value.record_kwargs(self.layer_idx, kwargs) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) attn_output = torch.zeros((bsz, q_len, self.config.num_attention_heads * self.head_dim), device=device, dtype=dtype) indices_b = torch.nonzero(doc_token_mask, as_tuple=False) if indices_b.shape[0] > 0: q_b, k_b, v_b = query_states[indices_b[:, 0], :, indices_b[:, 1]], key_states[indices_b[:, 0], :, indices_b[:, 1]], value_states[indices_b[:, 0], :, indices_b[:, 1]] doc_ids_b = doc_ids[indices_b[:, 0], indices_b[:, 1]] batch_indices_b = indices_b[:, 0] global_doc_ids_b = batch_indices_b * (max_doc_id + 1) + doc_ids_b _, counts_b = torch.unique_consecutive(global_doc_ids_b, return_counts=True) cu_seqlens_b = F.pad(torch.cumsum(counts_b, dim=0, dtype=torch.int32), (1, 0)) output_b_flat = flash_attn_varlen_func(q_b, k_b, v_b, cu_seqlens_q=cu_seqlens_b, cu_seqlens_k=cu_seqlens_b, max_seqlen_q=int(counts_b.max()), max_seqlen_k=int(counts_b.max()), dropout_p=self.attention_dropout if self.training else 0.0, causal=True).view(-1, self.config.num_attention_heads * self.head_dim) attn_output[indices_b[:, 0], indices_b[:, 1]] += output_b_flat template_mask = (doc_ids == -2) & (attention_mask == 1) template_indices = torch.nonzero(template_mask, as_tuple=False) if template_indices.shape[0] > 0: q_template = query_states.transpose(1, 2)[template_mask] k_template = key_states.transpose(1, 2)[template_mask] v_template = value_states.transpose(1, 2)[template_mask] template_counts_per_sample = torch.bincount(template_indices[:, 0], minlength=bsz) cu_seqlens_template = F.pad(torch.cumsum(template_counts_per_sample, dim=0, dtype=torch.int32), (1, 0)) output_template_flat = flash_attn_varlen_func(q_template, k_template, v_template, cu_seqlens_q=cu_seqlens_template, cu_seqlens_k=cu_seqlens_template, max_seqlen_q=int(template_counts_per_sample.max()), max_seqlen_k=int(template_counts_per_sample.max()), dropout_p=0.0, causal=True).view(-1, self.config.num_attention_heads * self.head_dim) attn_output[template_mask] = output_template_flat return self.o_proj(attn_output), None elif stage == "prefill_stage2": cache_kwargs = past_key_value.cache_kwargs[self.layer_idx] if self.memory_client is not None: if self.template_prefix_kcache is None: self.template_prefix_kcache , self.template_prefix_vcache = self.memory_client.get_template_prefix_kvcaches(self.layer_idx) if not self.template_prefix_kcache.is_cuda: self.template_prefix_kcache = self.template_prefix_kcache.to(device) if not self.template_prefix_vcache.is_cuda: self.template_prefix_vcache = self.template_prefix_vcache.to(device) template_prefix_kcache = self.template_prefix_kcache template_prefix_vcache = self.template_prefix_vcache else: template_prefix_kcache = cache_kwargs["template_prefix_kcache"].to(device) template_prefix_vcache = cache_kwargs["template_prefix_vcache"].to(device) final_k_to_scatter, final_v_to_scatter = None, None if self.is_router_layer: routing_q_for_scoring = self.router_q_proj(hidden_states).view(hidden_shape).transpose(1, 2) if self.decouple_router else query_states if self.aux_loss_method == "INFONCE": routing_q_for_scoring = F.normalize(routing_q_for_scoring, p=2, dim=-1) query_mask = ((doc_ids == 0) & (attention_mask == 1)) res = self.memory_client.doc_query(routing_q_for_scoring, query_mask, self.layer_idx) final_k_to_scatter, final_v_to_scatter, final_scores, num_selected_chunks_per_sample, final_selected_doc_ids = res if past_key_value.meta.get("require_recall_topk", False): recall_topk_list = [] for i in range(bsz): recall_topk_list.append({ "topk_doc_ids": final_selected_doc_ids[i].cpu().detach().tolist(), "score": final_scores[i].cpu().detach().tolist(), }) cache_kwargs["recall_topk"] = recall_topk_list else: num_selected_chunks_per_sample = torch.zeros(bsz, dtype=torch.long, device=device) num_q_per_sample = attention_mask.sum(dim=1) template_len = template_prefix_kcache.shape[2] kv_lengths = template_len + num_selected_chunks_per_sample + num_q_per_sample cu_seqlens_q = F.pad(num_q_per_sample.cumsum(0, dtype=torch.int32), (1, 0)) cu_seqlens_kv = F.pad(kv_lengths.cumsum(0, dtype=torch.int32), (1, 0)) total_q_tokens = cu_seqlens_q[-1].item() total_kv_tokens = cu_seqlens_kv[-1].item() q_final = torch.empty((total_q_tokens, self.config.num_attention_heads, self.head_dim), device=device, dtype=dtype) k_final_unrepeated = torch.empty((self.config.num_key_value_heads, total_kv_tokens, self.head_dim), device=device, dtype=dtype) v_final_unrepeated = torch.empty((self.config.num_key_value_heads, total_kv_tokens, self.head_dim), device=device, dtype=dtype) valid_q_mask = (attention_mask == 1) q_final = query_states.permute(0, 2, 1, 3)[valid_q_mask] offset_start_sample = cu_seqlens_kv[:-1] offset_start_template = offset_start_sample offset_start_chunks = offset_start_sample + template_len offset_start_question = offset_start_chunks + num_selected_chunks_per_sample template_indices = torch.arange(template_len, device=device).unsqueeze(0) + offset_start_template.unsqueeze(1) source_k_template = template_prefix_kcache.expand(bsz, -1, -1, -1).permute(1, 0, 2, 3).reshape(self.config.num_key_value_heads, -1, self.head_dim) k_final_unrepeated[:, template_indices.flatten(), :] = source_k_template source_v_template = template_prefix_vcache.expand(bsz, -1, -1, -1).permute(1, 0, 2, 3).reshape(self.config.num_key_value_heads, -1, self.head_dim) v_final_unrepeated[:, template_indices.flatten(), :] = source_v_template if self.is_router_layer and final_k_to_scatter is not None and final_k_to_scatter.shape[1] > 0: batch_indices_for_chunks = torch.arange(bsz, device=device).repeat_interleave(num_selected_chunks_per_sample) is_start_of_sample = torch.cat([torch.tensor([True], device=device), batch_indices_for_chunks[1:] != batch_indices_for_chunks[:-1]]) cumsum_ranks = torch.ones_like(batch_indices_for_chunks).cumsum(0) start_offsets = cumsum_ranks[is_start_of_sample].repeat_interleave(num_selected_chunks_per_sample) chunk_rank_in_sample = cumsum_ranks - start_offsets chunk_dest_indices = offset_start_chunks[batch_indices_for_chunks] + chunk_rank_in_sample k_final_unrepeated[:, chunk_dest_indices, :] = final_k_to_scatter if final_v_to_scatter.device == torch.device("cpu"): final_v_to_scatter = final_v_to_scatter.to(device) v_final_unrepeated[:, chunk_dest_indices, :] = final_v_to_scatter batch_indices_for_q = torch.arange(bsz, device=device).repeat_interleave(num_q_per_sample) q_rank_in_sample = (torch.cumsum(valid_q_mask.int(), dim=1) - 1)[valid_q_mask] q_dest_indices = offset_start_question[batch_indices_for_q] + q_rank_in_sample k_final_unrepeated[:, q_dest_indices, :] = key_states.permute(1, 0, 2, 3).reshape(self.config.num_key_value_heads, -1, self.head_dim)[:, valid_q_mask.flatten(), :] v_final_unrepeated[:, q_dest_indices, :] = value_states.permute(1, 0, 2, 3).reshape(self.config.num_key_value_heads, -1, self.head_dim)[:, valid_q_mask.flatten(), :] k_final = k_final_unrepeated v_final = v_final_unrepeated output_flat = flash_attn_varlen_func( q=q_final, k=k_final.transpose(0,1), v=v_final.transpose(0,1), cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_kv, max_seqlen_q=num_q_per_sample.max().item(), max_seqlen_k=kv_lengths.max().item(), dropout_p=0.0, causal=True ).view(-1, self.config.num_attention_heads * self.head_dim) attn_output = torch.zeros((bsz, q_len, self.config.num_attention_heads * self.head_dim), device=device, dtype=dtype) attn_output[valid_q_mask] = output_flat max_kv_len = kv_lengths.max().item() compacked_key_cache = torch.zeros((bsz, self.config.num_key_value_heads, max_kv_len, self.head_dim), dtype=dtype, device=device) compacked_value_cache = torch.zeros((bsz, self.config.num_key_value_heads, max_kv_len, self.head_dim), dtype=dtype, device=device) left_pad_mask = torch.arange(max_kv_len, device=device).unsqueeze(0) >= (max_kv_len - kv_lengths.unsqueeze(1)) compacked_key_cache.permute(0, 2, 1, 3)[left_pad_mask] = k_final_unrepeated.permute(1, 0, 2) compacked_value_cache.permute(0, 2, 1, 3)[left_pad_mask] = v_final_unrepeated.permute(1, 0, 2) cache_kwargs["compacked_key_cache"] = compacked_key_cache cache_kwargs["compacked_value_cache"] = compacked_value_cache cache_kwargs["kv_lengths"] = kv_lengths cache_kwargs["attention_mask"] = left_pad_mask past_key_value.record_kwargs(self.layer_idx, cache_kwargs) return self.o_proj(attn_output), None else: cache_kwargs = past_key_value.cache_kwargs[self.layer_idx] if "compacked_key_cache" not in cache_kwargs: raise ValueError("批次化紧凑KV缓存未找到。Prefill stage 2 是否正确运行?") compacked_key_cache = cache_kwargs["compacked_key_cache"] compacked_value_cache = cache_kwargs["compacked_value_cache"] kv_lengths = cache_kwargs["kv_lengths"] layer_attention_mask = cache_kwargs["attention_mask"] max_kv_len = compacked_key_cache.shape[2] full_k_unrepeated = torch.cat([compacked_key_cache, key_states], dim=2) full_v_unrepeated = torch.cat([compacked_value_cache, value_states], dim=2) if past_key_value.meta.get("qa_mode", False): cur_layer_attention_mask = torch.LongTensor([[1] * q_len for _ in range(bsz)]).to(device) cur_layer_attention_mask = (cur_layer_attention_mask * attention_mask).type(layer_attention_mask.dtype) layer_attention_mask = torch.cat([layer_attention_mask, cur_layer_attention_mask], dim=1) attn_mask_4d = layer_attention_mask[:, None, None, :].expand(-1, self.config.num_attention_heads, 1, -1) cache_kwargs["attention_mask"] = layer_attention_mask else: new_kv_lengths = kv_lengths + 1 max_new_kv_len = max_kv_len + 1 attn_mask_2d = torch.arange(max_new_kv_len, device=device).unsqueeze(0) >= (max_new_kv_len - new_kv_lengths.unsqueeze(1)) attn_mask_4d = attn_mask_2d[:, None, None, :].expand(-1, self.config.num_attention_heads, 1, -1) cache_kwargs["kv_lengths"] = new_kv_lengths key_states_gqa = repeat_kv(full_k_unrepeated, self.num_key_value_groups) value_states_gqa = repeat_kv(full_v_unrepeated, self.num_key_value_groups) attn_output = F.scaled_dot_product_attention( query_states, key_states_gqa, value_states_gqa, attn_mask=attn_mask_4d, dropout_p=0.0, is_causal=False ).transpose(1, 2).reshape(bsz, q_len, -1) cache_kwargs["compacked_key_cache"] = full_k_unrepeated cache_kwargs["compacked_value_cache"] = full_v_unrepeated past_key_value.record_kwargs(self.layer_idx, cache_kwargs) return self.o_proj(attn_output), None def _calculate_routing_scores_adaptive( self, query_states: torch.Tensor, # [B, H, Q_len, D] pooled_k_bched: torch.Tensor, # [B, C, H, D] routing_query_mask: torch.Tensor, # [B, Q_len] - 1 for valid, 0 for pad chunk_mask: torch.Tensor, # [B, C] - 1 for valid, 0 for pad ) -> torch.Tensor: bsz, num_heads, q_len, head_dim = query_states.shape _, max_chunks, _, _ = pooled_k_bched.shape dtype, device = query_states.dtype, query_states.device min_val = torch.finfo(dtype).min k_states_T = pooled_k_bched.permute(0, 2, 3, 1) current_scaling = 1.0 if self.decouple_router and "INFONCE" in self.aux_loss_method else self.scaling scores = torch.matmul(query_states, k_states_T) * current_scaling q_mask_expanded = routing_query_mask.view(bsz, 1, q_len, 1) k_mask_expanded = chunk_mask.view(bsz, 1, 1, max_chunks) final_mask = q_mask_expanded & k_mask_expanded scores.masked_fill_(~final_mask, min_val) if self.head_reduce_method == "max": scores = scores.max(dim=1).values elif self.head_reduce_method == "mean": scores = scores.mean(dim=1) else: raise NotImplementedError(f"Unsupported head reduce method: {self.head_reduce_method}") if self.query_reduce_method == "max": scores_final = scores.max(dim=1).values elif self.query_reduce_method == "mean": valid_mask = final_mask.squeeze(1) # [B, Q_len, C] scores_clean = torch.where(valid_mask, scores, torch.zeros_like(scores)) sum_scores = scores_clean.sum(dim=1) # [B, C] counts = valid_mask.sum(dim=1).to(dtype).clamp(min=1.0) mean_scores = sum_scores / counts scores_final = torch.where( chunk_mask, mean_scores, torch.tensor(min_val, device=device, dtype=dtype) ) elif self.query_reduce_method == "last": q_lens = routing_query_mask.sum(dim=1).long() last_indices = (q_lens - 1).clamp(min=0) gather_idx = last_indices.view(bsz, 1, 1).expand(-1, 1, max_chunks) scores_final = scores.gather(1, gather_idx).squeeze(1) scores_final.masked_fill_(~chunk_mask, min_val) else: raise NotImplementedError(f"Unsupported query reduce method: {self.query_reduce_method}") return scores_final def sequence_pooling_kv(self, key_states, value_states, doc_indices, global_chunk_ids): k_docs = key_states[doc_indices[:, 0], :, doc_indices[:, 1]] v_docs = value_states[doc_indices[:, 0], :, doc_indices[:, 1]] unique_global_chunk_ids, chunk_lengths = torch.unique_consecutive(global_chunk_ids, return_counts=True) num_unique_chunks = unique_global_chunk_ids.shape[0] chunk_counts_view = chunk_lengths.view(-1, 1, 1).to(dtype=torch.float32) def compute_pooled_states_via_cumsum(states, counts_view, lengths): b, h, d = states.shape states_flat = states.reshape(b, -1).to(dtype=torch.float32) states_cumsum = F.pad(torch.cumsum(states_flat, dim=0), (0, 0, 1, 0)) chunk_cu_seqlens = F.pad(torch.cumsum(lengths, 0), (1, 0)) state_sums_flat = states_cumsum[chunk_cu_seqlens[1:]] - states_cumsum[chunk_cu_seqlens[:-1]] state_sums = state_sums_flat.view(num_unique_chunks, h, d) return (state_sums / counts_view).to(dtype=states.dtype) pooled_k_chunks = compute_pooled_states_via_cumsum(k_docs, chunk_counts_view, chunk_lengths) pooled_v_chunks = compute_pooled_states_via_cumsum(v_docs, chunk_counts_view, chunk_lengths) return pooled_k_chunks, pooled_v_chunks def sequence_pooling_qkv(self, query_states, key_states, value_states, doc_indices, global_chunk_ids): q_docs = query_states[doc_indices[:, 0], :, doc_indices[:, 1]] k_docs = key_states[doc_indices[:, 0], :, doc_indices[:, 1]] v_docs = value_states[doc_indices[:, 0], :, doc_indices[:, 1]] unique_global_chunk_ids, chunk_lengths = torch.unique_consecutive(global_chunk_ids, return_counts=True) num_unique_chunks = unique_global_chunk_ids.shape[0] chunk_counts_view = chunk_lengths.view(-1, 1, 1).to(dtype=torch.float32) def compute_pooled_states_via_cumsum(states, counts_view, lengths): b, h, d = states.shape states_flat = states.reshape(b, -1).to(dtype=torch.float32) states_cumsum = F.pad(torch.cumsum(states_flat, dim=0), (0, 0, 1, 0)) chunk_cu_seqlens = F.pad(torch.cumsum(lengths, 0), (1, 0)) state_sums_flat = states_cumsum[chunk_cu_seqlens[1:]] - states_cumsum[chunk_cu_seqlens[:-1]] state_sums = state_sums_flat.view(num_unique_chunks, h, d) return (state_sums / counts_view).to(dtype=states.dtype) pooled_q_chunks = compute_pooled_states_via_cumsum(q_docs, chunk_counts_view, chunk_lengths) pooled_k_chunks = compute_pooled_states_via_cumsum(k_docs, chunk_counts_view, chunk_lengths) pooled_v_chunks = compute_pooled_states_via_cumsum(v_docs, chunk_counts_view, chunk_lengths) return pooled_q_chunks, pooled_k_chunks, pooled_v_chunks def count_chunks_per_batch(self, doc_ids, attention_mask, kernel_size): batch_size = doc_ids.size(0) chunk_counts = [] for i in range(batch_size): mask = attention_mask[i] ids = doc_ids[i] valid_ids = ids[mask == 1] if len(valid_ids) == 0: chunk_counts.append(0) continue _, counts = torch.unique_consecutive(valid_ids, return_counts=True) num_chunks = (counts + kernel_size - 1) // kernel_size total_chunks = num_chunks.sum().item() chunk_counts.append(total_chunks) return torch.LongTensor(chunk_counts).to(doc_ids.device) def _forward( self, hidden_states: torch.Tensor, doc_ids: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: bsz, q_len, _ = hidden_states.shape device, dtype = hidden_states.device, hidden_states.dtype hidden_shape = (bsz, q_len, -1, self.head_dim) query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) routing_query_mask = (doc_ids == 0) & (attention_mask == 1) doc_token_mask = (doc_ids > 0) & (attention_mask == 1) query_indices = torch.nonzero(routing_query_mask, as_tuple=False) doc_indices = torch.nonzero(doc_token_mask, as_tuple=False) if doc_indices.shape[0] == 0 or query_indices.shape[0] == 0: raise ValueError("No query or doc tokens found") max_doc_id = int(doc_ids.max().item()) attn_output = torch.zeros((bsz, q_len, self.config.num_attention_heads * self.head_dim), device=device, dtype=dtype) if self.is_router_layer: original_doc_ids = doc_ids[doc_token_mask] original_doc_batch_indices = doc_indices[:, 0] global_doc_ids = original_doc_batch_indices * (max_doc_id + 1) + original_doc_ids _, counts = torch.unique_consecutive(global_doc_ids, return_counts=True) total_doc_tokens = global_doc_ids.shape[0] offsets = torch.zeros(counts.shape[0] + 1, dtype=counts.dtype, device=device) offsets[1:] = counts.cumsum(0) offsets = offsets[:-1] expanded_offsets = torch.repeat_interleave(offsets, counts) original_order_ranks = torch.arange(total_doc_tokens, device=device) - expanded_offsets chunk_indices = original_order_ranks // self.pooling_kernel_size max_chunks_per_doc = (q_len // self.pooling_kernel_size) + 1 global_chunk_ids = global_doc_ids * max_chunks_per_doc + chunk_indices routing_q_states = None routing_pooled_k_chunks = None if self.decouple_router: routing_q_states = self.router_q_proj(hidden_states).view(hidden_shape).transpose(1, 2) if "INFONCE" in self.aux_loss_method: routing_q_states = F.normalize(routing_q_states, p=2, dim=-1) r_k_raw = self.router_k_proj(hidden_states).view(hidden_shape).transpose(1, 2) r_k_raw = repeat_kv(r_k_raw, self.num_key_value_groups) r_k_docs = r_k_raw[doc_indices[:, 0], :, doc_indices[:, 1]] unique_global_chunk_ids = torch.unique_consecutive(global_chunk_ids) _, chunk_lengths = torch.unique_consecutive(global_chunk_ids, return_counts=True) chunk_counts_view = chunk_lengths.view(-1, 1, 1).to(dtype=torch.float32) b_k, h_k, d_k = r_k_docs.shape k_flat = r_k_docs.reshape(b_k, -1).to(dtype=torch.float32) k_cumsum = F.pad(torch.cumsum(k_flat, dim=0), (0, 0, 1, 0)) chunk_cu_seqlens = F.pad(torch.cumsum(chunk_lengths, 0), (1, 0)) k_sums_flat = k_cumsum[chunk_cu_seqlens[1:]] - k_cumsum[chunk_cu_seqlens[:-1]] routing_pooled_k_chunks = (k_sums_flat.view(unique_global_chunk_ids.shape[0], h_k, d_k) / chunk_counts_view).to(dtype=r_k_docs.dtype) if "INFONCE" in self.aux_loss_method: routing_pooled_k_chunks = F.normalize(routing_pooled_k_chunks, p=2, dim=-1) pooled_q_chunks = query_states[doc_indices[:, 0], :, doc_indices[:, 1]] pooled_k_chunks = key_states[doc_indices[:, 0], :, doc_indices[:, 1]] pooled_v_chunks = value_states[doc_indices[:, 0], :, doc_indices[:, 1]] num_doc_tokens = pooled_q_chunks.shape[0] num_chunks = num_doc_tokens // self.pooling_kernel_size pooled_q_chunks = pooled_q_chunks.view(num_chunks, self.pooling_kernel_size, self.num_heads, self.head_dim).mean(dim=1) pooled_k_chunks = pooled_k_chunks.view(num_chunks, self.pooling_kernel_size, self.num_heads, self.head_dim).mean(dim=1) pooled_v_chunks = pooled_v_chunks.view(num_chunks, self.pooling_kernel_size, self.num_heads, self.head_dim).mean(dim=1) num_heads = self.config.num_attention_heads head_dim = self.head_dim else: pooled_q_chunks, pooled_k_chunks, pooled_v_chunks = self.sequence_pooling_qkv( query_states, key_states, value_states, doc_indices, global_chunk_ids, ) num_heads = self.config.num_attention_heads head_dim = self.head_dim routing_q_states = query_states routing_pooled_k_chunks = pooled_k_chunks if "INFONCE" in self.aux_loss_method: routing_q_states = F.normalize(routing_q_states, p=2, dim=-1) routing_pooled_k_chunks = F.normalize(routing_pooled_k_chunks, p=2, dim=-1) unique_global_chunk_ids = torch.unique_consecutive(global_chunk_ids) num_unique_chunks = unique_global_chunk_ids.shape[0] chunks_per_sample = self.count_chunks_per_batch(doc_ids, doc_token_mask, kernel_size=self.pooling_kernel_size) max_chunks = chunks_per_sample.max().item() pooled_router_k_bched = torch.zeros((bsz, max_chunks, num_heads, self.head_dim), device=device, dtype=dtype) chunk_mask = torch.arange(max_chunks, device=device).unsqueeze(0) < chunks_per_sample.unsqueeze(1) pooled_router_k_bched[chunk_mask] = routing_pooled_k_chunks q_lens = routing_query_mask.sum(dim=1) # (B,) max_q_len = int(q_lens.max().item()) if max_q_len == 0: max_q_len = 1 valid_q_flat = routing_q_states.transpose(1, 2)[routing_query_mask] # [Total_Valid_Q, H, D] compact_q_states_t = torch.zeros( bsz, max_q_len, self.config.num_attention_heads, self.head_dim, device=device, dtype=dtype ) idx_range = torch.arange(max_q_len, device=device).unsqueeze(0) mask_compact_q = idx_range < q_lens.unsqueeze(1) compact_q_states_t[mask_compact_q] = valid_q_flat compact_q_states = compact_q_states_t.transpose(1, 2) max_scores_per_chunk = self._calculate_routing_scores_adaptive( compact_q_states, # (B, H, S, D) pooled_router_k_bched, # (B, C, H, D) mask_compact_q, # (B, S) chunk_mask # (B, C) ) pooled_global_doc_ids = unique_global_chunk_ids // max_chunks_per_doc pooled_doc_ids_in_sample = pooled_global_doc_ids % (max_doc_id + 1) chunk_to_doc_id_flat = pooled_doc_ids_in_sample # 形状: (total_chunks, ) chunk_to_doc_id_bched = torch.full((bsz, max_chunks), 0, dtype=torch.long, device=device) chunk_to_doc_id_bched[chunk_mask] = chunk_to_doc_id_flat offsets = torch.arange(bsz, device=device) * (max_doc_id + 1) global_chunk_to_doc_id = chunk_to_doc_id_bched + offsets.unsqueeze(1) flat_doc_scores = torch.full((bsz * (max_doc_id + 1),), -float('inf'), device=device, dtype=dtype) valid_scores_flat = max_scores_per_chunk[chunk_mask] valid_global_doc_ids_flat = global_chunk_to_doc_id[chunk_mask] if self.chunk_reduce_method == "max": doc_scores = flat_doc_scores.scatter_reduce( dim=0, index=valid_global_doc_ids_flat, src=valid_scores_flat, reduce="amax", include_self=True ) elif self.chunk_reduce_method == "mean": flat_doc_sums = torch.zeros_like(flat_doc_scores) flat_doc_sums = flat_doc_sums.scatter_reduce( dim=0, index=valid_global_doc_ids_flat, src=valid_scores_flat, reduce="sum", include_self=False ) flat_doc_counts = torch.zeros_like(flat_doc_scores) ones = torch.ones_like(valid_scores_flat) flat_doc_counts = flat_doc_counts.scatter_reduce( dim=0, index=valid_global_doc_ids_flat, src=ones, reduce="sum", include_self=False ) flat_doc_counts_safe = flat_doc_counts.clamp(min=1.0) mean_scores = flat_doc_sums / flat_doc_counts_safe doc_scores = torch.where( flat_doc_counts > 0, mean_scores, flat_doc_scores # 这里是 -inf ) else: raise ValueError(f"Invalid chunk reduction method: {self.chunk_reduce_method}") scores_by_batch = doc_scores.view(bsz, -1) return_scores_by_batch = scores_by_batch.clone() num_docs_per_sample = (scores_by_batch > -1e9).sum(dim=1) # 为每个样本计算k值：取配置的top_k和实际文档数的较小者 k_per_sample = torch.min(num_docs_per_sample, torch.full_like(num_docs_per_sample, self.top_k_docs)) _, sorted_indices = torch.sort(scores_by_batch, dim=1, descending=True) range_tensor = torch.arange(scores_by_batch.shape[1], device=device).expand(bsz, -1) selection_mask = range_tensor < k_per_sample.unsqueeze(1) selected_docs_indices = sorted_indices.masked_fill(~selection_mask, -50) prompt_and_response_mask = (doc_ids < 1) & (attention_mask == 1) # 此处的 selected_docs_indices 已经是修复后的张量，所以这行代码无需修改 selected_docs_mask = torch.any(doc_ids.unsqueeze(-1) == selected_docs_indices.unsqueeze(1), dim=-1) & doc_token_mask pa_indices = torch.nonzero(prompt_and_response_mask, as_tuple=False) q_pa_flat = query_states[pa_indices[:, 0], :, pa_indices[:, 1]] k_pa_flat = key_states[pa_indices[:, 0], :, pa_indices[:, 1]] v_pa_flat = value_states[pa_indices[:, 0], :, pa_indices[:, 1]] sort_key_pa = pa_indices[:, 0] * q_len + pa_indices[:, 1] selected_doc_token_indices = torch.nonzero(selected_docs_mask, as_tuple=False) is_doc_token_mask_flat = doc_token_mask.flatten() global_chunk_ids_padded = torch.full((bsz * q_len,), -1, dtype=torch.long, device=device) global_chunk_ids_padded[is_doc_token_mask_flat] = global_chunk_ids selected_chunk_ids_flat = global_chunk_ids_padded.view(bsz, q_len)[selected_docs_mask] unique_selected_chunk_ids, inverse_indices_fix = torch.unique(selected_chunk_ids_flat, sorted=True, return_inverse=True) if unique_selected_chunk_ids.numel() > 0: first_occurrence_indices = torch.empty_like(unique_selected_chunk_ids, dtype=torch.long) first_occurrence_indices.scatter_reduce_(src=torch.arange(selected_chunk_ids_flat.numel(), device=device),index=inverse_indices_fix, dim=0, reduce='amin', include_self=False) representative_indices = selected_doc_token_indices[first_occurrence_indices] sort_key_chunks = representative_indices[:, 0] * q_len + representative_indices[:, 1] map_gcid_to_poolidx = torch.full((int(global_chunk_ids.max().item()) + 1,), -1, dtype=torch.long, device=device) map_gcid_to_poolidx[unique_global_chunk_ids] = torch.arange(num_unique_chunks, device=device) pool_indices_to_gather = map_gcid_to_poolidx[unique_selected_chunk_ids] assert (pool_indices_to_gather.sort().values != pool_indices_to_gather).sum() == 0 q_pooled_sel_flat = pooled_q_chunks[pool_indices_to_gather] k_pooled_sel_flat = pooled_k_chunks[pool_indices_to_gather] v_pooled_sel_flat = pooled_v_chunks[pool_indices_to_gather] batch_indices_chunks = representative_indices[:, 0] else: sort_key_chunks = torch.tensor([], dtype=torch.long, device=device) q_pooled_sel_flat = torch.tensor([], dtype=dtype, device=device).view(0, num_heads, head_dim) k_pooled_sel_flat = torch.tensor([], dtype=dtype, device=device).view(0, num_heads, head_dim) v_pooled_sel_flat = torch.tensor([], dtype=dtype, device=device).view(0, num_heads, head_dim) batch_indices_chunks = torch.tensor([], dtype=torch.long, device=device) q_combined = torch.cat([q_pa_flat, q_pooled_sel_flat], dim=0) k_combined = torch.cat([k_pa_flat, k_pooled_sel_flat], dim=0) v_combined = torch.cat([v_pa_flat, v_pooled_sel_flat], dim=0) combined_sort_keys = torch.cat([sort_key_pa, sort_key_chunks], dim=0) _, final_sort_indices = torch.sort(combined_sort_keys) q_a_final = q_combined[final_sort_indices] k_a_final = k_combined[final_sort_indices] v_a_final = v_combined[final_sort_indices] # 4.4 计算cu_seqlens (逻辑不变) batch_indices_pa = pa_indices[:, 0] batch_indices_combined = torch.cat([batch_indices_pa, batch_indices_chunks], dim=0) sorted_batch_indices = batch_indices_combined[final_sort_indices] batch_counts_a = torch.bincount(sorted_batch_indices, minlength=bsz) cu_seqlens_a = F.pad(torch.cumsum(batch_counts_a, dim=0, dtype=torch.int32), (1, 0)) else: prompt_and_response_mask = (doc_ids < 1) & (attention_mask == 1) pa_indices = torch.nonzero(prompt_and_response_mask, as_tuple=False) q_a_final = query_states[pa_indices[:, 0], :, pa_indices[:, 1]] k_a_final = key_states[pa_indices[:, 0], :, pa_indices[:, 1]] v_a_final = value_states[pa_indices[:, 0], :, pa_indices[:, 1]] batch_counts_a = prompt_and_response_mask.sum(dim=1) cu_seqlens_a = F.pad(torch.cumsum(batch_counts_a, dim=0, dtype=torch.int32), (1, 0)) return_scores_by_batch = None if q_a_final.shape[0] > 0: output_a_final = flash_attn_varlen_func( q_a_final, k_a_final, v_a_final, cu_seqlens_q=cu_seqlens_a, cu_seqlens_k=cu_seqlens_a, max_seqlen_q=int(batch_counts_a.max()), max_seqlen_k=int(batch_counts_a.max()), dropout_p=self.attention_dropout if self.training else 0.0, causal=True ).view(-1, self.config.num_attention_heads * self.head_dim) if self.is_router_layer: is_pa_mask_combined = torch.cat([ torch.ones(pa_indices.shape[0], dtype=torch.bool, device=device), torch.zeros(q_pooled_sel_flat.shape[0], dtype=torch.bool, device=device) # 修正为使用池化块的数量 ], dim=0) is_pa_mask_sorted = is_pa_mask_combined[final_sort_indices] output_pa_part = output_a_final[is_pa_mask_sorted] attn_output[pa_indices[:, 0], pa_indices[:, 1]] = output_pa_part else: attn_output[pa_indices[:, 0], pa_indices[:, 1]] = output_a_final indices_b = torch.nonzero(doc_token_mask, as_tuple=False) if indices_b.shape[0] > 0: q_b, k_b, v_b = query_states[indices_b[:, 0], :, indices_b[:, 1]], key_states[indices_b[:, 0], :, indices_b[:, 1]], value_states[indices_b[:, 0], :, indices_b[:, 1]] doc_ids_b = doc_ids[indices_b[:, 0], indices_b[:, 1]] batch_indices_b = indices_b[:, 0] global_doc_ids_b = batch_indices_b * (max_doc_id + 1) + doc_ids_b _, counts_b = torch.unique_consecutive(global_doc_ids_b, return_counts=True) cu_seqlens_b = F.pad(torch.cumsum(counts_b, dim=0, dtype=torch.int32), (1, 0)) output_b_flat = flash_attn_varlen_func( q_b, k_b, v_b, cu_seqlens_q=cu_seqlens_b, cu_seqlens_k=cu_seqlens_b, max_seqlen_q=int(counts_b.max()), max_seqlen_k=int(counts_b.max()), dropout_p=self.attention_dropout if self.training else 0.0, causal=True ).view(-1, self.config.num_attention_heads * self.head_dim) attn_output[indices_b[:, 0], indices_b[:, 1]] += output_b_flat return (self.o_proj(attn_output), return_scores_by_batch), None