iHore/vLLM-4x-RTX-3090
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
777e910dd0e54ffe8f95acd534c76c56e7320fe2
vLLM-4x-RTX-3090/patches/0016-gdn-bound-the-recurrent-state-index.patch
T
MOVZX 777e910dd0 Initial commit 
Signed-off-by: MOVZX <movzx@yahoo.com>
2026-10-15 01:01:52 +07:00

293 lines
12 KiB
Diff
Raw Blame History

From cda271d5ea26ff3c03c5d202041f9a09552c5ef3 Mon Sep 17 00:00:00 2001
From: halt95 <halt95@users.noreply.github.com>
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
Reference in New Issue View Git Blame Copy Permalink