小红书dots3 note开源深度解析：280B MoE多模态模型，内容平台入局开源大模型赛道

import torch
import torch.nn.functional as F

class Dots3MoELayer(torch.nn.Module):
    """dots3-note MoE层简化实现"""
    
    def __init__(self, hidden_dim=5120, n_experts=256, n_shared_experts=1, 
                 top_k=8, expert_dim=1536):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.n_experts = n_experts
        self.top_k = top_k
        
        # 路由网络（Gate）
        self.gate = torch.nn.Linear(hidden_dim, n_experts, bias=False)
        
        # 共享专家
        self.shared_expert = torch.nn.Sequential(
            torch.nn.Linear(hidden_dim, expert_dim * 4),
            torch.nn.SiLU(),
            torch.nn.Linear(expert_dim * 4, hidden_dim),
        )
        
        # 路由专家（示意：实际使用分组矩阵乘优化）
        self.experts = torch.nn.ModuleList([
            torch.nn.Sequential(
                torch.nn.Linear(hidden_dim, expert_dim * 4),
                torch.nn.SiLU(),
                torch.nn.Linear(expert_dim * 4, hidden_dim),
            )
            for _ in range(n_experts)
        ])
    
    def forward(self, x):
        # x: [batch_size, seq_len, hidden_dim]
        batch_size, seq_len, _ = x.shape
        
        # 1. 计算路由分数
        gate_logits = self.gate(x)  # [B, S, 256]
        gate_scores = F.softmax(gate_logits, dim=-1, dtype=torch.float32)
        
        # 2. Top-8 选择
        topk_scores, topk_indices = torch.topk(gate_scores, self.top_k, dim=-1)
        topk_scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True)
        
        # 3. 共享专家输出（始终激活）
        shared_out = self.shared_expert(x)
        
        # 4. 路由专家输出（仅对选中的专家计算）
        # 简化：实际实现使用分组GEMM + 分散/聚合
        final_hidden = shared_out  # 从共享专家开始
        final_hidden = final_hidden.to(x.dtype)
        
        # 实际推理框架会使用dispatch/combine融合算子
        # 这里仅展示路由逻辑
        routed_output = torch.zeros_like(x)
        for i in range(self.n_experts):
            mask = (topk_indices == i).any(dim=-1)
            if mask.any():
                expert_out = self.experts[i](x[mask])
                # 加权求和
                expert_weight = topk_scores[mask][topk_indices[mask] == i]
                routed_output[mask] += expert_out * expert_weight.unsqueeze(-1)
        
        return final_hidden + routed_output

import torch
import torch.nn.functional as F

class DSAttention(torch.nn.Module):
    """
    Dots Sparse Attention (DSA) 简化实现
    
    核心流程：
    1. 使用轻量级 Indexer 计算所有 token 的索引分数
    2. 选择 Top-2048 个最相关的 KV 条目
    3. 仅对选中的条目进行标准 MLA 注意力计算
    """
    
    def __init__(self, hidden_dim=5120, n_heads=64, head_dim=128,
                 index_n_heads=64, index_head_dim=128, index_topk=2048):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.index_topk = index_topk
        
        # MLA: QKV 低秩投影
        self.q_proj = torch.nn.Linear(hidden_dim, n_heads * head_dim, bias=False)
        self.kv_proj = torch.nn.Linear(hidden_dim, 2 * head_dim, bias=False)  # MLA共享KV
        self.o_proj = torch.nn.Linear(n_heads * head_dim, hidden_dim, bias=False)
        
        # DSA Indexer: 轻量级检索器
        self.indexer = DSAIndexer(
            hidden_dim=hidden_dim,
            n_heads=index_n_heads,
            head_dim=index_head_dim
        )
    
    def forward(self, x, attention_mask=None):
        batch_size, seq_len, _ = x.shape
        
        # 1. 计算 Q 和共享 KV
        q = self.q_proj(x)  # [B, S, n_heads * head_dim]
        q = q.view(batch_size, seq_len, self.n_heads, self.head_dim)
        
        kv = self.kv_proj(x)  # [B, S, 2 * head_dim]
        k, v = kv.chunk(2, dim=-1)  # 共享单头 KV
        k = k.unsqueeze(2)  # [B, S, 1, head_dim]
        v = v.unsqueeze(2)  # [B, S, 1, head_dim]
        
        # 2. 使用 Indexer 选择 Top-K 位置
        #    返回每个 query 需要关注的 key 索引
        selected_indices = self.indexer(x, self.index_topk)
        # selected_indices: [B, S, topk]
        
        # 3. 仅对选中的 KV 计算注意力
        # 使用 gather 操作收集选中的 K, V
        selected_k = torch.gather(
            k.unsqueeze(1).expand(-1, seq_len, -1, -1, -1),
            dim=2,
            index=selected_indices.unsqueeze(-1).unsqueeze(-1)
            .expand(-1, -1, -1, 1, self.head_dim)
        )  # [B, S, topk, 1, head_dim]
        
        selected_v = torch.gather(
            v.unsqueeze(1).expand(-1, seq_len, -1, -1, -1),
            dim=2,
            index=selected_indices.unsqueeze(-1).unsqueeze(-1)
            .expand(-1, -1, -1, 1, self.head_dim)
        )  # [B, S, topk, 1, head_dim]
        
        # 4. 注意力计算（仅 topk 个位置）
        attn_scores = torch.einsum(
            "bqhd,bqkhd->bqhk",
            q.unsqueeze(2),  # [B, S, 1, n_heads, head_dim]
            selected_k      # [B, S, topk, 1, head_dim]
        ) / (self.head_dim ** 0.5)
        
        if attention_mask is not None:
            attn_scores = attn_scores + attention_mask
        
        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_output = torch.einsum(
            "bqhk,bqkhd->bqhd",
            attn_weights,
            selected_v
        )
        
        # 5. 输出投影
        attn_output = attn_output.reshape(batch_size, seq_len, -1)
        return self.o_proj(attn_output)


class DSAIndexer(torch.nn.Module):
    """
    DSA 索引器 (Indexer)
    
    使用独立的低维投影快速编码所有 token 的位置与语义，
    为每个 query 计算与所有 key 的索引分数。
    """
    
    def __init__(self, hidden_dim, n_heads=64, head_dim=128):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = head_dim
        
        # 独立的索引器参数
        self.wq = torch.nn.Linear(hidden_dim, n_heads * head_dim, bias=False)
        self.wk = torch.nn.Linear(hidden_dim, n_heads * head_dim, bias=False)
        self.weights_proj = torch.nn.Linear(hidden_dim, n_heads, bias=False)
        self.k_norm = torch.nn.LayerNorm(head_dim)
    
    def forward(self, x, topk):
        batch_size, seq_len, _ = x.shape
        
        # 索引器 Q 和 K
        idx_q = self.wq(x)  # [B, S, n_heads * head_dim]
        idx_k = self.wk(x)  # [B, S, n_heads * head_dim]
        
        idx_q = idx_q.view(batch_size, seq_len, self.n_heads, self.head_dim)
        idx_k = idx_k.view(batch_size, seq_len, self.n_heads, self.head_dim)
        idx_k = self.k_norm(idx_k)
        
        # 计算索引分数
        # 公式: I(t,s) = sum_j(w_j * ReLU(q_j · k_j))
        idx_scores = torch.einsum(
            "bqhd,bkhd->bqkh",
            idx_q, idx_k
        )  # [B, S, S, n_heads]
        
        idx_scores = F.relu(idx_scores)  # ReLU 激活
        
        # 多头加权求和
        idx_weights = self.weights_proj(x)  # [B, S, n_heads]
        idx_scores = torch.einsum(
            "bqkh,bqh->bqk",
            idx_scores, idx_weights
        )  # [B, S, S]
        
        # 选择 Top-K
        topk_values, topk_indices = torch.topk(
            idx_scores, topk, dim=-1
        )  # [B, S, topk]
        
        return topk_indices

import torch
import torchvision.transforms as T
from PIL import Image

class Dots3MultimodalProcessor:
    """
    dots3-note 多模态处理管线
    支持文本、图像、视频、音频的统一输入处理
    """
    
    def __init__(self, model, processor):
        self.model = model
        self.processor = processor
        self.device = next(model.parameters()).device
        
        # 图像预处理
        self.image_transform = T.Compose([
            T.Resize((384, 384)),
            T.ToTensor(),
            T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
        ])
    
    def process_text_only(self, text, enable_thinking=False):
        """纯文本推理"""
        messages = [{"role": "user", "content": text}]
        inputs = self.processor.tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt",
            return_dict=True,
            enable_thinking=enable_thinking,
        ).to(self.device)
        
        outputs = self.model.generate(
            **inputs,
            max_new_tokens=2048,
            temperature=0.7,
            top_p=0.95,
        )
        return self.processor.decode(
            outputs[0, inputs.input_ids.shape[1]:],
            skip_special_tokens=True
        )
    
    def process_image_text(self, image_path, text):
        """图文理解"""
        image = Image.open(image_path).convert("RGB")
        image_tensor = self.image_transform(image).unsqueeze(0).to(self.device)
        
        # 构建多模态输入
        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "image": image_tensor},
                {"type": "text", "text": text},
            ]
        }]
        
        inputs = self.processor.tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt",
            return_dict=True,
        ).to(self.device)
        
        outputs = self.model.generate(
            **inputs,
            max_new_tokens=2048,
            temperature=0.7,
        )
        return self.processor.decode(
            outputs[0, inputs.input_ids.shape[1]:],
            skip_special_tokens=True
        )
    
    def process_video(self, video_path, text):
        """视频理解（含音频轨）"""
        # 视频处理器会自动执行：
        # 1. 帧采样（根据token预算确定帧数）
        # 2. 音频轨提取并编码
        # 3. 时间戳交错排列
        # 4. 帧、时间戳、音频段的token预算分配
        messages = [{
            "role": "user",
            "content": [
                {"type": "video_url", "video_url": {"url": video_path}},
                {"type": "text", "text": text},
            ]
        }]
        # ... 后续处理与图像类似
        pass

# 使用示例
def run_multimodal_example():
    """使用OpenAI兼容接口进行多模态推理"""
    from openai import OpenAI
    
    client = OpenAI(
        base_url="http://127.0.0.1:8000/v1",
        api_key="EMPTY"
    )
    
    # 图像理解
    response = client.chat.completions.create(
        model="dots3-note-prev",
        messages=[{
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {
                    "url": "https://example.com/photo.jpg"
                }},
                {"type": "text", "text": "请描述这张图片中的细节"},
            ]
        }],
        temperature=0.7,
        max_tokens=1024,
    )
    print(response.choices[0].message.content)
    
    # 启用思考模式（推理模式）
    response = client.chat.completions.create(
        model="dots3-note-prev",
        messages=[{"role": "user", "content": "证明: 根号2是无理数"}],
        temperature=1.0,
        max_tokens=4096,
        extra_body={"chat_template_kwargs": {"enable_thinking": True}},
    )
    print(response.choices[0].message.content)

import torch
from typing import List, Dict, Any, Optional

class TEMPOTrainer:
    """
    TEMPO: Test-time scaled Value Estimation with Macro-step Policy Optimization
    
    核心思想：将长程任务拆分为多个macro-step，
    在每个macro-step结束时，Agent切换为critic角色，
    通过test-time scaling推理来估计当前状态的预期剩余回报。
    """
    
    def __init__(self, model, tokenizer, 
                 macro_step_size: int = 10,
                 n_critic_rollouts: int = 8):
        self.model = model
        self.tokenizer = tokenizer
        self.macro_step_size = macro_step_size
        self.n_critic_rollouts = n_critic_rollouts
    
    def train_epoch(self, environment, n_episodes: int):
        """一轮TEMPO训练"""
        for episode in range(n_episodes):
            trajectory = self.rollout_with_macro_steps(environment)
            self.update_policy(trajectory)
    
    def rollout_with_macro_steps(self, env) -> List[Dict]:
        """执行带有macro-step划分的 rollout"""
        trajectory = []
        state = env.reset()
        done = False
        step_count = 0
        
        while not done:
            # Actor阶段：执行一个macro-step
            macro_actions = []
            for _ in range(self.macro_step_size):
                if done:
                    break
                action = self.act(state)
                next_state, reward, done, info = env.step(action)
                macro_actions.append({
                    "state": state,
                    "action": action,
                    "reward": reward,
                    "next_state": next_state,
                    "done": done,
                })
                state = next_state
                step_count += 1
            
            # Critic阶段：评估当前状态的预期剩余回报
            if not done:
                value_estimate = self.estimate_value(
                    state, 
                    trajectory, 
                    env.get_privileged_info()
                )
            else:
                value_estimate = 0.0
            
            # 记录macro-step
            trajectory.append({
                "macro_actions": macro_actions,
                "value_estimate": value_estimate,
                "final_reward": reward if done else None,
            })
        
        return trajectory
    
    def act(self, state) -> str:
        """Actor: 生成动作"""
        prompt = self.build_actor_prompt(state)
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
        outputs = self.model.generate(
            **inputs, 
            max_new_tokens=256,
            temperature=1.0,
            top_p=0.95,
        )
        action = self.tokenizer.decode(
            outputs[0, inputs.input_ids.shape[1]:],
            skip_special_tokens=True
        )
        return action
    
    def estimate_value(self, state, history, privileged_info) -> float:
        """
        Critic: 通过test-time scaling推理估计价值
        
        通过多次rollout和推理分析，估计当前状态的预期剩余回报。
        关键发现：在长程任务中，"评价比生成更简单"，
        即使Agent当前无法解决问题，作为critic时仍能做出准确的价值估计。
        """
        prompt = self.build_critic_prompt(state, history, privileged_info)
        
        value_estimates = []
        for _ in range(self.n_critic_rollouts):
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=512,
                temperature=0.7,
                # 启用思考模式进行深度分析
                chat_template_kwargs={"enable_thinking": True},
            )
            analysis = self.tokenizer.decode(
                outputs[0, inputs.input_ids.shape[1]:],
                skip_special_tokens=True
            )
            # 从分析文本中提取价值估计
            value = self.extract_value_from_analysis(analysis)
            value_estimates.append(value)
        
        # 聚合多个critic rollouts的估计
        return torch.mean(torch.tensor(value_estimates)).item()
    
    def update_policy(self, trajectory):
        """基于TEMPO的trajectory更新策略"""
        # 对每个macro-step计算优势函数
        advantages = []
        returns = []
        G = 0
        
        for macro_step in reversed(trajectory):
            if macro_step["final_reward"] is not None:
                G = macro_step["final_reward"]
            else:
                G = macro_step["value_estimate"] + 0.99 * G  # 折扣因子
            returns.insert(0, G)
            advantages.insert(0, G - macro_step["value_estimate"])
        
        # 使用PPO-like的目标函数更新策略
        # 实际实现中，会基于advantages进行策略梯度更新
        loss = self.compute_policy_loss(trajectory, advantages)
        loss.backward()
        # ... 优化器step

import torch
import torch.nn.functional as F

class MTPDecoding:
    """
    MTP (Multi-Token Prediction) 多token投机解码
    
    内置的MTP层可以在单次前向中并行生成多个候选token，
    通过校验机制接受正确预测，大幅减少解码前向次数。
    """
    
    def __init__(self, model, mtp_head, n_speculative_tokens=3):
        self.model = model
        self.mtp_head = mtp_head  # 1.13B共享MTP层
        self.n_speculative_tokens = n_speculative_tokens
    
    @torch.no_grad()
    def generate_with_mtp(self, input_ids, max_new_tokens=1024, 
                          temperature=0.7, top_p=0.95):
        """
        使用MTP投机解码加速生成
        
        流程:
        1. 主模型前向 → 获取当前token的logits
        2. MTP Head → 并行预测后续n个token
        3. 校验 → 接受正确预测的token，丢弃错误的
        4. 重复直到完成
        """
        device = input_ids.device
        batch_size = input_ids.shape[0]
        
        for _ in range(max_new_tokens // (self.n_speculative_tokens + 1) + 1):
            # Step 1: 主模型前向 → 获取当前token logits
            main_outputs = self.model(input_ids, use_cache=True)
            main_logits = main_outputs.logits[:, -1, :]  # [B, vocab_size]
            
            # 采样当前token
            next_token = self.sample(main_logits, temperature, top_p)
            
            # Step 2: MTP Head → 并行预测后续token
            # MTP基于当前层的隐藏状态，预测后续n个token
            draft_tokens = []
            hidden_state = main_outputs.last_hidden_state[:, -1:, :]
            
            for _ in range(self.n_speculative_tokens):
                mtp_logits = self.mtp_head(hidden_state)
                draft_token = self.sample(mtp_logits[:, -1, :], temperature, top_p)
                draft_tokens.append(draft_token)
                
                # 更新隐藏状态用于下一个MTP预测
                # MTP层共享主模型的embedding
                draft_embed = self.model.get_input_embeddings()(draft_token)
                hidden_state = self.mtp_head.forward_embed(draft_embed)
            
            # Step 3: 校验 — 并行前向所有草稿token
            draft_ids = torch.cat([next_token] + draft_tokens, dim=-1)
            verify_outputs = self.model(
                torch.cat([input_ids, draft_ids[:, :-1]], dim=-1),
                use_cache=False
            )
            
            verify_logits = verify_outputs.logits[:, -self.n_speculative_tokens:, :]
            
            # 校验每个草稿token
            accepted_tokens = [next_token]
            for i, draft_token in enumerate(draft_tokens):
                target_logits = main_logits if i == 0 else verify_logits[:, i-1, :]
                target_prob = F.softmax(target_logits, dim=-1)
                draft_prob = F.softmax(
                    verify_logits[:, i, :] if i < len(draft_tokens) - 1 
                    else self.mtp_head(hidden_state)[:, -1, :], 
                    dim=-1
                )
                
                # 拒绝采样校验
                p_accept = torch.min(
                    torch.ones_like(target_prob[0, draft_token[0]]),
                    draft_prob[0, draft_token[0]] / target_prob[0, draft_token[0]]
                )
                
                if torch.rand(1).item() < p_accept.item():
                    accepted_tokens.append(draft_token)
                else:
                    # 拒绝: 从target分布重新采样
                    corrected = self.sample(target_logits, temperature, top_p)
                    accepted_tokens.append(corrected)
                    break
            
            # 将接受的token追加到输入
            new_tokens = torch.cat(accepted_tokens, dim=-1)
            input_ids = torch.cat([input_ids, new_tokens], dim=-1)
            
            # 检查是否达到EOS
            if (new_tokens == self.model.config.eos_token_id).any():
                break
        
        return input_ids
    
    def sample(self, logits, temperature, top_p):
        """带temperature和top-p的采样"""
        logits = logits / temperature
        
        # Top-p (nucleus) sampling
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
        
        # 移除累积概率超过top_p的token
        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0
        
        indices_to_remove = sorted_indices_to_remove.scatter(
            -1, sorted_indices, sorted_indices_to_remove
        )
        logits[indices_to_remove] = float('-inf')
        
        probs = F.softmax(logits, dim=-1)
        return torch.multinomial(probs, num_samples=1)

# 使用vLLM Ascend部署dots3-note（Atlas 800 A3集群）
vllm serve dots-studio/dots3-note-prev-fp8 \
  --served-model-name dots3-note-prev \
  --host 0.0.0.0 \
  --tensor-parallel-size 8 \
  --enable-expert-parallel \
  --moe-backend deep_gemm \
  --max-model-len 262144 \
  --enable-auto-tool-choice \
  --tool-call-parser dots \
  --speculative-config '{"method":"mtp","num_speculative_tokens":3}'

# 1. 使用SGLang部署（推荐，支持完整的MTP投机解码）
docker run --gpus all --ipc=host -p 8000:8000 \
  lmsysorg/sglang:dev-dots3-note \
  sglang serve \
  --model-path dots-studio/dots3-note-prev-fp8 \
  --served-model-name dots3-note-prev \
  --host 0.0.0.0 \
  --port 8000 \
  --context-length 524288 \
  --enable-dp-attention \
  --dp-size 8 \
  --tp-size 8 \
  --ep-size 8 \
  --moe-dense-tp-size 1 \
  --page-size 64 \
  --trust-remote-code \
  --attention-backend fa3 \
  --moe-a2a-backend deepep \
  --enable-multimodal \
  --speculative-algorithm NEXTN \
  --speculative-num-steps 3 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 4 \
  --speculative-draft-model-path dots-studio/dots3-note-prev-fp8

# 2. 使用vLLM部署（成熟稳定）
vllm serve dots-studio/dots3-note-prev-fp8 \
  --served-model-name dots3-note-prev \
  --host 0.0.0.0 \
  --tensor-parallel-size 8 \
  --enable-expert-parallel \
  --moe-backend deep_gemm \
  --max-model-len 262144 \
  --speculative-config '{"method":"mtp","num_speculative_tokens":3}' \
  --enable-auto-tool-choice --tool-call-parser dots

from openai import OpenAI

client = OpenAI(
    base_url="http://127.0.0.1:8000/v1",
    api_key="EMPTY"
)

# 文本推理
response = client.chat.completions.create(
    model="dots3-note-prev",
    messages=[{"role": "user", "content": "解释MoE架构的工作原理"}],
    temperature=1.0,
    top_p=0.95,
    max_tokens=2048,
    extra_body={"chat_template_kwargs": {"enable_thinking": True}},
)
print(response.choices[0].message.content)

# 多模态：图像理解
response = client.chat.completions.create(
    model="dots3-note-prev",
    messages=[{
        "role": "user",
        "content": [
            {
                "type": "image_url",
                "image_url": {"url": "https://example.com/chart.png"}
            },
            {"type": "text", "text": "分析这张图表中的趋势"},
        ]
    }],
    max_tokens=1024,
)
print(response.choices[0].message.content)

# 工具调用（Agent模式）
response = client.chat.completions.create(
    model="dots3-note-prev",
    messages=[{"role": "user", "content": "查询北京今天天气，并帮我安排行程"}],
    tools=[{
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "获取某城市的天气信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {"type": "string", "description": "城市名称"}
                },
                "required": ["city"]
            }
        }
    }],
    tool_choice="auto",
)
print(response.choices[0].message)