[CI] stabilize GDN prefill CuTeDSL test (#49388)
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Codex <noreply@openai.com>
Thien Tran committed
Jul 23, 2026 at 00:03 UTC
b44311b6ef9232d1f345f4b55adef7abc223f0e7
1 file changed
+14
-19
tests/kernels/mamba/test_gdn_prefill_cutedsl.py
+14
-19
@@ -33,12 +33,10 @@ from vllm.third_party.flash_linear_attention.ops.index import ( # noqa: E402
33
@pytest.mark.parametrize("num_seqs", [1, 5, 257])
34
@pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32])
35
def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
36
- seq_lens = torch.randint(
37
- 1,
38
- 130,
39
- (num_seqs,),
40
- dtype=torch.int32,
41
- )
36
+ rng_cpu = torch.Generator("cpu").manual_seed(1234)
37
+ rng = torch.Generator("cuda").manual_seed(2345)
38
+
39
+ seq_lens = torch.randint(1, 130, (num_seqs,), dtype=torch.int32, generator=rng_cpu)
40
cu_seqlens = torch.zeros(num_seqs + 1, device="cuda", dtype=torch.int32)
41
cu_seqlens[1:] = seq_lens.to(device="cuda").cumsum(0)
42
total_tokens = int(cu_seqlens[-1].item())
@@ -56,8 +54,9 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
54
head_k_dim,
55
device="cuda",
56
dtype=dtype,
57
+ generator=rng,
58
)
60
- k = torch.randn_like(q)
59
+ k = torch.randn_like(q, generator=rng)
60
v = torch.randn(
61
1,
62
total_tokens,
@@ -65,29 +64,24 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
64
head_v_dim,
65
device="cuda",
66
dtype=dtype,
67
+ generator=rng,
68
)
69
q = F.normalize(q.float(), p=2, dim=-1).to(dtype)
70
k = F.normalize(k.float(), p=2, dim=-1).to(dtype)
71
a = torch.randn(
72
- 1,
73
- total_tokens,
74
- num_v_heads,
75
- device="cuda",
76
- dtype=dtype,
72
+ 1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng
73
)
74
b = torch.randn(
79
- 1,
80
- total_tokens,
81
- num_v_heads,
82
- device="cuda",
83
- dtype=dtype,
75
+ 1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng
76
)
77
# Match upstream FLA GatedDeltaNet synthetic initialization:
78
# https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/gated_deltanet.py
87
- A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(0, 16)
79
+ A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(
80
+ 0, 16, generator=rng
81
+ )
82
A_log = torch.log(A)
83
dt = torch.exp(
90
- torch.rand(num_v_heads, device="cuda", dtype=torch.float32)
84
+ torch.rand(num_v_heads, device="cuda", dtype=torch.float32, generator=rng)
85
* (math.log(0.1) - math.log(0.001))
86
+ math.log(0.001)
87
)
@@ -105,6 +99,7 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
99
head_k_dim,
100
device="cuda",
101
dtype=state_dtype,
102
+ generator=rng,
103
)
104
* 0.05
105
)