From b44311b6ef9232d1f345f4b55adef7abc223f0e7 Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Thu, 23 Jul 2026 00:03:46 +0800 Subject: [PATCH] [CI] stabilize GDN prefill CuTeDSL test (#49388) Signed-off-by: Thien Tran Co-authored-by: Codex --- .../kernels/mamba/test_gdn_prefill_cutedsl.py | 33 ++++++++----------- 1 file changed, 14 insertions(+), 19 deletions(-) diff --git a/tests/kernels/mamba/test_gdn_prefill_cutedsl.py b/tests/kernels/mamba/test_gdn_prefill_cutedsl.py index 1f5a24fd81..861b686e32 100644 --- a/tests/kernels/mamba/test_gdn_prefill_cutedsl.py +++ b/tests/kernels/mamba/test_gdn_prefill_cutedsl.py @@ -33,12 +33,10 @@ from vllm.third_party.flash_linear_attention.ops.index import ( # noqa: E402 @pytest.mark.parametrize("num_seqs", [1, 5, 257]) @pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype): - seq_lens = torch.randint( - 1, - 130, - (num_seqs,), - dtype=torch.int32, - ) + rng_cpu = torch.Generator("cpu").manual_seed(1234) + rng = torch.Generator("cuda").manual_seed(2345) + + seq_lens = torch.randint(1, 130, (num_seqs,), dtype=torch.int32, generator=rng_cpu) cu_seqlens = torch.zeros(num_seqs + 1, device="cuda", dtype=torch.int32) cu_seqlens[1:] = seq_lens.to(device="cuda").cumsum(0) total_tokens = int(cu_seqlens[-1].item()) @@ -56,8 +54,9 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype): head_k_dim, device="cuda", dtype=dtype, + generator=rng, ) - k = torch.randn_like(q) + k = torch.randn_like(q, generator=rng) v = torch.randn( 1, total_tokens, @@ -65,29 +64,24 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype): head_v_dim, device="cuda", dtype=dtype, + generator=rng, ) q = F.normalize(q.float(), p=2, dim=-1).to(dtype) k = F.normalize(k.float(), p=2, dim=-1).to(dtype) a = torch.randn( - 1, - total_tokens, - num_v_heads, - device="cuda", - dtype=dtype, + 1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng ) b = torch.randn( - 1, - total_tokens, - num_v_heads, - device="cuda", - dtype=dtype, + 1, total_tokens, num_v_heads, device="cuda", dtype=dtype, generator=rng ) # Match upstream FLA GatedDeltaNet synthetic initialization: # https://github.com/fla-org/flash-linear-attention/blob/main/fla/layers/gated_deltanet.py - A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_(0, 16) + A = torch.empty(num_v_heads, device="cuda", dtype=torch.float32).uniform_( + 0, 16, generator=rng + ) A_log = torch.log(A) dt = torch.exp( - torch.rand(num_v_heads, device="cuda", dtype=torch.float32) + torch.rand(num_v_heads, device="cuda", dtype=torch.float32, generator=rng) * (math.log(0.1) - math.log(0.001)) + math.log(0.001) ) @@ -105,6 +99,7 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype): head_k_dim, device="cuda", dtype=state_dtype, + generator=rng, ) * 0.05 )