[Bugfix][GDN] Fix CuteDSL BF16 KKT inversion divergence (#53864)
Signed-off-by: Zoe923 <871844693@qq.com>
Zoe923 committed
Sep 17, 2026 at 23:49 UTC
6a7de3ba3ad95965e41618568ff803fa01ab168d
2 files changed
+85
-2
tests/kernels/mamba/test_gdn_prefill_cutedsl.py
+83
@@ -192,3 +192,86 @@ def test_gdn_chunk_cutedsl_correctness(num_seqs: int, state_dtype: torch.dtype):
192
assert no_buffer_state_error.mean().item() < 6e-4
193
assert buffer_o_error.max().item() == 0
194
assert buffer_state_error.max().item() == 0
195
+
196
+
197
+def test_gdn_chunk_cutedsl_correlated_keys_stability():
198
+ """Cover a KKT block where three BF16 Newton rounds diverge."""
199
+ total_tokens = 64
200
+ num_k_heads = 4
201
+ num_v_heads = 8
202
+ head_dim = 128
203
+ rng = torch.Generator("cuda").manual_seed(3456)
204
+
205
+ q = torch.zeros(
206
+ 1,
207
+ total_tokens,
208
+ num_k_heads,
209
+ head_dim,
210
+ device="cuda",
211
+ dtype=torch.bfloat16,
212
+ )
213
+ k = torch.zeros_like(q)
214
+ q[..., 0] = 1
215
+ k[..., 0] = 1
216
+ v = torch.randn(
217
+ 1,
218
+ total_tokens,
219
+ num_v_heads,
220
+ head_dim,
221
+ device="cuda",
222
+ dtype=torch.bfloat16,
223
+ generator=rng,
224
+ )
225
+ g = torch.zeros(
226
+ 1,
227
+ total_tokens,
228
+ num_v_heads,
229
+ device="cuda",
230
+ dtype=torch.float32,
231
+ )
232
+ beta = torch.full_like(g, 0.9)
233
+ initial_state = torch.zeros(
234
+ 1,
235
+ num_v_heads,
236
+ head_dim,
237
+ head_dim,
238
+ device="cuda",
239
+ dtype=torch.float32,
240
+ )
241
+ cu_seqlens = torch.tensor([0, total_tokens], device="cuda", dtype=torch.int32)
242
+ chunk_indices, chunk_offsets = prepare_metadata_cutedsl(cu_seqlens, total_tokens)
243
+
244
+ ref_o, ref_state = chunk_gated_delta_rule(
245
+ q=q,
246
+ k=k,
247
+ v=v,
248
+ g=g,
249
+ beta=beta,
250
+ initial_state=initial_state,
251
+ output_final_state=True,
252
+ cu_seqlens=cu_seqlens,
253
+ use_qk_l2norm_in_kernel=False,
254
+ )
255
+ actual_o, actual_state = chunk_gated_delta_rule_cutedsl(
256
+ q=q,
257
+ k=k,
258
+ v=v,
259
+ g=g,
260
+ beta=beta,
261
+ initial_state=initial_state,
262
+ cu_seqlens=cu_seqlens,
263
+ chunk_indices=chunk_indices,
264
+ chunk_offsets=chunk_offsets,
265
+ )
266
+ torch.accelerator.synchronize()
267
+
268
+ assert torch.isfinite(actual_o).all()
269
+ assert torch.isfinite(actual_state).all()
270
+ output_error = (actual_o.float() - ref_o.float()).abs()
271
+ state_error = (
272
+ actual_state.float() - ref_state.to(actual_state.dtype).float()
273
+ ).abs()
274
+ assert output_error.max().item() < 1e-1
275
+ assert output_error.mean().item() < 2e-3
276
+ assert state_error.max().item() < 2e-1
277
+ assert state_error.mean().item() < 6e-4
vllm/model_executor/layers/mamba/ops/gdn_chunk_cutedsl/kernel_kkt_inv_uw.py
+2
-2
@@ -573,8 +573,8 @@ class Sm100ChunkUWKernel:
573
for i in cutlass.range_constexpr(4):
574
M[i] = _bf16x2_sub(_bf16x2_neg(eye[i]), M[i])
575
576
- # 3 rounds of Newton-Schulz
577
- for _ in cutlass.range_constexpr(3):
576
+ # The fourth round keeps correlated-key KKT blocks stable in BF16.
577
+ for _ in cutlass.range_constexpr(4):
578
# First MMA: -AiM = Ai @ (-M)
579
cute.copy(stsm_atom, Ai_bf16, sA_ldsm[None, (warp_id_, warp_id_)])
580
cute.arch.sync_warp()