[Bugfix][GDN] Fix CuteDSL BF16 KKT inversion divergence (#53864)

Signed-off-by: Zoe923 <871844693@qq.com>

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

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()