diff --git a/tests/kernels/mamba/test_gdn_prefill_cutedsl.py b/tests/kernels/mamba/test_gdn_prefill_cutedsl.py index 861b686e32..d86cfde2c7 100644 --- a/tests/kernels/mamba/test_gdn_prefill_cutedsl.py +++ b/tests/kernels/mamba/test_gdn_prefill_cutedsl.py @@ -192,3 +192,86 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype): assert no_buffer_state_error.mean().item() < 6e-4 assert buffer_o_error.max().item() == 0 assert buffer_state_error.max().item() == 0 + + +def test_gdn_chunk_cutedsl_correlated_keys_stability(): + """Cover a KKT block where three BF16 Newton rounds diverge.""" + total_tokens = 64 + num_k_heads = 4 + num_v_heads = 8 + head_dim = 128 + rng = torch.Generator("cuda").manual_seed(3456) + + q = torch.zeros( + 1, + total_tokens, + num_k_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + ) + k = torch.zeros_like(q) + q[..., 0] = 1 + k[..., 0] = 1 + v = torch.randn( + 1, + total_tokens, + num_v_heads, + head_dim, + device="cuda", + dtype=torch.bfloat16, + generator=rng, + ) + g = torch.zeros( + 1, + total_tokens, + num_v_heads, + device="cuda", + dtype=torch.float32, + ) + beta = torch.full_like(g, 0.9) + initial_state = torch.zeros( + 1, + num_v_heads, + head_dim, + head_dim, + device="cuda", + dtype=torch.float32, + ) + cu_seqlens = torch.tensor([0, total_tokens], device="cuda", dtype=torch.int32) + chunk_indices, chunk_offsets = prepare_metadata_cutedsl(cu_seqlens, total_tokens) + + ref_o, ref_state = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + output_final_state=True, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=False, + ) + actual_o, actual_state = chunk_gated_delta_rule_cutedsl( + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + ) + torch.accelerator.synchronize() + + assert torch.isfinite(actual_o).all() + assert torch.isfinite(actual_state).all() + output_error = (actual_o.float() - ref_o.float()).abs() + state_error = ( + actual_state.float() - ref_state.to(actual_state.dtype).float() + ).abs() + assert output_error.max().item() < 1e-1 + assert output_error.mean().item() < 2e-3 + assert state_error.max().item() < 2e-1 + assert state_error.mean().item() < 6e-4 diff --git a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py index ba61ec68b5..4b4046a093 100644 --- a/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py +++ b/vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py @@ -573,8 +573,8 @@ class Sm100ChunkUWKernel: for i in cutlass.range_constexpr(4): M[i] = _bf16x2_sub(_bf16x2_neg(eye[i]), M[i]) - # 3 rounds of Newton-Schulz - for _ in cutlass.range_constexpr(3): + # The fourth round keeps correlated-key KKT blocks stable in BF16. + for _ in cutlass.range_constexpr(4): # First MMA: -AiM = Ai @ (-M) cute.copy(stsm_atom, Ai_bf16, sA_ldsm[None, (warp_id_, warp_id_)]) cute.arch.sync_warp()