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

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

【免费下载链接】sglangSGLang is a high-performance serving framework for large language models and multimodal models.项目地址: https://gitcode.com/GitHub_Trending/sg/sglang

SGLang 在test/registered/attention/unittests/gdn/下为 GDN（Gated 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 backends）：torch_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 阶段）。

覆盖矩阵：单元格语义与运行器模式解读

单元格标记约定

blocked与deferred的区别至关重要：blocked意味着该组合在架构上就不可达，不应作为待办事项去"补测试"；deferred才是有意延后的能力。

运行器模式（列）说明

矩阵的 12 列覆盖了 SGLang 前向执行的主要运行器形态：

• Eager Phase 2：eager 模式下的扩展（EXTEND）执行，对应 GDN 代表性输入全量扫描；
• CG decode：CUDA Graph 捕获/回放的 decode 执行；
• PCG extend / BCG extend：piecewise 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 runner：EAGLE 草稿模型运行器、EAGLE 与草稿扩展结合的运行器、Frozen-KV MTP（frozen_kv_mtp）运行器。

完整覆盖矩阵

矩阵的读法

• 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_dim=64的前提：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_dim=64, head_v_dim=64)）。

逐后端深度解析

torch_native：覆盖 eager 与扩展路径

test_torch_native.py 由make_gdn_cases("torch_native")生成代表性输入扫描（EXTEND/DECODE 等核心形态），并额外覆盖两类 case：

• split-op 扩展：runner_split_op_gdn_extend_ragged_page_boundary，prefix_lens=(0, 8, 16)、extend_lens=(15, 8, 1)，page_size=16，静态 token 缓冲数为 32，同时以breakable=False/True分别驱动pcg与bcg两条 CUDA Graph 路径，用于验证live-token 切片在更大静态 token 缓冲下的正确性；
• 布局鲁棒性（layout robustness）：interleaved_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 decode：runner_cuda_graph_gdn_decode_page_boundary，prefix_lens=(14, 15, 16)，page_size=16，覆盖 decode 跨越页边界时状态索引的处理；
• Verify（eager + CG）：EAGLE_VERIFY_CASES与EAGLE_VERIFY_CUDA_GRAPH_CASES两个元组覆盖了EAGLE chain（topk=1）与 EAGLE tree（topk=2），并进一步覆盖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缓冲中取出被追踪的状态，以atol=3e-2, rtol=3e-2断言两路输出与两路追踪状态一致——这是"线性注意力内核固定为 Triton"这一矩阵约定之外的交叉验证特例。

Hybrid 分发 fan-out 测试：用 MagicMock 验证分发切片

矩阵之外，README 单独记录了三个仅针对triton后端、基于MagicMock的分发 fan-out 测试。它们验证的不是数值正确性，而是HybridLinearAttnBackend分发层本身：每个测试用两个MagicMock子后端构造一个HybridLinearAttnBackend，断言两个子后端都收到了匹配的调用。

为何不直接断言前向输出、而要在MagicMock层面 spy？test_triton.py 的注释给出了关键理由：分发层的切片变更（slice mutation）在 fixture 恰好使用相同的 capture/replay 元数据时，可能被前向输出断言漏掉——只有直接监视每个子后端的init_forward_metadata*方法，才能在缺失调用时立即暴露问题。

三个测试的验证逻辑值得注意：

• eager 测试通过SimpleNamespace(forward_mode=SimpleNamespace(is_draft_extend_v2=lambda: False))构造 sentinel forward batch，强制进入"两个子后端都分发"的路径（因为DRAFT_EXTEND_V2模式下分发层会跳过线性后端，见 hybrid_linear_attn_backend.py）；
• replay/capture 测试使用_make_sentinel_fb()构造带batch_size、seq_lens、spec_info等属性的 sentinel batch；
• 断言采用sentinel 对象身份（identity）匹配而非精确的(args, kwargs)形状匹配：_assert_fanout_forwarded只要求每个 sentinel 出现在调用的位置参数或关键字参数中，从而容忍生产代码在 positional ↔ keyword 转发间的重构，避免测试因参数传递形式变化而脆弱失效。

输入与配置覆盖范围

README 列出的输入/配置覆盖包括：

• 页大小 1（gdn_extend_page_size_1）；
• 整页对齐（exact-page，如零前缀 + 恰好一页extend_lens=(16,)）；
• 跨页（crossing-page）；
• ragged 页边界（ragged page-boundary，如prefix_lens=(0, 8, 16)与extend_lens=(15, 8, 1)的组合）；
• 页大小 32 的跨页（page-size-32 crossing）；
• decode 边界（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_heads=2、num_v_heads=2，head dim 默认 32，dtype 默认torch.bfloat16。

两条与验证相关的补充约定：

• GDN 使用推测式 Mamba 状态缓冲来覆盖 target verify（TARGET_VERIFY）路径——验证阶段的状态索引来自推测解码的中间状态索引构建逻辑；
• split-op 测试用更大的静态 token 缓冲验证 live-token 切片（如静态缓冲 32、实际 token 更少的组合），确保 CUDA Graph 静态形状下切出真实 token 的逻辑正确。

生产环境不支持的组合：_capture_metadata/_replay_metadata的硬约束

矩阵中多个blocked单元格指向同一个底层契约：MambaAttnBackendBase._capture_metadata/_replay_metadata（hybrid_linear_attn_backend.py）只接受两类 forward mode：

1. DECODE_OR_IDLE（forward_mode.is_decode_or_idle()）：decode 或空闲；
2. TARGET_VERIFY（forward_mode.is_target_verify()）：目标模型验证。

任何其他模式都会抛出ValueError(f"Invalid 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_metadata（hybrid_linear_attn_backend.py:246附近）：非 decode、非 extend 的模式同样抛出ValueError。合法模式集合为：

• is_decode_or_idle()：DECODE / IDLE；
• is_extend(include_draft_extend_v2=True)：按 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-2；
• EAGLE tree（topk=2）用例：GDN_TREE_ATOL = 5e-2——矩阵中 Verify CG 列标注的"scoped5e-2atol"即由此而来，原因是 bf16 递推在树形验证布局下会有更大的累积误差，需要更宽松的容忍度；
• FlashInfer 线性 prefill 与 Triton 内核的交叉对比同样使用atol=3e-2, rtol=3e-2。

后续工作方向

README 末尾列出两项待办，可作为 GDN 测试能力演进的参考：

1. 在可用时补充更多线性注意力内核后端变体（当前矩阵线性注意力一侧固定为 Triton GDN 内核，FlashInfer 线性 prefill 已作为专项测试先行覆盖）；
2. 在 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