fix(gdn): keep packed decode beta in fp32

tests/unit_tests/test_gdn_packed_decode_patch.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2026 BAAI. All rights reserved.
+
+from types import SimpleNamespace
+
+import vllm_fl
+from vllm_fl.patches import gdn_packed_decode
+
+
+def _vulnerable_kernel():
+    beta_val = "tl.sigmoid(b_val).to(b.dtype.element_ty).to(tl.float32)"
+    return beta_val
+
+
+def _fixed_kernel():
+    beta_val = "tl.sigmoid(b_val)"
+    return beta_val
+
+
+def test_patch_replaces_vulnerable_kernel(monkeypatch):
+    target = SimpleNamespace(
+        fused_recurrent_gated_delta_rule_packed_decode_kernel=_vulnerable_kernel
+    )
+    monkeypatch.setattr(
+        gdn_packed_decode.importlib, "import_module", lambda _module: target
+    )
+
+    assert gdn_packed_decode.patch_vllm_packed_gdn_beta() is True
+    replacement = target.fused_recurrent_gated_delta_rule_packed_decode_kernel
+    assert replacement is not _vulnerable_kernel
+    assert replacement._fl_fp32_beta is True
+    assert gdn_packed_decode.patch_vllm_packed_gdn_beta() is False
+
+
+def test_patch_preserves_already_fixed_upstream_kernel(monkeypatch):
+    target = SimpleNamespace(
+        fused_recurrent_gated_delta_rule_packed_decode_kernel=_fixed_kernel
+    )
+    monkeypatch.setattr(
+        gdn_packed_decode.importlib, "import_module", lambda _module: target
+    )
+
+    assert gdn_packed_decode.patch_vllm_packed_gdn_beta() is False
+    assert target.fused_recurrent_gated_delta_rule_packed_decode_kernel is _fixed_kernel
+
+
+def test_patch_is_optional_when_symbol_is_unavailable(monkeypatch):
+    monkeypatch.setattr(
+        gdn_packed_decode.importlib,
+        "import_module",
+        lambda _module: SimpleNamespace(),
+    )
+
+    assert gdn_packed_decode.patch_vllm_packed_gdn_beta() is False
+
+
+def test_registration_is_optional_when_gdn_module_is_unavailable(monkeypatch):
+    def missing_module(_module):
+        raise ModuleNotFoundError("vendor vLLM image does not provide FLA")
+
+    monkeypatch.setattr(vllm_fl.importlib, "import_module", missing_module)
+
+    assert vllm_fl._register_gdn_packed_decode_patch() is False
+
+
+def test_registration_is_not_vendor_gated(monkeypatch):
+    calls = []
+    patch_module = SimpleNamespace(
+        patch_vllm_packed_gdn_beta=lambda: calls.append("patched") or True
+    )
+    monkeypatch.setattr(
+        vllm_fl.importlib,
+        "import_module",
+        lambda module: patch_module,
+    )
+
+    assert vllm_fl._register_gdn_packed_decode_patch() is True
+    assert calls == ["patched"]

vllm_fl/__init__.py

@@ -136,6 +136,25 @@ def register_router():
     from vllm_fl.ops.fused_moe.router import replace_router_with_fl
     replace_router_with_fl()
 
+
+def _register_gdn_packed_decode_patch() -> bool:
+    """Install the packed GDN fix when this vLLM build provides it.
+
+    Vendor images may omit vLLM's FLA package or route GDN through a different
+    implementation. Keep the compatibility hook capability-based: any build
+    carrying the vulnerable kernel is patched, while builds without the
+    required module or symbol remain untouched.
+    """
+    try:
+        patch_module = importlib.import_module("vllm_fl.patches.gdn_packed_decode")
+        patch_fn = patch_module.patch_vllm_packed_gdn_beta
+    except (ImportError, AttributeError) as exc:
+        logger.debug("Packed GDN decode patch is unavailable: %s", exc)
+        return False
+
+    return patch_fn()
+
+
 def register_model():
     """Register FL-specific models not yet upstream."""
     # General plugins are loaded independently in spawned model-inspection and
@@ -152,6 +171,8 @@ def register_model():
     register_quant_linear()
     register_router()
 
+    _register_gdn_packed_decode_patch()
+
     # Register GLM-5 (GlmMoeDsa) — config not yet upstream
     try:
         from vllm.transformers_utils.config import _CONFIG_REGISTRY

vllm_fl/patches/gdn_packed_decode.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# The Triton kernel below is adapted from flash-linear-attention via vLLM.
+# The original source was distributed under the MIT license.
+"""Numerical compatibility patch for vLLM's packed GDN decode kernel."""
+
+from __future__ import annotations
+
+import importlib
+import inspect
+import logging
+
+from vllm.model_executor.layers.fla.ops.op import exp
+from vllm.triton_utils import tl, triton
+
+logger = logging.getLogger(__name__)
+
+_TARGET_MODULE = "vllm.model_executor.layers.fla.ops.fused_recurrent"
+_TARGET_NAME = "fused_recurrent_gated_delta_rule_packed_decode_kernel"
+_VULNERABLE_BETA_EXPRESSION = (
+    "tl.sigmoid(b_val).to(b.dtype.element_ty).to(tl.float32)"
+)
+
+
+@triton.jit
+def _fused_recurrent_gated_delta_rule_packed_decode_kernel_fp32_beta(
+    mixed_qkv,
+    a,
+    b,
+    A_log,
+    dt_bias,
+    o,
+    h0,
+    ht,
+    ssm_state_indices,
+    scale,
+    stride_mixed_qkv_tok: tl.constexpr,
+    stride_a_tok: tl.constexpr,
+    stride_b_tok: tl.constexpr,
+    stride_init_state_token: tl.constexpr,
+    stride_final_state_token: tl.constexpr,
+    stride_indices_seq: tl.constexpr,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    SOFTPLUS_THRESHOLD: tl.constexpr,
+    USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
+):
+    i_v, i_nh = tl.program_id(0), tl.program_id(1)
+    i_n, i_hv = i_nh // HV, i_nh % HV
+    i_h = i_hv // (HV // H)
+
+    o_k = tl.arange(0, BK)
+    o_v = i_v * BV + tl.arange(0, BV)
+    mask_k = o_k < K
+    mask_v = o_v < V
+    mask_h = mask_v[:, None] & mask_k[None, :]
+
+    state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq).to(tl.int64)
+    p_o = o + (i_n * HV + i_hv) * V + o_v
+
+    if state_idx <= 0:
+        zero = tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty)
+        tl.store(p_o, zero, mask=mask_v)
+        return
+
+    p_h0 = h0 + state_idx * stride_init_state_token
+    p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
+    b_h = tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
+
+    p_mixed = mixed_qkv + i_n * stride_mixed_qkv_tok
+    q_off = i_h * K + o_k
+    k_off = (H * K) + i_h * K + o_k
+    v_off = (2 * H * K) + i_hv * V + o_v
+    b_q = tl.load(p_mixed + q_off, mask=mask_k, other=0).to(tl.float32)
+    b_k = tl.load(p_mixed + k_off, mask=mask_k, other=0).to(tl.float32)
+    b_v = tl.load(p_mixed + v_off, mask=mask_v, other=0).to(tl.float32)
+
+    if USE_QK_L2NORM_IN_KERNEL:
+        b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
+        b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
+    b_q = b_q * scale
+
+    a_val = tl.load(a + i_n * stride_a_tok + i_hv).to(tl.float32)
+    b_val = tl.load(b + i_n * stride_b_tok + i_hv).to(tl.float32)
+    A_log_val = tl.load(A_log + i_hv).to(tl.float32)
+    dt_bias_val = tl.load(dt_bias + i_hv).to(tl.float32)
+    x = a_val + dt_bias_val
+    softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x)
+    g_val = -tl.exp(A_log_val) * softplus_x
+
+    # Keep sigmoid(b) in FP32. Rounding it to the input dtype perturbs every
+    # recurrent state update and compounds across decode steps.
+    beta_val = tl.sigmoid(b_val)
+
+    b_h *= exp(g_val)
+    b_v -= tl.sum(b_h * b_k[None, :], 1)
+    b_v *= beta_val
+    b_h += b_v[:, None] * b_k[None, :]
+    b_o = tl.sum(b_h * b_q[None, :], 1)
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
+
+    p_ht = ht + state_idx * stride_final_state_token
+    p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
+    tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
+
+
+_fused_recurrent_gated_delta_rule_packed_decode_kernel_fp32_beta._fl_fp32_beta = True
+
+
+def _kernel_needs_beta_patch(kernel) -> bool:
+    if getattr(kernel, "_fl_fp32_beta", False):
+        return False
+
+    python_fn = getattr(kernel, "fn", kernel)
+    try:
+        source = inspect.getsource(python_fn)
+    except (OSError, TypeError):
+        # vLLM 0.24.0 is known to need the fix. If source inspection is not
+        # available (for example in a stripped wheel), prefer the corrected
+        # implementation over silently retaining the precision bug.
+        return True
+    return _VULNERABLE_BETA_EXPRESSION in source
+
+
+def patch_vllm_packed_gdn_beta() -> bool:
+    """Replace the vulnerable vLLM packed GDN kernel with the FP32-beta one.
+
+    Returns ``True`` only when the replacement is applied. The symbol/source
+    checks make the hook idempotent and allow it to no-op once vLLM contains an
+    equivalent upstream fix.
+    """
+    try:
+        target_module = importlib.import_module(_TARGET_MODULE)
+        current_kernel = getattr(target_module, _TARGET_NAME)
+    except (ImportError, AttributeError) as exc:
+        logger.debug("Packed GDN decode kernel is unavailable: %s", exc)
+        return False
+
+    if not _kernel_needs_beta_patch(current_kernel):
+        return False
+
+    setattr(
+        target_module,
+        _TARGET_NAME,
+        _fused_recurrent_gated_delta_rule_packed_decode_kernel_fp32_beta,
+    )
+    logger.info("Patched vLLM packed GDN decode to keep beta in FP32")
+    return True