执行摘要

• 一句话：Cosmos3 Edge 与蒸馏变体支持，变体自动识别
• 推荐动作：值得精读。重点看三处设计：① 变体检测一次性预计算（update_config_from_dict + lru_cache）避免 per-request 下载 checkpoint 配置，这是 diffusion 多变体服务的通用模式；② Ulysses SP 下 SDE 噪声生成器跨 rank 同步的思路，可复用到其它扩散模型；③ _predict_noise_cfg 对 CFG padding 在无 mask cross-attention 下语义的分析（零 K/V 仍占 softmax 权重），是少见的数值质量权衡。另外 review 中 mickqian 对防御性代码的两次拒绝（fail-fast、显式契约）展示了该仓库的工程偏好，值得团队对齐。

功能与动机

PR body 明确："All variants continue to be served by the single Cosmos3 pipeline, with the variant selected from the checkpoint config at load time." 此前 PR#24994、PR#26492、PR#27168 建立了 Cosmos3 家族基础支持，本 PR 在此基础上扩展 Edge 与 Distilled 变体。PR 评论中 nv-dmajchrowski 补充："after the previous approval we had to add few more changes (edge checkpoint format changed slightly after the release + we had small accuracy issues)"，说明 Edge 发布后 checkpoint 格式有微调，本 PR 是完整适配。

实现拆解

1. 变体检测与预计算：在 configs/pipeline_configs/cosmos3.py 新增 is_edge_checkpoint（读取 transformer/config.json 的 backbone_type 或 hidden_act=relu2）、_distilled_sampler_config（读取 scheduler config 的 _class_name 与 fixed_step_sampler_config.t_list）、get_distilled_sigmas，并使用 functools.lru_cache 缓存；Cosmos3Config.update_config_from_dict 在 model_path 就绪后一次性算出 is_edge 与 distilled_sigmas 存入配置对象，避免每个请求重复下载 checkpoint 子目录（commit 3b51f64 的优化）。

2. Edge 架构支持：configs/models/dits/cosmos3video.py 将 _build_cosmos3_param_names_mapping 参数化为 gated_mlp 模式，dense 变体跳过 gate/up 合并规则，up_proj/down_proj 直通；Cosmos3VideoArchConfig 新增 hidden_act、qk_norm_for_text、use_und_k_norm_for_gen 字段。runtime/models/dits/cosmos3video.py 新增 Cosmos3DenseMLP（up → ReLU → 平方 → down）与 _build_mlp 工厂，并按 hidden_act 选择 MLP 结构；Cosmos3CausalAttention 支持可选 QK-norm 与独立的 k_norm_und_for_gen（GEN 交叉注意力单独归一化 K）。

3. 采样默认值：configs/sample/cosmos3.py 新增 _resolve_variant_defaults，按 is_edge / is_distilled 解析；height/width 默认改为 None 由变体填充；Edge 使用 256p/480p 分辨率白名单与对应默认尺寸，T2I/T2V 分别用 guidance 7.0/5.0，T2I 默认 CFG 区间 (400, 1000)；蒸馏变体强制 guidance_scale = 1.0，num_inference_steps 固定为 sigma 表长度（期间修复了默认 None 导致 warmup 崩溃的问题）。

4. 调度与去噪：Cosmos3TimestepPreparationStage 蒸馏路径先 set_shift(1.0) 再用 set_timesteps(sigmas=distilled_sigmas) 驱动 scheduler，防止双重重偏移；_default_flow_shift_for_mode 对 Edge 视频模式统一返回 3.0。Cosmos3DenoisingStage 将批次 seed 构造的 torch.Generator（device 取 latents.device）传入 scheduler step，保证 Ulysses SP 下各 rank 噪声一致；_predict_noise_cfg 在正负 prompt 长度不等时把 CFG 两个分支拆成独立 forward，避免零 padding 的 K/V 稀释注意力。scheduler_loader.py 新增 _supported_init_kwargs，按构造函数签名过滤 scheduler_config 不接受的键。

5. 兼容与配套：registry.py 注册 Cosmos3OmniPipeline 并保留 Cosmos3OmniDiffusersPipeline 别名；hf_diffusers_utils.py 新增 _split_hf_subfolder 支持 namespace/repo/subfolder 局部下载；权重加载时丢弃 *_quantizer.* ModelOpt 校准缓存；action/api.py 增加 multipart/form-data 上传支持；video_api.py 新增 form_text_value 防止 JSON prompt 被反序列化。测试 test_cosmos3.py 新增 dense 参数映射、Edge flow_shift/采样默认值、_class_name 双拼写注册等用例。

关键文件：

• python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py（模块 管线配置；类别 source；类型 core-logic；符号 is_edge_checkpoint, _distilled_sampler_config, is_distilled_checkpoint, get_distilled_sigmas）: 变体检测核心：is_edge_checkpoint / _distilled_sampler_config / get_distilled_sigmas，并在 update_config_from_dict 中一次性预计算，per-request 路径都从这里读。
• python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py（模块 模型层；类别 source；类型 data-contract；符号 Cosmos3DenseMLP, _build_mlp, forward, init）: Edge dense 骨干的模型侧实现：新增 Cosmos3DenseMLP 与 _build_mlp 工厂，Cosmos3CausalAttention 支持可选 QK-norm 与 GEN 侧独立 K 归一化。
• python/sglang/multimodal_gen/configs/sample/cosmos3.py（模块 采样配置；类别 source；类型 core-logic；符号 _guidance_is_explicit, _resolve_variant_defaults）: 变体采样默认值（分辨率、guidance、CFG 区间、步数）在此解析，直接影响每个请求的采样行为。
• python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3.py（模块 管线阶段；类别 source；类型 data-contract；符号 _inject_caption_metadata, _default_flow_shift_for_mode, _predict_noise_cfg, _cfg_combine）: 调度与去噪主路径：distilled sigma 驱动 scheduler、SP 同步 generator、CFG 分支拆分、结构化 caption 元数据注入。
• python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py（模块 模型配置；类别 source；类型 data-contract；符号 _build_cosmos3_param_names_mapping, gated_mlp）: 权重映射契约：_build_cosmos3_param_names_mapping 参数化 gated_mlp，决定 gate/up 是否合并。
• python/sglang/multimodal_gen/test/unit/test_cosmos3.py（模块 单元测试；类别 test；类型 test-coverage；符号 TestCosmos3DenseParamNamesMapping, test_edge_flow_shift_default, test_und_mlp_up_proj_unmerged, test_gen_mlp_down_proj_unmerged）: 测试覆盖 dense 权重映射、Edge 采样默认值、_class_name 新旧注册、flow_shift 默认值。
• python/sglang/multimodal_gen/runtime/loader/component_loaders/scheduler_loader.py（模块 组件加载；类别 source；类型 core-logic；符号 _supported_init_kwargs）: _supported_init_kwargs 按构造函数签名过滤 scheduler_config 键，蒸馏 checkpoint 才能被加载。
• python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py（模块 工具层；类别 source；类型 core-logic；符号 _split_hf_subfolder）: _split_hf_subfolder 支持 namespace/repo/subfolder 形式的组件路径下载，解决新格式 checkpoint 的子目录加载。
• python/sglang/multimodal_gen/runtime/entrypoints/action/api.py（模块 入口服务；类别 source；类型 entrypoint；符号 _multipart_action_payload, _parse_form_value, _is_form_upload, _save_action_upload）: 附带新增 multipart/form-data 支持（kediwu0331 提交），让 action 端点能处理视频/图片上传。
• python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py（模块 入口服务；类别 source；类型 entrypoint；符号 form_text_value）: form_text_value 避免结构化 JSON prompt 被 form 解析反序列化，Edge 的 JSON caption 才能通过请求校验。

关键符号：is_edge_checkpoint, is_distilled_checkpoint, get_distilled_sigmas, update_config_from_dict, Cosmos3DenseMLP, _build_mlp, _build_cosmos3_param_names_mapping, _resolve_variant_defaults, _guidance_is_explicit, _inject_caption_metadata, _default_flow_shift_for_mode, _predict_noise_cfg, _supported_init_kwargs, _split_hf_subfolder, _multipart_action_payload, form_text_value

关键源码片段

python/sglang/multimodal_gen/configs/pipeline_configs/cosmos3.py

变体检测核心：is_edge_checkpoint / _distilled_sampler_config / get_distilled_sigmas，并在 update_config_from_dict 中一次性预计算，per-request 路径都从这里读。

# Cosmos3 变体检测：Edge 与蒸馏 checkpoints 的识别全部基于
# checkpoint 自身配置，而不是硬编码模型名，保证同一管线可服务整个家族。
# 函数带 lru_cache，model_path 首次解析后结果常驻内存。
​
@functools.lru_cache(maxsize=None)
def is_edge_checkpoint(model_path: str) -> bool:
    """判断 checkpoint 是否为 Edge（dense 骨干）变体。
​
    直接从 transformer/config.json 读取，早于权重加载，
    因此客户端进程在解析采样默认值时无需把权重搬上设备。
    """
    from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
        get_diffusers_component_config,
    )
​
    config = get_diffusers_component_config(
        component_path=os.path.join(model_path, "transformer")
    )
    # Edge 的 dense 骨干可能以 backbone_type 或 hidden_act=relu2 之一标记
    return (
        config.get("backbone_type") == COSMOS3_EDGE_BACKBONE_TYPE
        or config.get("hidden_act") == "relu2"
    )
​
​
@functools.lru_cache(maxsize=None)
def _distilled_sampler_config(model_path: str) -> dict | None:
    """读取蒸馏 checkpoint 的固定步 sigma 表，非蒸馏返回 None。
​
    蒸馏只体现在 scheduler：checkpoint 自带
    FlowMatchEulerDiscreteScheduler 与 fixed_step_sampler_config.t_list，
    而其余变体使用多步 FlowUniPC。
    """
    from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
        get_diffusers_component_config,
    )
​
    config = get_diffusers_component_config(
        component_path=os.path.join(model_path, "scheduler")
    )
    # 只有显式声明了固定步表才判定为蒸馏变体
    if config.get("_class_name") != "FlowMatchEulerDiscreteScheduler":
        return None
    sampler = config.get("fixed_step_sampler_config")
    if not sampler or not sampler.get("t_list"):
        return None
    return sampler
​
​
def get_distilled_sigmas(model_path: str) -> list[float] | None:
    """向调度阶段暴露蒸馏 sigma 表；非蒸馏 checkpoint 返回 None。"""
    sampler = _distilled_sampler_config(model_path)
    return list(sampler["t_list"]) if sampler is not None else None

python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py

Edge dense 骨干的模型侧实现：新增 Cosmos3DenseMLP 与 _build_mlp 工厂，Cosmos3CausalAttention 支持可选 QK-norm 与 GEN 侧独立 K 归一化。

class Cosmos3DenseMLP(nn.Module):
    """dense MLP：down(relu(up(x)) ** 2)，Edge 变体的非门控结构。
​
    与 SwiGLU 的 gate_up_proj 合并权重布局不同，dense 变体只有独立的
    up_proj / down_proj，权重加载走未合并的直通路径。
    """
​
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        prefix: str = "",
        quant_config: QuantizationConfig | None = None,
    ):
        super().__init__()
        # up_proj 不合并 gate，直接展开放大中间维度
        self.up_proj = ColumnParallelLinear(
            hidden_size,
            intermediate_size,
            bias=False,
            gather_output=False,
            quant_config=quant_config,
            prefix=add_prefix("up_proj", prefix),
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
            bias=False,
            input_is_parallel=True,
            quant_config=quant_config,
            prefix=add_prefix("down_proj", prefix),
        )
​
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        up, _ = self.up_proj(x)
        up = F.relu(up)
        out, _ = self.down_proj(up * up) # squared-ReLU：relu(x) ** 2
        return out
​
​
def _build_mlp(
    hidden_act: str,
    hidden_size: int,
    intermediate_size: int,
    prefix: str,
    quant_config: QuantizationConfig | None,
) -> nn.Module:
    """按 hidden_act 选择 MLP 结构：relu2 走 dense，其余走 SwiGLU 门控。"""
    mlp_cls = Cosmos3DenseMLP if hidden_act == "relu2" else Cosmos3GatedMLP
    return mlp_cls(
        hidden_size=hidden_size,
        intermediate_size=intermediate_size,
        prefix=prefix,
        quant_config=quant_config,
    )

python/sglang/multimodal_gen/configs/sample/cosmos3.py

变体采样默认值（分辨率、guidance、CFG 区间、步数）在此解析，直接影响每个请求的采样行为。

def _resolve_variant_defaults(self, is_edge: bool, is_distilled: bool = False) -> None:
    """按变体补齐未显式指定的分辨率与 guidance 默认值。
​
    Edge 只在 256p/480p 分辨率上训练过，更大尺寸会把空间 mRoPE
    网格推到训练范围外导致画面破碎，因此必须收窄分辨率白名单并
    给出对应的默认尺寸。
    """
    is_t2i = self.num_frames == 1
    if is_distilled:
        # 蒸馏模型把 guidance 烘焙进权重，推理只需单次 forward，
        # 关闭 CFG 并强制 guidance_scale = 1.0
        self.guidance_scale = 1.0
    elif is_edge and not self._guidance_is_explicit():
        # 用户显式传入 guidance_scale 时尊重显式值，避免被默认覆盖
        self.guidance_scale = (
            COSMOS3_EDGE_T2I_GUIDANCE_SCALE
            if is_t2i
            else COSMOS3_EDGE_T2V_GUIDANCE_SCALE
        )
    # T2I 只在噪声高的 [400, 1000] 区间施加 CFG；
    # 低噪声步的 guidance 已被证明会损害采样质量
    if is_t2i and not is_distilled and self.guidance_interval is None:
        self.guidance_interval = COSMOS3_T2I_GUIDANCE_INTERVAL
    if is_edge:
        self.supported_resolutions = COSMOS3_EDGE_SUPPORTED_RESOLUTIONS
        if self.height is None and self.width is None:
            if is_t2i:
                self.width = self.height = COSMOS3_EDGE_T2I_SIZE
            else:
                self.width, self.height = (
                    COSMOS3_EDGE_T2V_WIDTH,
                    COSMOS3_EDGE_T2V_HEIGHT,
                )

评论区精华

围绕 5 个话题展开，其中 3 个被采纳修复、2 个被维护者明确拒绝。核心交锋：① gemini-code-assist[bot] 指出蒸馏 sigma 双重 shift 问题，mickqian 确认在 commit 78ce372 修复；② kediwu0331 发现 guidance_scale 默认 None 导致 warmup 崩溃，作者随即修复；③ mickqian 两次拒绝 bot 的防御性建议（model_path 空值防护、inspect.signature try-except），理由是 fail-fast 优于静默选错变体、过滤契约必须可检查；④ bot 建议 generator 用 latents.device，已被采纳。

• 蒸馏 checkpoint 的 sigma 双重 shift (correctness): mickqian 确认已在 commit 78ce372 修复：蒸馏路径先 set_shift(1.0) 再应用 checkpoint 提供的 sigma 表。
• guidance_scale 默认 None 导致 warmup 崩溃 (correctness): nv-dmajchrowski 已推送修复，恢复有值默认并配合 _guidance_is_explicit 判断。
• is_edge_checkpoint 是否需要空 model_path 防护 (question): 不添加防护，保持 fail-fast。
• inspect.signature 是否加 try-except 兜底 (design): 不添加兜底。
• 随机噪声 generator 的 device 选择 (correctness): mickqian 确认已在 commit 78ce372 改为 latents.device。

风险与影响

• 风险：
  1. 调度器配置键过滤：scheduler_loader._supported_init_kwargs 依赖 inspect.signature 过滤配置键，若调度器类不可检查（如动态生成类），当前实现会直接抛异常；mickqian 明确拒绝加 try-except 兜底，该假设是硬性的，未来接入新调度器时需要留意。
  2. checkpoint 配置读取崩溃：is_edge_checkpoint 与 _distilled_sampler_config 在 model_path 为空或路径非法时 os.path.join 会抛 TypeError，且 get_diffusers_component_config 读取失败会中断服务启动——这是有意的 fail-fast，但对间接调用路径（如测试 mock 或非 Cosmos3 配置）可能引入意外崩溃。
  3. CFG 分支性能回退：_predict_noise_cfg 在正负 prompt 长度不一致时拆成两次独立 forward，Edge 默认 negative_prompt="" 与长 prompt 之间几乎必然触发该路径，推理耗时约为合并批次的两倍（但避免了 padding 造成的质量损失）。
  4. 蒸馏步数契约：distilled 路径忽略用户传入的 num_inference_steps，步数固定为 sigma 表长度，用户期望与实际行为可能不一致。
  5. SP 噪声行为改变：SP 下 SDE 噪声从各 rank 独立随机改为种子一致，修复 shard-boundary artifact 的同时改变了既有输出的随机性分布，属于有意的行为变更。
  6. multipart form 新逻辑：action 端点的 multipart 解析、上传文件落盘与临时目录清理是新增入口逻辑，文件路径拼接使用 request_id + filename，存在路径注入与文件清理风险，需要验证。
     - 影响：功能影响：Cosmos3 全家族（Nano/Super/Edge，普通与蒸馏）由一个管线统一服务，checkpoint 变体在加载时自动识别；旧 checkpoint 的 _class_name 拼写（Cosmos3OmniDiffusersPipeline）继续可用，无需用户迁移。系统影响：扩散管线的时间步准备与噪声生成路径改动影响所有 Cosmos3 请求；蒸馏请求默认走固定 sigma 表，不再受 num_inference_steps 控制。团队影响：diffusion 维护者需要理解变体检测的缓存时机与 CFG 拆分语义；后续新增变体只需扩展检测函数与默认值表。用户影响：Edge 分辨率被限制在 256p/480p，显式请求更高分辨率会被拒绝或警告，属于有意的产品约束，需要文档说明。
     - 风险标记：核心路径变更, checkpoint 格式兼容, 蒸馏调度契约, CFG 分支性能回退, SP 噪声一致性

关联脉络

• PR #24994 Add Cosmos3 model family support: PR body 引用的 Cosmos3 基础支持，本 PR 是对该家族的变体扩展。
• PR #26492 Cosmos3 follow-ups: PR body 引用的 Cosmos3 后续支持，与本 PR 同属 Cosmos3 功能线。
• PR #27168 Cosmos3 extensions: PR body 引用的 Cosmos3 扩展，与本 PR 同属 Cosmos3 功能线。
• PR #30569 New model format support: PR body 明确说明 copied changes from #30569 to make it work with new model format。
• PR #34249 [diffusion] move DiT execution capabilities to runtime models: 同属 diffusion 架构向 runtime 模型层集中的演进脉络，涉及相同文件 runtime/models/dits 与 configs/models/dits。
• PR #34401 Fix model-driven DiT layerwise offload auto policy: 同属 diffusion 模块的 configs/models/dits 与 runtime/platforms 演进，涉及 DiT 配置与运行时一致性。