From 25dd4897d4378bca33e56a2f332723d599ec3e57 Mon Sep 17 00:00:00 2001 From: halt95 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