尧图网页设计 
MLX-VLM 运行 NVIDIA Nemotron Labs DiffusionAR、块扩散与线性自推测三种解码模式实战指南【免费下载链接】mlx-vlmMLX-VLM is a package for inference and fine-tuning of Vision Language Models (VLMs) on your Mac using MLX.项目地址: https://gitcode.com/GitHub_Trending/ml/mlx-vlm本篇技术指南以mlx_vlm/models/nemotron_labs_diffusion/README.md为骨架结合该模型在 mlx-vlm 仓库中的完整源码实现模型配置、生成后端、CLI 参数定义系统讲解如何在自己的 Mac 上通过 MLX 运行 NVIDIA 的 8B 纯文本扩散语言模型。读完本文你将掌握普通自回归生成、带掩码块去噪的可视化扩散生成、线性自推测加速以及 Thinking 模式共四种使用方式并理解generation_mode、sampler、head_scoring等模型专属参数背后的底层行为。一、模型与能力总览Nemotron Labs Diffusion 是 NVIDIA 出品的纯文本扩散语言模型text-only diffusion language model其最特别之处在于同一个 checkpoint 同时支持三种解码范式无需切换权重文件Text generation自回归生成走标准的mlx_vlm.generate自回归路径与普通 LLM 用法完全一致Diffusion generation扩散生成基于掩码块去噪masked block denoising开启--verbose时可实时观察掩码块被逐步填满的可视化过程Linear self-speculation线性自推测用扩散过程作为草稿draft再用自回归进行验证verify以更少的前向次数获得更多 token。此外还支持Thinking mode通过 chat template 的--enable-thinking开关启用思考模式。官方支持的唯一公开权重如下表所示来源README 原文ModelTypeParamsContextModalitiesnvidia/Nemotron-Labs-Diffusion-8BDense diffusion LM8B262kText值得注意的是该模型不支持任何视觉/音频/视频输入。在仓库的 prompt_utils.py 中nemotron_labs_diffusion被注册为MessageFormat.TEXT_ONLY模型包装类也会在收到pixel_values时直接抛出ValueError(Nemotron Labs Diffusion is a text-only model.)见 nemotron_labs_diffusion.py。二、安装与 mlx-vlm 其他模型一致通过 pip 安装最新版本即可pip install -U mlx-vlm安装后模型解析由仓库 utils.py 中的get_model_and_args完成读取 checkpoint 的config.json中model_type字段此处为nemotron_labs_diffusion再通过importlib动态导入mlx_vlm.models.nemotron_labs_diffusion包其init.py 导出Model与ModelConfig。因此首次运行时只需要--model nvidia/Nemotron-Labs-Diffusion-8B无需手动注册。三、CLI 实战四种运行模式3.1 自回归生成默认路径不传任何扩散相关参数时Nemotron 走普通自回归生成路径mlx_vlm.generate \ --model nvidia/Nemotron-Labs-Diffusion-8B \ --prompt Write a short story about a clockmaker. \ --max-tokens 256 \ --temperature 0.0--temperature 0.0表示贪心解码该路径与普通 LLM 的逐 token 生成完全一致。从源码看AR 路径由 language.py 的ar_generate实现prefill 阶段用因果掩码use_causal_maskTrue逐层填充 KV cache然后逐 token 采样。3.2 扩散生成masked block denoising通过--gen-kwargs传入generation_modediffusion即可切换到掩码扩散路径mlx_vlm.generate \ --model nvidia/Nemotron-Labs-Diffusion-8B \ --prompt Write a short story about a clockmaker. \ --max-tokens 256 \ --max-denoising-steps 16 \ --temperature 0.0 \ --gen-kwargs {generation_mode: diffusion} \ --verbose要点说明generation_mode是模型专属的生成参数由 Nemotron 后端解释不适用于其他模型--max-denoising-steps 16将默认的 32 步去噪上限降为 16适合做速度实验详见下文默认调度说明--verbose会启用DiffusionUnmaskingVisualizer实现在 diffusion_visualizer.py在掩码逐步被填满时输出块级可视化需要终端为 TTY 才会激活sys.stdout.isatty()。默认调度Nemotron 遵循上游/Transformers 的转移策略——默认 32 步去噪上限default_diffusion_steps 32与 0.9 置信度转移阈值default_diffusion_threshold 0.9这两个默认值直接定义在 config.py 中。同时原生模式native sampler还会使用Transformers 对齐运行时transformers-parity runtime来执行去噪编码器——见 language.py该运行时把注意力切回 Transformers eager 实现、并关闭 MLX 特化的 BM32 GEMM kernel以保证与上游数值行为一致。上游模式别名generation_modedlm同样被接受dlm 即 diffusion language model。3.3 线性自推测生成mlx_vlm.generate \ --model nvidia/Nemotron-Labs-Diffusion-8B \ --prompt Write a short story about a clockmaker. \ --max-tokens 256 \ --temperature 0.0 \ --gen-kwargs {generation_mode: linear_speculative}generation_modelinear_speculative透传给 Nemotron 后端启用线性自推测路径Hugging Face 仓库中随权重附带的linear_spec_lora适配器会在可用时自动加载见 language.py 的load_linear_spec_loraLoRA rank128、scale4.0只挂在每层self_attn.o_proj上上游模式别名generation_modelinear_spec同样被接受。3.4 Thinking 模式mlx_vlm.generate \ --model nvidia/Nemotron-Labs-Diffusion-8B \ --prompt Solve this step by step: if a train travels 180 km in 2.5 hours, what is its average speed? \ --enable-thinking \ --max-tokens 512 \ --temperature 0.0--enable-thinking在 dispatch.py 中定义启用后支持thinking_mode的 chat template 会收到thinking_modeenabled从而在回答前先产出思考内容。四、Python API 编程接口除 CLI 外四种能力都可以通过 Python 直接调用与仓库顶层mlx_vlm.generate/mlx_vlm.load/mlx_vlm.prompt_utils.apply_chat_template组合使用。4.1 基础文本生成from mlx_vlm import generate, load from mlx_vlm.prompt_utils import apply_chat_template model, processor load(nvidia/Nemotron-Labs-Diffusion-8B) prompt apply_chat_template( processor, model.config, Write a short story about a clockmaker., ) result generate( modelmodel, processorprocessor, promptprompt, max_tokens256, temperature0.0, ) print(result.text)4.2 扩散生成from mlx_vlm import generate, load from mlx_vlm.prompt_utils import apply_chat_template model, processor load(nvidia/Nemotron-Labs-Diffusion-8B) prompt apply_chat_template( processor, model.config, Write a short story about a clockmaker., ) result generate( modelmodel, processorprocessor, promptprompt, max_tokens256, max_denoising_steps16, temperature0.0, generation_modediffusion, ) print(result.text)注意 Python 接口中generation_mode是generate()的关键字参数而 CLI 中需要放进--gen-kwargs的 JSON 里。max_denoising_steps会被 diffusion.py 映射为后端generate()的steps参数。4.3 线性自推测生成from mlx_vlm import generate, load from mlx_vlm.prompt_utils import apply_chat_template model, processor load(nvidia/Nemotron-Labs-Diffusion-8B) prompt apply_chat_template( processor, model.config, Write a short story about a clockmaker., ) result generate( modelmodel, processorprocessor, promptprompt, max_tokens256, temperature0.0, generation_modelinear_speculative, ) print(result.text)五、模型专属生成参数详解generation_mode、sampler、head_scoring、ar_weight等都是通过--gen-kwargs透传到 Nemotron 后端的模型专属参数其解析与校验逻辑全部集中在 language.py 的generate()方法开头。5.1 generation_mode 路由生成管线的统一入口是 diffusion.py 的_uses_model_diffusion_generator当模型拥有扩散生成器Nemotron 满足该条件且generation_mode ! ar时请求会进入扩散路径。模型默认模式由配置项default_generation_mode ar决定config.py因此不传参时保持自回归行为。三个有效取值如下generation_mode取值含义别名ar普通自回归生成默认无需显式指定diffusion掩码块扩散去噪dlmlinear_speculative扩散起草 AR 验证linear_spec5.2 sampler采样器变体从 NVIDIA 官方评测工具链移植了多种采样器变体通过sampler指定。源码中的合法集合见 language.py同时还有一批别名optimized/bound→confidence_threshold_boundhf/upstream→native等。sampler 取值说明默认值相关native默认值0.9 置信度阈值 Transformers 对齐运行时默认 32 步、阈值 0.9confidence_threshold_bound有界置信度采样min_threshold默认 0.45保留优化过的 MLX kernel适合快速实验采样缩放因子默认 2.0配置default_diffusion_sampling_scaling_factorfixed每步固定转移num_to_transfer个 token—confidence_threshold_ref参考式置信度阈值带位置相关的动态阈值—cumulative_error累积对数置信度误差采样—以samplerconfidence_threshold_bound为例其有界逻辑见 language.py是token 置信度要么达到全局阈值 0.9要么不低于min_threshold0.45且同时不低于按位置计算的动态阈值1 - scaling_factor / (pos 2)。5.3 head_scoring打分策略用于性能剖析head_scoringfull默认物化完整词表 logits 打分。之所以是默认值是因为在 MLX 优化过的 matmul 路径上通常更快head_scoringchunked不拼接完整词表 logits而是分块对掩码行打分实现为_chunked_greedy_score_weight见 language.py专用于剖析非物化打分器的内存/速度特征。从配置看词表被分为 4 块进行 greedy 打分greedy_score_chunks 4config.py 对应逻辑见 language.py。5.4 ar_weight混合 AR 扩散实验ar_weight取值范围0.0~1.0默认0.0禁用。当大于 0 时每次去噪前向之外会额外执行一次因果块前向causal block forward然后把双向隐藏状态与因果隐藏状态按(1 - ar_weight) * bidirectional ar_weight * causal混合language.py。超出范围会直接抛出ValueError。5.5 其他可透传参数CLI 层面 dispatch.py 还提供了--block-length、--num-to-transfer、--max-transfer-per-step、--editing-threshold、--max-post-steps、--stability-steps、--threshold、--min-threshold、--diffusion-min-canvas-length等扩散相关开关它们会经diffusion_kwargs_from_argsdiffusion.py注入生成 kwargs与--gen-kwargs共同生效。六、扩散生成原理从源码看 masked block denoisingNemotron 的扩散生成并非整段文本一次性去噪而是按 32 token 的块block逐步推进。主循环位于 language.py核心流程如下Preload先用因果 prefill 计算提示词的隐藏状态采样出第一个 token 作为块锚点建块初始化一个长度为block_length默认 32即配置block_size32的块全部位置填入mask_token_id100仅第 0 位放入锚点 token迭代去噪对每个去噪步用双向掩码_make_bidirectional_mask见 language.py对整块做一次前向得到掩码位置的 logits置信度转移对掩码位置的采样置信度排序按所选 sampler 的策略决定本步解封哪些位置原生模式即置信度 0.9最后一步强制补全剩余所有掩码块间衔接块填满后用因果前向计算块尾状态采样出下一个块的锚点 token重复直到达到gen_length或遇到 EOS。每一步转移的 token 数会被记录diffusion_accepted_tokens用于后续统计diffusion_tokens_per_denoise_forward每个去噪前向平均产出的 token 数这是评估扩散路径吞吐的关键指标。七、架构细节README 给出的架构参数与 config.py 完全对应汇总如下Backbone密集 decoder-only 的 Ministral 风格 transformerLayers34 层num_hidden_layers34Hidden size4096hidden_size4096Attention32 个 query 头、8 个 KV 头、head 维度 128num_attention_heads32、num_key_value_heads8、head_dim128MLPSwiGLU中间层 14336intermediate_size14336激活函数siluRoPE从 checkpoint 读取的长上下文 YaRN / Llama 4 风格缩放参数。实现上language.py 的_llama4_attention_scale会按llama_4_scaling_beta与original_max_position_embeddings对注意力分数施加随位置增长的对数缩放配置默认rope_theta1000000.0、max_position_embeddings262144Diffusion head独立的输出投影tie_word_embeddingsFalse覆盖 131072 词表vocab_size131072模型包装类的sanitize会把 checkpoint 中的encoder.*与diffusion_head.*键名规整为 MLX 内部结构nemotron_labs_diffusion.pyMask tokenmask_token_id100EOS 为 11且dlm_paradigmbidirectional。值得一提的工程细节该实现针对 MLX Metal 后端内置了多个手写 Metal kernel仅当mx.metal.is_available()时编译包括_make_bm32_linear_kernel基于 MLX Steel GEMM 的 BM32 专用线性 kernellanguage.py当序列长度恰为 32即整块去噪且为 bf16 时触发_SMALL_ROW_GEMV_KERNEL/_SMALL_ROW_SWIGLU_KERNEL面向 2~8 短序列的融合 GEVM/SwiGLU kernellanguage.pyMLP 还会按序列长度自动分块tiny/small/medium 三档见 language.py。这些 kernel 均有严格的触发条件dtype、序列长度、维度对齐不满足时自动回退到通用实现保证正确性优先。八、线性自推测的工作方式线性自推测的完整实现是linear_spec_generatelanguage.py可概括为扩散起草 AR 验证循环起草阶段用掩码块 双向注意力在linear_spec_lora适配器启用set_linear_spec_lora_enabled(True)的状态下反复前向把掩码逐步填满得到一块草稿 token验证阶段关闭 LoRAenabledFalse用因果自回归前向对整块草稿做一次验证逐位置比对草稿与 AR 输出得到接受长度accepted自适应窗口draft_window从 8 起步若整块全部被接受则翻倍上限 32若接受率 ≤ 一半则减半动态平衡草稿质量与计算量language.py缓存回滚验证后通过_trim_cache把 KV cache 截断回接受长度保证下一轮从正确位置继续。其中 LoRA 适配器仅用于扩散起草阶段且只修改o_proj层——这正是 README 所强调的linear_spec_lora adapter is used only during the diffusion draft phase。仓库还提供了流式版本stream_linear_spec_generatelanguage.py用于逐块产出结果的场景。九、运行统计与调参建议扩散生成会在stats字典中记录模型级指标其中核心三项为diffusion_denoise_nfe去噪阶段的前向次数NFEnumber of forward evaluationsdiffusion_post_block_nfe块间衔接post-block前向次数diffusion_tokens_per_denoise_forward每个去噪前向平均接受的 token 数 diffusion_accepted_tokens / diffusion_denoise_nfe。此外还有diffusion_blocks、diffusion_confidence_steps、diffusion_argmax_only_steps、diffusion_masked_rows_scored、diffusion_mixed_ar_forwards等细粒度计数language.py。若想剖析非物化打分器使用head_scoringchunked。默认调度与速度实验默认 32 步去噪 0.9 置信度阈值。调低--max-denoising-steps可以换取速度但 README 明确警告质量可能快速劣化——建议从 16 起步逐档尝试。原生 sampler 因启用 Transformers 对齐运行时会更保守追求更快 MLX 实验时改用samplerconfidence_threshold_boundmin_threshold默认 0.45可保留优化 kernel。十、测试与验证仓库通过单元测试保证了该实现的数值正确性与 dtype 一致性test_diffusion_models.py 的test_nemotron_labs_diffusion用小配置构造模型验证 bf16 下 logits/top-k/采样概率的 dtype 保持并把_chunked_greedy_score_weight的结果与完整hidden weight.T的 argmax 逐一对照test_models.py 的test_nemotron_labs_diffusion_language_model验证语言模型前向与 dtype 一致性。十一、注意事项与限制汇总模型纯文本不支持图像、音频、视频输入AR 生成应使用不带扩散参数的普通 CLI不传--gen-kwargs中与扩散相关的键扩散生成使用掩码块去噪--verbose会实时展示掩码填充过程需 TTY 终端默认调度为 32 步去噪 0.9 置信度阈值降低--max-denoising-steps提速时需警惕质量下降当前后端生成仅支持batch size 1language.py 会显式抛错扩散与线性自推测都通过模型专属generation_mode暴露CLI 用法为--gen-kwargs {generation_mode: diffusion}Python 接口则为直接传关键字参数上游模式名作为别名同样有效dlm扩散、linear_spec线性自推测可选linear_spec_lora适配器随 Hugging Face 仓库权重发布仅在线性自推测的扩散起草阶段使用其余模式自动忽略。若需深入了解实现细节可继续阅读 模型说明文档、生成后端完整实现 与 统一扩散路由。【免费下载链接】mlx-vlmMLX-VLM is a package for inference and fine-tuning of Vision Language Models (VLMs) on your Mac using MLX.项目地址: https://gitcode.com/GitHub_Trending/ml/mlx-vlm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 MLX-VLM 运行 NVIDIA Nemotron Labs Diffusion：AR、块扩散与线性自推测三种解码模式实战指南
MLX-VLM 运行 NVIDIA Nemotron Labs DiffusionAR、块扩散与线性自推测三种解码模式实战指南【免费下载链接】mlx-vlmMLX-VLM is a package for inference and fine-tuning of Vision Language Models (VLMs) on your Mac using MLX.项目地址: https://gitcode.com/GitHub_Trending/ml/mlx-vlm本篇技术指南以mlx_vlm/models/nemotron_labs_diffusion/README.md为骨架结合该模型在 mlx-vlm 仓库中的完整源码实现模型配置、生成后端、CLI 参数定义系统讲解如何在自己的 Mac 上通过 MLX 运行 NVIDIA 的 8B 纯文本扩散语言模型。读完本文你将掌握普通自回归生成、带掩码块去噪的可视化扩散生成、线性自推测加速以及 Thinking 模式共四种使用方式并理解generation_mode、sampler、head_scoring等模型专属参数背后的底层行为。一、模型与能力总览Nemotron Labs Diffusion 是 NVIDIA 出品的纯文本扩散语言模型text-only diffusion language model其最特别之处在于同一个 checkpoint 同时支持三种解码范式无需切换权重文件Text generation自回归生成走标准的mlx_vlm.generate自回归路径与普通 LLM 用法完全一致Diffusion generation扩散生成基于掩码块去噪masked block denoising开启--verbose时可实时观察掩码块被逐步填满的可视化过程Linear self-speculation线性自推测用扩散过程作为草稿draft再用自回归进行验证verify以更少的前向次数获得更多 token。此外还支持Thinking mode通过 chat template 的--enable-thinking开关启用思考模式。官方支持的唯一公开权重如下表所示来源README 原文ModelTypeParamsContextModalitiesnvidia/Nemotron-Labs-Diffusion-8BDense diffusion LM8B262kText值得注意的是该模型不支持任何视觉/音频/视频输入。在仓库的 prompt_utils.py 中nemotron_labs_diffusion被注册为MessageFormat.TEXT_ONLY模型包装类也会在收到pixel_values时直接抛出ValueError(Nemotron Labs Diffusion is a text-only model.)见 nemotron_labs_diffusion.py。二、安装与 mlx-vlm 其他模型一致通过 pip 安装最新版本即可pip install -U mlx-vlm安装后模型解析由仓库 utils.py 中的get_model_and_args完成读取 checkpoint 的config.json中model_type字段此处为nemotron_labs_diffusion再通过importlib动态导入mlx_vlm.models.nemotron_labs_diffusion包其init.py 导出Model与ModelConfig。因此首次运行时只需要--model nvidia/Nemotron-Labs-Diffusion-8B无需手动注册。三、CLI 实战四种运行模式3.1 自回归生成默认路径不传任何扩散相关参数时Nemotron 走普通自回归生成路径mlx_vlm.generate \ --model nvidia/Nemotron-Labs-Diffusion-8B \ --prompt Write a short story about a clockmaker. \ --max-tokens 256 \ --temperature 0.0--temperature 0.0表示贪心解码该路径与普通 LLM 的逐 token 生成完全一致。从源码看AR 路径由 language.py 的ar_generate实现prefill 阶段用因果掩码use_causal_maskTrue逐层填充 KV cache然后逐 token 采样。3.2 扩散生成masked block denoising通过--gen-kwargs传入generation_modediffusion即可切换到掩码扩散路径mlx_vlm.generate \ --model nvidia/Nemotron-Labs-Diffusion-8B \ --prompt Write a short story about a clockmaker. \ --max-tokens 256 \ --max-denoising-steps 16 \ --temperature 0.0 \ --gen-kwargs {generation_mode: diffusion} \ --verbose要点说明generation_mode是模型专属的生成参数由 Nemotron 后端解释不适用于其他模型--max-denoising-steps 16将默认的 32 步去噪上限降为 16适合做速度实验详见下文默认调度说明--verbose会启用DiffusionUnmaskingVisualizer实现在 diffusion_visualizer.py在掩码逐步被填满时输出块级可视化需要终端为 TTY 才会激活sys.stdout.isatty()。默认调度Nemotron 遵循上游/Transformers 的转移策略——默认 32 步去噪上限default_diffusion_steps 32与 0.9 置信度转移阈值default_diffusion_threshold 0.9这两个默认值直接定义在 config.py 中。同时原生模式native sampler还会使用Transformers 对齐运行时transformers-parity runtime来执行去噪编码器——见 language.py该运行时把注意力切回 Transformers eager 实现、并关闭 MLX 特化的 BM32 GEMM kernel以保证与上游数值行为一致。上游模式别名generation_modedlm同样被接受dlm 即 diffusion language model。3.3 线性自推测生成mlx_vlm.generate \ --model nvidia/Nemotron-Labs-Diffusion-8B \ --prompt Write a short story about a clockmaker. \ --max-tokens 256 \ --temperature 0.0 \ --gen-kwargs {generation_mode: linear_speculative}generation_modelinear_speculative透传给 Nemotron 后端启用线性自推测路径Hugging Face 仓库中随权重附带的linear_spec_lora适配器会在可用时自动加载见 language.py 的load_linear_spec_loraLoRA rank128、scale4.0只挂在每层self_attn.o_proj上上游模式别名generation_modelinear_spec同样被接受。3.4 Thinking 模式mlx_vlm.generate \ --model nvidia/Nemotron-Labs-Diffusion-8B \ --prompt Solve this step by step: if a train travels 180 km in 2.5 hours, what is its average speed? \ --enable-thinking \ --max-tokens 512 \ --temperature 0.0--enable-thinking在 dispatch.py 中定义启用后支持thinking_mode的 chat template 会收到thinking_modeenabled从而在回答前先产出思考内容。四、Python API 编程接口除 CLI 外四种能力都可以通过 Python 直接调用与仓库顶层mlx_vlm.generate/mlx_vlm.load/mlx_vlm.prompt_utils.apply_chat_template组合使用。4.1 基础文本生成from mlx_vlm import generate, load from mlx_vlm.prompt_utils import apply_chat_template model, processor load(nvidia/Nemotron-Labs-Diffusion-8B) prompt apply_chat_template( processor, model.config, Write a short story about a clockmaker., ) result generate( modelmodel, processorprocessor, promptprompt, max_tokens256, temperature0.0, ) print(result.text)4.2 扩散生成from mlx_vlm import generate, load from mlx_vlm.prompt_utils import apply_chat_template model, processor load(nvidia/Nemotron-Labs-Diffusion-8B) prompt apply_chat_template( processor, model.config, Write a short story about a clockmaker., ) result generate( modelmodel, processorprocessor, promptprompt, max_tokens256, max_denoising_steps16, temperature0.0, generation_modediffusion, ) print(result.text)注意 Python 接口中generation_mode是generate()的关键字参数而 CLI 中需要放进--gen-kwargs的 JSON 里。max_denoising_steps会被 diffusion.py 映射为后端generate()的steps参数。4.3 线性自推测生成from mlx_vlm import generate, load from mlx_vlm.prompt_utils import apply_chat_template model, processor load(nvidia/Nemotron-Labs-Diffusion-8B) prompt apply_chat_template( processor, model.config, Write a short story about a clockmaker., ) result generate( modelmodel, processorprocessor, promptprompt, max_tokens256, temperature0.0, generation_modelinear_speculative, ) print(result.text)五、模型专属生成参数详解generation_mode、sampler、head_scoring、ar_weight等都是通过--gen-kwargs透传到 Nemotron 后端的模型专属参数其解析与校验逻辑全部集中在 language.py 的generate()方法开头。5.1 generation_mode 路由生成管线的统一入口是 diffusion.py 的_uses_model_diffusion_generator当模型拥有扩散生成器Nemotron 满足该条件且generation_mode ! ar时请求会进入扩散路径。模型默认模式由配置项default_generation_mode ar决定config.py因此不传参时保持自回归行为。三个有效取值如下generation_mode取值含义别名ar普通自回归生成默认无需显式指定diffusion掩码块扩散去噪dlmlinear_speculative扩散起草 AR 验证linear_spec5.2 sampler采样器变体从 NVIDIA 官方评测工具链移植了多种采样器变体通过sampler指定。源码中的合法集合见 language.py同时还有一批别名optimized/bound→confidence_threshold_boundhf/upstream→native等。sampler 取值说明默认值相关native默认值0.9 置信度阈值 Transformers 对齐运行时默认 32 步、阈值 0.9confidence_threshold_bound有界置信度采样min_threshold默认 0.45保留优化过的 MLX kernel适合快速实验采样缩放因子默认 2.0配置default_diffusion_sampling_scaling_factorfixed每步固定转移num_to_transfer个 token—confidence_threshold_ref参考式置信度阈值带位置相关的动态阈值—cumulative_error累积对数置信度误差采样—以samplerconfidence_threshold_bound为例其有界逻辑见 language.py是token 置信度要么达到全局阈值 0.9要么不低于min_threshold0.45且同时不低于按位置计算的动态阈值1 - scaling_factor / (pos 2)。5.3 head_scoring打分策略用于性能剖析head_scoringfull默认物化完整词表 logits 打分。之所以是默认值是因为在 MLX 优化过的 matmul 路径上通常更快head_scoringchunked不拼接完整词表 logits而是分块对掩码行打分实现为_chunked_greedy_score_weight见 language.py专用于剖析非物化打分器的内存/速度特征。从配置看词表被分为 4 块进行 greedy 打分greedy_score_chunks 4config.py 对应逻辑见 language.py。5.4 ar_weight混合 AR 扩散实验ar_weight取值范围0.0~1.0默认0.0禁用。当大于 0 时每次去噪前向之外会额外执行一次因果块前向causal block forward然后把双向隐藏状态与因果隐藏状态按(1 - ar_weight) * bidirectional ar_weight * causal混合language.py。超出范围会直接抛出ValueError。5.5 其他可透传参数CLI 层面 dispatch.py 还提供了--block-length、--num-to-transfer、--max-transfer-per-step、--editing-threshold、--max-post-steps、--stability-steps、--threshold、--min-threshold、--diffusion-min-canvas-length等扩散相关开关它们会经diffusion_kwargs_from_argsdiffusion.py注入生成 kwargs与--gen-kwargs共同生效。六、扩散生成原理从源码看 masked block denoisingNemotron 的扩散生成并非整段文本一次性去噪而是按 32 token 的块block逐步推进。主循环位于 language.py核心流程如下Preload先用因果 prefill 计算提示词的隐藏状态采样出第一个 token 作为块锚点建块初始化一个长度为block_length默认 32即配置block_size32的块全部位置填入mask_token_id100仅第 0 位放入锚点 token迭代去噪对每个去噪步用双向掩码_make_bidirectional_mask见 language.py对整块做一次前向得到掩码位置的 logits置信度转移对掩码位置的采样置信度排序按所选 sampler 的策略决定本步解封哪些位置原生模式即置信度 0.9最后一步强制补全剩余所有掩码块间衔接块填满后用因果前向计算块尾状态采样出下一个块的锚点 token重复直到达到gen_length或遇到 EOS。每一步转移的 token 数会被记录diffusion_accepted_tokens用于后续统计diffusion_tokens_per_denoise_forward每个去噪前向平均产出的 token 数这是评估扩散路径吞吐的关键指标。七、架构细节README 给出的架构参数与 config.py 完全对应汇总如下Backbone密集 decoder-only 的 Ministral 风格 transformerLayers34 层num_hidden_layers34Hidden size4096hidden_size4096Attention32 个 query 头、8 个 KV 头、head 维度 128num_attention_heads32、num_key_value_heads8、head_dim128MLPSwiGLU中间层 14336intermediate_size14336激活函数siluRoPE从 checkpoint 读取的长上下文 YaRN / Llama 4 风格缩放参数。实现上language.py 的_llama4_attention_scale会按llama_4_scaling_beta与original_max_position_embeddings对注意力分数施加随位置增长的对数缩放配置默认rope_theta1000000.0、max_position_embeddings262144Diffusion head独立的输出投影tie_word_embeddingsFalse覆盖 131072 词表vocab_size131072模型包装类的sanitize会把 checkpoint 中的encoder.*与diffusion_head.*键名规整为 MLX 内部结构nemotron_labs_diffusion.pyMask tokenmask_token_id100EOS 为 11且dlm_paradigmbidirectional。值得一提的工程细节该实现针对 MLX Metal 后端内置了多个手写 Metal kernel仅当mx.metal.is_available()时编译包括_make_bm32_linear_kernel基于 MLX Steel GEMM 的 BM32 专用线性 kernellanguage.py当序列长度恰为 32即整块去噪且为 bf16 时触发_SMALL_ROW_GEMV_KERNEL/_SMALL_ROW_SWIGLU_KERNEL面向 2~8 短序列的融合 GEVM/SwiGLU kernellanguage.pyMLP 还会按序列长度自动分块tiny/small/medium 三档见 language.py。这些 kernel 均有严格的触发条件dtype、序列长度、维度对齐不满足时自动回退到通用实现保证正确性优先。八、线性自推测的工作方式线性自推测的完整实现是linear_spec_generatelanguage.py可概括为扩散起草 AR 验证循环起草阶段用掩码块 双向注意力在linear_spec_lora适配器启用set_linear_spec_lora_enabled(True)的状态下反复前向把掩码逐步填满得到一块草稿 token验证阶段关闭 LoRAenabledFalse用因果自回归前向对整块草稿做一次验证逐位置比对草稿与 AR 输出得到接受长度accepted自适应窗口draft_window从 8 起步若整块全部被接受则翻倍上限 32若接受率 ≤ 一半则减半动态平衡草稿质量与计算量language.py缓存回滚验证后通过_trim_cache把 KV cache 截断回接受长度保证下一轮从正确位置继续。其中 LoRA 适配器仅用于扩散起草阶段且只修改o_proj层——这正是 README 所强调的linear_spec_lora adapter is used only during the diffusion draft phase。仓库还提供了流式版本stream_linear_spec_generatelanguage.py用于逐块产出结果的场景。九、运行统计与调参建议扩散生成会在stats字典中记录模型级指标其中核心三项为diffusion_denoise_nfe去噪阶段的前向次数NFEnumber of forward evaluationsdiffusion_post_block_nfe块间衔接post-block前向次数diffusion_tokens_per_denoise_forward每个去噪前向平均接受的 token 数 diffusion_accepted_tokens / diffusion_denoise_nfe。此外还有diffusion_blocks、diffusion_confidence_steps、diffusion_argmax_only_steps、diffusion_masked_rows_scored、diffusion_mixed_ar_forwards等细粒度计数language.py。若想剖析非物化打分器使用head_scoringchunked。默认调度与速度实验默认 32 步去噪 0.9 置信度阈值。调低--max-denoising-steps可以换取速度但 README 明确警告质量可能快速劣化——建议从 16 起步逐档尝试。原生 sampler 因启用 Transformers 对齐运行时会更保守追求更快 MLX 实验时改用samplerconfidence_threshold_boundmin_threshold默认 0.45可保留优化 kernel。十、测试与验证仓库通过单元测试保证了该实现的数值正确性与 dtype 一致性test_diffusion_models.py 的test_nemotron_labs_diffusion用小配置构造模型验证 bf16 下 logits/top-k/采样概率的 dtype 保持并把_chunked_greedy_score_weight的结果与完整hidden weight.T的 argmax 逐一对照test_models.py 的test_nemotron_labs_diffusion_language_model验证语言模型前向与 dtype 一致性。十一、注意事项与限制汇总模型纯文本不支持图像、音频、视频输入AR 生成应使用不带扩散参数的普通 CLI不传--gen-kwargs中与扩散相关的键扩散生成使用掩码块去噪--verbose会实时展示掩码填充过程需 TTY 终端默认调度为 32 步去噪 0.9 置信度阈值降低--max-denoising-steps提速时需警惕质量下降当前后端生成仅支持batch size 1language.py 会显式抛错扩散与线性自推测都通过模型专属generation_mode暴露CLI 用法为--gen-kwargs {generation_mode: diffusion}Python 接口则为直接传关键字参数上游模式名作为别名同样有效dlm扩散、linear_spec线性自推测可选linear_spec_lora适配器随 Hugging Face 仓库权重发布仅在线性自推测的扩散起草阶段使用其余模式自动忽略。若需深入了解实现细节可继续阅读 模型说明文档、生成后端完整实现 与 统一扩散路由。【免费下载链接】mlx-vlmMLX-VLM is a package for inference and fine-tuning of Vision Language Models (VLMs) on your Mac using MLX.项目地址: https://gitcode.com/GitHub_Trending/ml/mlx-vlm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考