179 lines
7.9 KiB
Diff
179 lines
7.9 KiB
Diff
From 25dd4897d4378bca33e56a2f332723d599ec3e57 Mon Sep 17 00:00:00 2001
|
|
From: halt95 <halt95@users.noreply.github.com>
|
|
Date: Tue, 22 Sep 2026 04:13:07 -0400
|
|
Subject: [PATCH 010/103] gdn: bound the decode state index
|
|
|
|
Ports the check from upstream vLLM PR #50021.
|
|
---
|
|
.../test_fused_sigmoid_gating_delta_rule.py | 81 +++++++++++++++++++
|
|
.../ops/fused_sigmoid_gating.py | 42 ++++++++--
|
|
2 files changed, 118 insertions(+), 5 deletions(-)
|
|
|
|
diff --git a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py
|
|
index fc923cb..7d2e399 100644
|
|
--- a/tests/kernels/test_fused_sigmoid_gating_delta_rule.py
|
|
+++ b/tests/kernels/test_fused_sigmoid_gating_delta_rule.py
|
|
@@ -197,3 +197,84 @@ def test_fused_sigmoid_gating_delta_rule_update_spec(
|
|
torch.testing.assert_close(
|
|
last_recurrent_state, last_recurrent_state_ref, atol=1e-2, rtol=1e-2
|
|
)
|
|
+
|
|
+
|
|
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
|
+@pytest.mark.parametrize("num_accepted", [0, 5])
|
|
+def test_spec_invalid_accepted_count_zeroes_output(num_accepted: int) -> None:
|
|
+ """Invalid active-request state selection must not expose empty output.
|
|
+
|
|
+ Ported from upstream vLLM PR #50021, covering BOTH kernels: the decode guard
|
|
+ (fused_sigmoid_gating) and the recurrent guard (fused_recurrent). A prior version dropped the
|
|
+ recurrent arm because that kernel was unguarded; it now has the
|
|
+ guard rather than the caveat, so the arm is back.
|
|
+
|
|
+ num_accepted 0 gives i_t == -1, a read before this request's row;
|
|
+ num_accepted 5 exceeds the one-column ssm_state_indices, a read past it.
|
|
+ ``fill_uninitialized_memory`` makes the old behaviour visible: the kernel
|
|
+ returned early without writing ``o``, which is ``q.new_empty``.
|
|
+ """
|
|
+ torch.set_default_device(DEVICE)
|
|
+ set_random_seed(0)
|
|
+
|
|
+ num_tokens, num_k_heads, num_v_heads = 4, 2, 4
|
|
+ head_k_dim = head_v_dim = 16
|
|
+ query = torch.rand(1, num_tokens, num_k_heads, head_k_dim)
|
|
+ key = torch.rand_like(query)
|
|
+ value = torch.rand(1, num_tokens, num_v_heads, head_v_dim)
|
|
+ a = torch.rand(num_tokens, num_v_heads)
|
|
+ b = torch.rand_like(a)
|
|
+ A_log = torch.rand(num_v_heads)
|
|
+ dt_bias = torch.rand_like(A_log)
|
|
+ initial_state = torch.rand(num_tokens + 1, num_v_heads, head_v_dim, head_k_dim)
|
|
+ state_indices = torch.arange(1, num_tokens + 1, dtype=torch.int32).view(
|
|
+ 1, num_tokens
|
|
+ )
|
|
+ cu_seqlens = torch.tensor([0, num_tokens], dtype=torch.int32)
|
|
+ accepted = torch.tensor([num_accepted], dtype=torch.int32)
|
|
+
|
|
+ previous_deterministic = torch.are_deterministic_algorithms_enabled()
|
|
+ previous_warn_only = torch.is_deterministic_algorithms_warn_only_enabled()
|
|
+ previous_fill = torch.utils.deterministic.fill_uninitialized_memory
|
|
+ torch.use_deterministic_algorithms(True)
|
|
+ torch.utils.deterministic.fill_uninitialized_memory = True
|
|
+ try:
|
|
+ recurrent_state = initial_state.clone()
|
|
+ recurrent_out, recurrent_state = fused_recurrent_gated_delta_rule(
|
|
+ q=query,
|
|
+ k=key,
|
|
+ v=value,
|
|
+ g=torch.rand(1, num_tokens, num_v_heads),
|
|
+ beta=torch.rand(1, num_tokens, num_v_heads),
|
|
+ initial_state=recurrent_state,
|
|
+ inplace_final_state=True,
|
|
+ ssm_state_indices=state_indices,
|
|
+ cu_seqlens=cu_seqlens,
|
|
+ num_accepted_tokens=accepted,
|
|
+ )
|
|
+
|
|
+ sigmoid_state = initial_state.clone()
|
|
+ sigmoid_out, sigmoid_state = fused_sigmoid_gating_delta_rule_update(
|
|
+ A_log=A_log,
|
|
+ a=a,
|
|
+ b=b,
|
|
+ dt_bias=dt_bias,
|
|
+ q=query,
|
|
+ k=key,
|
|
+ v=value,
|
|
+ initial_state=sigmoid_state,
|
|
+ inplace_final_state=True,
|
|
+ ssm_state_indices=state_indices,
|
|
+ cu_seqlens=cu_seqlens,
|
|
+ num_accepted_tokens=accepted,
|
|
+ )
|
|
+ finally:
|
|
+ torch.utils.deterministic.fill_uninitialized_memory = previous_fill
|
|
+ torch.use_deterministic_algorithms(
|
|
+ previous_deterministic, warn_only=previous_warn_only
|
|
+ )
|
|
+
|
|
+ torch.testing.assert_close(recurrent_out, torch.zeros_like(recurrent_out))
|
|
+ torch.testing.assert_close(sigmoid_out, torch.zeros_like(sigmoid_out))
|
|
+ torch.testing.assert_close(recurrent_state, initial_state)
|
|
+ torch.testing.assert_close(sigmoid_state, initial_state)
|
|
diff --git a/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py b/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py
|
|
index 7e0c7e0..c5e75ec 100644
|
|
--- a/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py
|
|
+++ b/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py
|
|
@@ -10,6 +10,14 @@
|
|
import torch
|
|
|
|
from vllm.triton_utils import tl, triton
|
|
+from vllm.v1.worker import guard_telemetry
|
|
+
|
|
+# Bias so a negative row/index survives atomic_max on a zeroed counter
|
|
+# -- it MUST be a tl.constexpr: a plain module-level int is not readable
|
|
+# from a @triton.jit body -- an earlier build shipped one and every guarded kernel died
|
|
+# with "NameError: Cannot access global variable _GUARD_BIAS" on the smoke.
|
|
+_GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS)
|
|
+
|
|
|
|
|
|
@triton.heuristics(
|
|
@@ -37,6 +45,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
|
|
cu_seqlens,
|
|
ssm_state_indices,
|
|
num_accepted_tokens,
|
|
+ guard_ctr, # bounds-guard instrumentation: persistent int64[3]
|
|
scale,
|
|
N: tl.int64, # num of sequences
|
|
T: tl.int64, # num of tokens
|
|
@@ -106,12 +115,32 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
|
|
i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1
|
|
else:
|
|
i_t = 0
|
|
- # Load state index and check for invalid entries
|
|
- state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to(
|
|
- tl.int64
|
|
- )
|
|
- # Skip if state index is invalid (NULL_BLOCK_ID=0)
|
|
+ # PR #50021: i_t is num_accepted_tokens - 1 and is otherwise
|
|
+ # unbounded against a stride_indices_seq-column tensor. A zero
|
|
+ # accepted count gives i_t == -1, a read before this request's row
|
|
+ # (and before the tensor for i_n == 0); a stale or too-large count
|
|
+ # reads past the row. The state_idx <= 0 guard below rejects only
|
|
+ # non-positive values, so an out-of-range read returning a garbage
|
|
+ # positive int flowed into h0 + state_idx * stride_init_state_token
|
|
+ # and faulted the SM.
|
|
+ idx_in_row = (i_t >= 0) & (i_t < stride_indices_seq)
|
|
+ if (i_t < 0) or (i_t >= stride_indices_seq):
|
|
+ tl.atomic_add(guard_ctr, 1)
|
|
+ tl.atomic_max(guard_ctr + 1, (i_n).to(tl.int64) + _GUARD_BIAS)
|
|
+ tl.atomic_max(guard_ctr + 2, (i_t).to(tl.int64) + _GUARD_BIAS)
|
|
+ state_idx = tl.load(
|
|
+ ssm_state_indices + i_n * stride_indices_seq + i_t,
|
|
+ mask=idx_in_row,
|
|
+ other=0,
|
|
+ ).to(tl.int64)
|
|
+ # Skip invalid state indices without exposing uninitialized output.
|
|
+ # o is q.new_empty in the entry below, so the previous early return
|
|
+ # left this request's rows holding whatever the allocator returned.
|
|
if state_idx <= 0:
|
|
+ zero = tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty)
|
|
+ for _ in range(0, T):
|
|
+ tl.store(p_o, zero, mask=mask_v)
|
|
+ p_o += HV * V
|
|
return
|
|
p_h0 = h0 + state_idx * stride_init_state_token
|
|
else:
|
|
@@ -255,6 +284,9 @@ def fused_sigmoid_gating_delta_rule_update(
|
|
cu_seqlens=cu_seqlens,
|
|
ssm_state_indices=ssm_state_indices,
|
|
num_accepted_tokens=num_accepted_tokens,
|
|
+ guard_ctr=guard_telemetry.counter(
|
|
+ guard_telemetry.SITE_GDN_DECODE_STATE, q.device
|
|
+ ),
|
|
scale=scale,
|
|
N=N,
|
|
T=T,
|
|
--
|
|
2.43.0
|
|
|