From cda271d5ea26ff3c03c5d202041f9a09552c5ef3 Mon Sep 17 00:00:00 2001 From: halt95 Date: Tue, 22 Sep 2026 04:13:09 -0400 Subject: [PATCH 016/103] gdn: bound the recurrent state index Both callers of the guarded kernel (fused recurrent and KDA) are updated in the same commit. --- .../worker/gpu/test_guard_telemetry_replay.py | 167 ++++++++++++++++++ .../ops/fused_recurrent.py | 44 ++++- .../flash_linear_attention/ops/kda.py | 7 + 3 files changed, 213 insertions(+), 5 deletions(-) create mode 100644 tests/v1/worker/gpu/test_guard_telemetry_replay.py diff --git a/tests/v1/worker/gpu/test_guard_telemetry_replay.py b/tests/v1/worker/gpu/test_guard_telemetry_replay.py new file mode 100644 index 0000000..52327ab --- /dev/null +++ b/tests/v1/worker/gpu/test_guard_telemetry_replay.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The test that distinguishes "architecturally fixed" from "fixed". + +The first bounds-guard instrumentation design allocated a counter per launch and read +it back inline. Under CUDA-graph REPLAY the Python host code never runs, so +that counter was re-zeroed every step and never read: the instrument was +silently dead in exactly the configuration the crash happens in (QSA declares +AttentionCGSupport.UNIFORM_BATCH, so the guarded kernels sit inside the +captured forward). + +The current design rebuilds it on a persistent buffer plus captured atomics. That is the +right architecture, but "right architecture" is a claim about code. This is the +measurement: capture a graph around a guarded kernel with a seeded bad value, +replay it N times WITHOUT re-entering Python, and require the counter to have +advanced N times. + +A test that only ran the kernel eagerly would pass on the first design too, +which is why it must be a replay. +""" + +import pytest +import torch + +from vllm.third_party.flash_linear_attention.ops import ( + fused_sigmoid_gating_delta_rule_update, +) +from vllm.v1.worker import guard_telemetry + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA required" +) + +REPLAYS = 5 + + +def _seeded_call(device): + """One guarded launch whose guard MUST fire. + + num_accepted_tokens = 0 makes i_t == -1 in + fused_sigmoid_gating_delta_rule_update_kernel, a read before this request's + row of ssm_state_indices. That is the decode-state guard. + """ + num_tokens, num_k_heads, num_v_heads = 4, 2, 4 + head_k_dim = head_v_dim = 16 + q = torch.rand(1, num_tokens, num_k_heads, head_k_dim, device=device) + k = torch.rand_like(q) + v = torch.rand(1, num_tokens, num_v_heads, head_v_dim, device=device) + a = torch.rand(num_tokens, num_v_heads, device=device) + b = torch.rand_like(a) + A_log = torch.rand(num_v_heads, device=device) + dt_bias = torch.rand_like(A_log) + state = torch.rand( + num_tokens + 1, num_v_heads, head_v_dim, head_k_dim, device=device + ) + state_indices = torch.arange( + 1, num_tokens + 1, dtype=torch.int32, device=device + ).view(1, num_tokens) + cu_seqlens = torch.tensor([0, num_tokens], dtype=torch.int32, device=device) + accepted = torch.zeros(1, dtype=torch.int32, device=device) # -> i_t == -1 + + def run(): + fused_sigmoid_gating_delta_rule_update( + A_log=A_log, + a=a, + b=b, + dt_bias=dt_bias, + q=q, + k=k, + v=v, + initial_state=state, + inplace_final_state=True, + ssm_state_indices=state_indices, + cu_seqlens=cu_seqlens, + num_accepted_tokens=accepted, + ) + + return run + + +def _count(device): + ctr = guard_telemetry.counter(guard_telemetry.SITE_GDN_DECODE_STATE, device) + torch.cuda.synchronize() + return int(ctr[0].item()) + + +def test_guard_counts_under_cudagraph_replay(): + device = torch.device("cuda", torch.cuda.current_device()) + # Exactly what the model runner does, and for the same reason: allocate the + # counter before anything can capture. + guard_telemetry.prealloc(device) + guard_telemetry.reset_for_test() + + run = _seeded_call(device) + + # Warm up on a side stream, the standard pre-capture dance. This also + # compiles the kernel: a JIT compile inside a capture would abort it. + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + + warm = _count(device) + assert warm > 0, ( + "the seeded fault did not fire the guard even eagerly -- the control is " + "broken, so the replay result below would mean nothing" + ) + per_launch = warm // 3 + assert per_launch > 0 + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + # Capturing itself must not have aborted, and must not have counted. + torch.cuda.synchronize() + before = _count(device) + + for _ in range(REPLAYS): + graph.replay() + torch.cuda.synchronize() + after = _count(device) + + delta = after - before + assert delta == REPLAYS * per_launch, ( + f"counter advanced by {delta} over {REPLAYS} replays, expected " + f"{REPLAYS * per_launch}. A delta of 0 means the guard telemetry is " + f"dead under replay -- the replay defect is NOT fixed." + ) + + +def test_counter_pointer_is_stable_across_replays(): + """The captured atomic addresses one fixed buffer. + + If counter() ever returned a fresh tensor, the captured graph would keep + writing to a freed address -- the other half of the replay defect, and the one that + corrupts memory rather than merely losing evidence. + """ + device = torch.device("cuda", torch.cuda.current_device()) + guard_telemetry.prealloc(device) + first = guard_telemetry.counter(guard_telemetry.SITE_GDN_DECODE_STATE, device) + second = guard_telemetry.counter(guard_telemetry.SITE_GDN_DECODE_STATE, device) + assert first.data_ptr() == second.data_ptr() + # and the unindexed alias must resolve to the same buffer + third = guard_telemetry.counter( + guard_telemetry.SITE_GDN_DECODE_STATE, torch.device("cuda") + ) + assert third.data_ptr() == first.data_ptr() + + # ...and it must survive the reader giving up on it. An + # earlier failure path popped the buffer, which drops the only durable + # reference: counter() hands out transient views, and under replay nothing + # holds one at all, so the storage would return to the allocator, be + # reused, and the CAPTURED atomics would keep writing to the baked address. + # Stop polling, never stop holding. + try: + guard_telemetry._read_disabled.add(device) + after = guard_telemetry.counter(guard_telemetry.SITE_GDN_DECODE_STATE, device) + assert after.data_ptr() == first.data_ptr(), ( + "the counter moved after the reader disabled it -- captured graphs " + "are now writing to freed storage" + ) + assert device in guard_telemetry._buffers + finally: + guard_telemetry._read_disabled.discard(device) diff --git a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py index eb08b93..2379fcb 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_recurrent.py @@ -11,6 +11,12 @@ 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. +_GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) from .op import exp, log @@ -36,6 +42,7 @@ def fused_recurrent_gated_delta_rule_fwd_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 +113,36 @@ def fused_recurrent_gated_delta_rule_fwd_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, while ssm_state_indices has only stride_indices_seq + # columns per request. 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. + # + # This is the SAME defect as the decode-state guard, in the twin kernel. An + # earlier revision left it unguarded and dropped the upstream test's arm for it; + # it is now ported rather than caveated. + 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: @@ -227,6 +258,9 @@ def fused_recurrent_gated_delta_rule_fwd( 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_RECURRENT_STATE, q.device + ), scale=scale, N=N, T=T, diff --git a/vllm/third_party/flash_linear_attention/ops/kda.py b/vllm/third_party/flash_linear_attention/ops/kda.py index 548f600..5ba7213 100644 --- a/vllm/third_party/flash_linear_attention/ops/kda.py +++ b/vllm/third_party/flash_linear_attention/ops/kda.py @@ -14,6 +14,7 @@ import torch.nn as nn from vllm.model_executor.custom_op import CustomOp from vllm.triton_utils import tl, triton +from vllm.v1.worker import guard_telemetry from vllm.utils.math_utils import RCP_LN2, cdiv, next_power_of_2 from .chunk_delta_h import chunk_gated_delta_rule_fwd_h @@ -81,6 +82,12 @@ def fused_recurrent_kda_fwd( cu_seqlens=cu_seqlens, ssm_state_indices=ssm_state_indices, num_accepted_tokens=num_accepted_tokens, + # Bounds-guard instrumentation: the same kernel object as fused_recurrent.py:249, so it + # takes the same guard counter. Same site: the defect being guarded is + # the kernel's, not the caller's. + guard_ctr=guard_telemetry.counter( + guard_telemetry.SITE_GDN_RECURRENT_STATE, q.device + ), scale=scale, N=N, T=T, -- 2.43.0