SGLang GDN 混合线性注意力测试能力矩阵:全注意力后端与运行器模式的覆盖解析

发布时间:2026/9/10 13:32:34
SGLang GDN 混合线性注意力测试能力矩阵:全注意力后端与运行器模式的覆盖解析 SGLang GDN 混合线性注意力测试能力矩阵全注意力后端与运行器模式的覆盖解析【免费下载链接】sglangSGLang is a high-performance serving framework for large language models and multimodal models.项目地址: https://gitcode.com/GitHub_Trending/sg/sglangSGLang 在test/registered/attention/unittests/gdn/下为 GDNGated Delta Net混合线性注意力维护了一套系统的单元测试能力矩阵用于验证全注意力后端 Triton GDN 线性注意力内核的组合在各种运行器模式eager、CUDA Graph、piecewise/breakable CUDA Graph、EAGLE 推测验证等下的数值正确性与分发正确性。本文基于该目录下的 README.md结合仓库内测试用例与 hybrid_linear_attn_backend.py 源码实现完整解读这套矩阵的覆盖范围、单元格语义、被阻断的生产路径以及测试设计背后的关键约定帮助开发者快速定位 GDN 注意力在后端与运行器维度上的能力边界。GDN 混合线性注意力与测试定位GDN 注意力是 SGLang 中一类混合线性注意力hybrid linear attention机制模型由若干全注意力层full-attention layers与若干线性注意力层linear-attention layers即 GDN 层混合构成。在 SGLang 的注意力后端体系中这由HybridLinearAttnBackend统一管理——它内部持有两个子后端self.attn_backend_list [full_attn_backend, linear_attn_backend]并通过_is_full_attn()按layer_id ∈ full_attn_layers判断某层应走全注意力后端还是线性注意力后端见 hybrid_linear_attn_backend.py。本目录测试的核心组合约定如下矩阵列是运行器模式runner modes矩阵行是全注意力后端full-attention backendstorch_native、triton、flashinfer线性注意力内核固定为 Triton GDN 内核除 flashinfer 线性 prefill 专项测试外期望输出使用独立的纯 PyTorch gated-delta 递推参考实现pure-PyTorch gated-delta recurrence reference来对比而不是用 Triton/FLA GDN 内核互相印证从而避免用被测实现验证被测实现的同源偏差。测试目录包含五个测试文件gdn 测试目录test_torch_native.py、test_triton.py、test_flashinfer.py、test_gdn_cutedsl_ring_verify.py、test_gdn_replayssm_spec_fold.py、test_linear_replayssm_decode.py。其中前三个对应 README 中的三行后端矩阵均通过register_cuda_ci/register_amd_ci注册进 CI例如 test_triton.py 中注册了4-gpu-b200与1-gpu-large的 CUDA 阶段以及stage-b-test-1-gpu-large-amd的 AMD 阶段。覆盖矩阵单元格语义与运行器模式解读单元格标记约定标记含义✓ variants该组合已被测试覆盖单元格列出覆盖的配置变体—不适用该组合没有生产路径blocked: reason生产环境不支持不是后续跟进项deferred: reason未来可能落地当前被禁用blocked与deferred的区别至关重要blocked意味着该组合在架构上就不可达不应作为待办事项去补测试deferred才是有意延后的能力。运行器模式列说明矩阵的 12 列覆盖了 SGLang 前向执行的主要运行器形态Eager Phase 2eager 模式下的扩展EXTEND执行对应 GDN 代表性输入全量扫描CG decodeCUDA Graph 捕获/回放的 decode 执行PCG extend / BCG extendpiecewise CUDA Graph 与 breakable CUDA Graph 两种扩展执行路径。在 test_triton.py 的 split-op 测试中通过runner bcg if breakable else pcg区分二者Verify eager / Verify CG推测解码中目标模型对草稿 token 的验证TARGET_VERIFY分 eager 与 CUDA Graph 两种形态DE eager / DE CG / DE-V2 CG草稿扩展DRAFT_EXTEND相关的运行器形态EAGLE-draft runner / EAGLE-DE runner / FKVMTP runnerEAGLE 草稿模型运行器、EAGLE 与草稿扩展结合的运行器、Frozen-KV MTPfrozen_kv_mtp运行器。完整覆盖矩阵全注意力后端Eager Phase 2CG decodePCG extendBCG extendVerify eagerVerify CGDE eagerDE CGDE-V2 CGEAGLE-draft runnerEAGLE-DE runnerFKVMTP runnertorch_native✓ 完整代表性 GDN 输入扫描—TorchNativeAttnBackend无 CUDA Graph 钩子✓ ragged 页边界扩展✓ ragged 页边界扩展————————triton✓ 完整代表性 GDN 输入扫描✓ decode 页边界✓ ragged 页边界扩展✓ ragged 页边界扩展✓ EAGLE chaintopk1 EAGLE treetopk2✓ EAGLE chain EAGLE treetree 使用 scoped5e-2atol 容忍 bf16 递推累积误差—blockedHybridLinearAttnBackend._replay_metadata拒绝DECODE_OR_IDLE/TARGET_VERIFY之外的模式hybrid_linear_attn_backend.py:509,572blocked同上—blocked同上—flashinfer✓ 完整 GDN 扫描head_dim64FlashInfer SM90 prefill 约束✓ decode 页边界✓ ragged 页边界扩展✓ ragged 页边界扩展✓ EAGLE chaintopk1 EAGLE treetopk2✓ EAGLE chain EAGLE treescoped5e-2atol—blocked同上blocked同上—blocked同上—矩阵的读法torch_native后端只有 eager 与 extend 类路径被覆盖TorchNativeAttnBackend没有实现 CUDA Graph 钩子因此 CG decode 与所有 Verify 列均为—不适用这是能力缺失而非被阻断triton与flashinfer后端覆盖面一致eager、CG decode、PCG/BCG extend、Verify eager、Verify CG 全部覆盖而DE CG、DE-V2 CG、EAGLE-DE runner 三个草稿扩展相关的 CUDA Graph 列被blocked原因是同一处底层约束flashinfer行在 eager 列多了一个head_dim64的前提FlashInfer SM90 prefill 内核要求 value head dim ∈ {64, 128, 256}test_flashinfer.py 中即固定HEAD_K_DIM 64、HEAD_V_DIM 64并把该维度贯穿到所有用例调用run_gdn_attention_case(..., head_k_dim64, head_v_dim64)。逐后端深度解析torch_native覆盖 eager 与扩展路径test_torch_native.py 由make_gdn_cases(torch_native)生成代表性输入扫描EXTEND/DECODE 等核心形态并额外覆盖两类 casesplit-op 扩展runner_split_op_gdn_extend_ragged_page_boundaryprefix_lens(0, 8, 16)、extend_lens(15, 8, 1)page_size16静态 token 缓冲数为 32同时以breakableFalse/True分别驱动pcg与bcg两条 CUDA Graph 路径用于验证live-token 切片在更大静态 token 缓冲下的正确性布局鲁棒性layout robustnessinterleaved_pages与non_monotonic_extend两种激进布局默认测试使用shuffled_pages验证后端对非规整页布局的处理。注意该文件的 split-op 测试在 ROCm 上被跳过split-op extend 运行器走的是 piecewise-CUDA-Graph 路径TcPiecewiseForwardContext.num_tokens该路径尚未在 ROCm 上接通。triton覆盖最完整的基准后端test_triton.py 是覆盖度最高的文件矩阵中triton行的每个 ✓ 单元格都有对应实现CG decoderunner_cuda_graph_gdn_decode_page_boundaryprefix_lens(14, 15, 16)page_size16覆盖 decode 跨越页边界时状态索引的处理Verifyeager CGEAGLE_VERIFY_CASES与EAGLE_VERIFY_CUDA_GRAPH_CASES两个元组覆盖了EAGLE chaintopk1与 EAGLE treetopk2并进一步覆盖frozen_kv_mtp、dflash、ngram三种非 EAGLE 的链式规格类型——注释明确指出这三种规格均通过spec_info的 custom/tree mask 被 GDN 后端统一处理且全部通过纯 PyTorch gated-delta 递推参考实现验证布局鲁棒性与 torch_native 相同的interleaved_pages/non_monotonic_extend布局测试decode 模式跳过non_monotonic_extend因为该布局对 decode 无意义。flashinfer硬件约束与线性 prefill 专项test_flashinfer.py 在矩阵覆盖之外还包含一个独立测试类TestFlashInferLinearGDNBackendCorrectness专门验证FlashInfer 线性 prefill 后端与 Triton GDN 内核的等价性要求 SM90或 SM100/SM103 搭配 CUDA 13_supports_flashinfer_linear_gdn判定且 FlashInfer DSL prefill 内核在 SM90/SM100 上要求 head size 128因此该测试HEAD_DIM 128测试用例flashinfer_gdn_prefill_state_checkpoints以prefix_lens(0, 64, 128)、extend_lens(64, 65, 129)模拟extra-buffer 调度器产生的 tracking 元数据mamba_track_mask、mamba_track_indices、mamba_track_seqlens覆盖 Mamba 状态检查点映射与状态拷贝逻辑测试将同一 fixture 先后用 FlashInfer 与TritonGDNKernel()跑 eager 前向分别从mamba2_layer_cache的conv[0]/temporal缓冲中取出被追踪的状态以atol3e-2, rtol3e-2断言两路输出与两路追踪状态一致——这是线性注意力内核固定为 Triton这一矩阵约定之外的交叉验证特例。Hybrid 分发 fan-out 测试用 MagicMock 验证分发切片矩阵之外README 单独记录了三个仅针对triton后端、基于MagicMock的分发 fan-out 测试。它们验证的不是数值正确性而是HybridLinearAttnBackend分发层本身每个测试用两个MagicMock子后端构造一个HybridLinearAttnBackend断言两个子后端都收到了匹配的调用。测试覆盖的变更mutationtest_hybrid_dispatch_eager_init_forward_metadata_fan_outM20 —init_forward_metadata中attn_backend_list[1:]切片hybrid_linear_attn_backend.py:825-827test_hybrid_dispatch_replay_init_forward_metadata_fan_outM19 —init_forward_metadata_replay_cuda_graph中attn_backend_list[:1]切片hybrid_linear_attn_backend.py:879-900test_hybrid_dispatch_capture_init_forward_metadata_fan_out对称的 capture 覆盖未收录进 mutation journal为何不直接断言前向输出、而要在MagicMock层面 spytest_triton.py 的注释给出了关键理由分发层的切片变更slice mutation在 fixture 恰好使用相同的 capture/replay 元数据时可能被前向输出断言漏掉——只有直接监视每个子后端的init_forward_metadata*方法才能在缺失调用时立即暴露问题。三个测试的验证逻辑值得注意eager 测试通过SimpleNamespace(forward_modeSimpleNamespace(is_draft_extend_v2lambda: False))构造 sentinel forward batch强制进入两个子后端都分发的路径因为DRAFT_EXTEND_V2模式下分发层会跳过线性后端见 hybrid_linear_attn_backend.pyreplay/capture 测试使用_make_sentinel_fb()构造带batch_size、seq_lens、spec_info等属性的 sentinel batch断言采用sentinel 对象身份identity匹配而非精确的(args, kwargs)形状匹配_assert_fanout_forwarded只要求每个 sentinel 出现在调用的位置参数或关键字参数中从而容忍生产代码在 positional ↔ keyword 转发间的重构避免测试因参数传递形式变化而脆弱失效。输入与配置覆盖范围README 列出的输入/配置覆盖包括页大小 1gdn_extend_page_size_1整页对齐exact-page如零前缀 恰好一页extend_lens(16,)跨页crossing-pageragged 页边界ragged page-boundary如prefix_lens(0, 8, 16)与extend_lens(15, 8, 1)的组合页大小 32 的跨页page-size-32 crossingdecode 边界decode boundary如prefix_lens(14, 15, 16)让 decode 恰好落在页边界附近batch-size-1 decode用例。这些 case 由 gdn_attention.py 中的make_gdn_cases(backend)生成GDNAttentionCase数据类字段包括name、backend、forward_mode、num_k_heads、num_v_heads、page_size、prefix_lens、extend_lens、linear_attn_prefill_backend默认num_k_heads2、num_v_heads2head dim 默认 32dtype 默认torch.bfloat16。两条与验证相关的补充约定GDN 使用推测式 Mamba 状态缓冲来覆盖 target verifyTARGET_VERIFY路径——验证阶段的状态索引来自推测解码的中间状态索引构建逻辑split-op 测试用更大的静态 token 缓冲验证 live-token 切片如静态缓冲 32、实际 token 更少的组合确保 CUDA Graph 静态形状下切出真实 token 的逻辑正确。生产环境不支持的组合_capture_metadata/_replay_metadata的硬约束矩阵中多个blocked单元格指向同一个底层契约MambaAttnBackendBase._capture_metadata/_replay_metadatahybrid_linear_attn_backend.py只接受两类 forward modeDECODE_OR_IDLEforward_mode.is_decode_or_idle()decode 或空闲TARGET_VERIFYforward_mode.is_target_verify()目标模型验证。任何其他模式都会抛出ValueError(fInvalid forward mode: {forward_mode})。这是 GDN 的Mamba2AttnBackend、KDA、Lightning、Mamba2 共享的基础约定因此DRAFT_EXTEND/DRAFT_EXTEND_V2的 CUDA Graph 捕获/回放对 GDN 线性注意力一侧是结构性不可达的structurally unreachable——不是暂时未实现而是被后端契约明确拒绝这解释了矩阵中 DE CG、DE-V2 CG、EAGLE-DE runner 三列在 triton/flashinfer 行均为blocked而非deferred的原因。另一条约束在_forward_metadatahybrid_linear_attn_backend.py:246附近非 decode、非 extend 的模式同样抛出ValueError。合法模式集合为is_decode_or_idle()DECODE / IDLEis_extend(include_draft_extend_v2True)按 forward_batch_info.py 中ForwardMode枚举的定义该谓词覆盖EXTEND / MIXED / DRAFT_EXTEND / DRAFT_EXTEND_V2 / TARGET_VERIFY / SPLIT_PREFILL / DLLM_EXTEND。理解这两条约束就能快速判断某运行器模式 × GDN 后端组合是否值得写测试凡涉及 GDN 线性注意力侧的 CUDA Graph 捕获/回放且模式不在上述白名单内都属于生产不可达路径不应作为测试欠覆盖项。Caveats初始 SSM 状态恒为零README 明确记录了一个测试语义上的注意事项build_gdn_attention_fixture不会像 dense 的_populate_prefix_kv那样把前缀 token 真正跑一遍模块因此 SSM 状态缓冲始终停留在运行器初始化时的零状态。这带来两个推论prefix_lens 0的用例中实际路径与参考路径都从零状态开始因此二者匹配是平凡成立的换句话说非零prefix_lens只用于锻炼 metadata 路径页布局、状态索引、track 逻辑并不验证递推状态的延续recurrent-state continuation。在阅读本目录测试结果时务必记住这一边界extend 测试验证的是零初始状态下给定分页/边界布局时 GDN 内核与纯 PyTorch 参考的一致性而非长序列前缀状态跨 chunk 正确累积。后者由矩阵之外的状态检查点测试如 FlashInfer 线性 prefill 的 checkpoint 用例以显式模拟 tracking 元数据的方式覆盖。数值容差约定参考 gdn_attention.py 中的默认容差普通用例GDN_ATOL 3e-2、GDN_RTOL 3e-2EAGLE treetopk2用例GDN_TREE_ATOL 5e-2——矩阵中 Verify CG 列标注的scoped5e-2atol即由此而来原因是 bf16 递推在树形验证布局下会有更大的累积误差需要更宽松的容忍度FlashInfer 线性 prefill 与 Triton 内核的交叉对比同样使用atol3e-2, rtol3e-2。后续工作方向README 末尾列出两项待办可作为 GDN 测试能力演进的参考在可用时补充更多线性注意力内核后端变体当前矩阵线性注意力一侧固定为 Triton GDN 内核FlashInfer 线性 prefill 已作为专项测试先行覆盖在 EAGLE chain/tree 在多个内核上保持稳定之前暂不扩大推测式 worker 的标签范围——即先确保现有 EAGLE 验证覆盖在 kernel 维度上的稳定性再考虑更多 speculative worker 组合。小结这张能力矩阵的价值在于把后端 × 运行器模式二维空间中的覆盖状态显式化✓ 表示已由数值或分发测试覆盖— 表示无生产路径blocked 表示架构性不可达deferred 表示有意延后。对于需要在 SGLang 上开发或调试 GDN 注意力或同类 Mamba 系线性注意力如 KDA、Lightning、Mamba2的工程师这张矩阵与 test_triton.py、test_torch_native.py、test_flashinfer.py 三个文件共同构成了理解哪些组合值得测、哪些组合测不了、为什么测不了的完整参考。【免费下载链接】sglangSGLang is a high-performance serving framework for large language models and multimodal models.项目地址: https://gitcode.com/GitHub_Trending/sg/sglang创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考