Purpose

End-to-end impact

This is an end-to-end correctness bug, not only a synthetic numerical issue.

On Qwen3.5-397B-A17B-FP8, the current three-round CuteDSL GDN prefill kernel can cause deterministic greedy generation to diverge catastrophically from the FLA backend. In one affected real business request with a 9,991-token prompt, the three-round implementation returned an empty output under an otherwise identical serving configuration.

Changing only the Newton-Schulz correction count from three to four restored the affected request to the expected FLA output. After the fix:

• 200/200 responses across client concurrency 1, 2, 4, and 8 matched their corresponding FLA responses character-for-character.
• 2,000/2,000 responses across client concurrency 16 through 200 matched their corresponding FLA responses character-for-character.
• No empty output, malformed text, Unicode replacement character, timeout, OOM, or kernel error was observed.

Therefore, the fourth Newton-Schulz round fixes a model-level correctness failure while preserving the performance advantage of the CuteDSL backend.

Here, "exact" refers to generated response strings under deterministic greedy decoding. The intermediate GDN output and state tensors are numerically close to FLA, but are not bitwise identical.

Numerical root cause and fix

The CuteDSL GDN prefill kernel currently computes the inverse of each 16x16 lower-triangular KKT diagonal block using three Newton-Schulz correction rounds with BF16 intermediates.

Three rounds are sufficient for this block size in exact arithmetic. However, they are not always sufficient when the intermediate inverse is stored back in BF16 after every round. Highly correlated keys can produce non-normal KKT blocks with transient growth. In these cases, the third round may finish before the BF16 cancellation has converged, causing a large inverse error that propagates into the GDN state and output.

This was observed in both synthetic and real-model inputs:

• A deterministic correlated-key case produced an output maximum absolute error of 752.116699 with the three-round implementation.
• Across 131 real Qwen3.5 chunks and 32 heads, the worst KKT inverse absolute error with three rounds was 98.5014.
• The same failure was reproduced with both CUTLASS DSL 4.5.2 and 4.6.2, so it is not specific to one CUTLASS version or to an internal port.

This PR adds one Newton-Schulz correction round, changing the iteration count from three to four.

The fourth round does not change the GDN mathematical definition, tensor layout, metadata, or public API. It only improves convergence of the existing BF16 inverse approximation.

A deterministic regression test is also added. It constructs a 64-token case with highly correlated normalized keys, beta=0.9, g=0, and a zero FP32 initial state. This input reliably exposes the three-round failure and verifies that the output and final state remain finite and within the expected error bounds.

Four rounds were selected as the minimum sufficient fix:

The fifth round further reduces the maximum error, but provides almost no improvement in the mean error. Four rounds therefore provide the best correctness/performance trade-off observed in the tested inputs.

Performance impact

The additional correction round adds approximately 0.38% average latency to the local CuteDSL GDN prefill kernel compared with the three-round implementation.

Despite this small cost, the corrected four-round CuteDSL kernel remains 2.09x to 3.20x faster than the FLA reference across the tested Qwen3.5 GDN shapes.

End-to-end Qwen3.5-397B-A17B-FP8 testing also retained a positive performance improvement over FLA:

The backend was also tested with 200 requests at each client concurrency from 16 to 200. Both sides used TP2, MTP disabled, a 10k-token scheduler budget, max_num_seqs=220, unique cache salts, and zero prefix-cache hits.

Across these ten points, total throughput improved by 8.49% on average and mean latency improved by 8.10% on average. C16 mean TTFT regressed by 1.01%, which is within the observed run-to-run variance, while its mean latency still improved by 8.23%. Every concurrency point completed 200/200 requests successfully, and all 2,000 CuteDSL responses matched their FLA counterparts character-for-character.

The three-round implementation was not used as an end-to-end performance baseline because it produced incorrect model outputs. Its incremental cost was measured at the kernel level instead.

Test Plan

GPU regression tests

Run the existing parameterized CuteDSL/FLA comparisons together with the new correlated-key stability regression:

pytest -q tests/kernels/mamba/test_gdn_prefill_cutedsl.py

The tests cover:

• num_seqs = 1, 5, 257
• BF16 and FP32 initial states
• Output and final-state comparisons against the FLA reference
• A deterministic correlated-key KKT block that diverges with three Newton-Schulz rounds
• Finite-value checks and explicit output/state error thresholds

Real hidden-state replay

Replay captured GDN inputs from all 45 GDN layers of Qwen3.5-397B-A17B-FP8 on both TP ranks:

• 45 layers x 2 TP ranks = 90 captures
• Verify that all candidate outputs and states are finite
• Compare output/state mean absolute error and relative L2 error against the FLA reference

End-to-end validation

Run 50 real prompts at concurrency 1, 2, 4, and 8 with:

• Greedy decoding
• Thinking disabled
• Prefix-cache hits fixed at zero
• Identical model and serving configuration for FLA and CuteDSL

Compare all generated responses character-for-character.

For high-concurrency validation, run 200 requests at client concurrency 16, 32, 64, 128, 136, 144, 160, 176, 192, and 200. Use unique cache salts and compare each CuteDSL response against the corresponding FLA response.

Test Result

GPU regression

With the original three-round implementation, the new regression test fails as expected:

output max absolute error: 752.116699
expected: < 0.1

With the fourth correction round:

7 passed in 10.09s

Real hidden-state replay

All 90 layer/rank captures remained finite and passed the error gates.

Worst observed values:

End-to-end correctness

The four concurrency groups produced 200 responses in total:

50/50 successful at concurrency 1
50/50 successful at concurrency 2
50/50 successful at concurrency 4
50/50 successful at concurrency 8

All 200 CuteDSL responses matched their corresponding FLA responses character-for-character. No empty output, malformed text, Unicode replacement characters, or prefix-cache hits were observed.

The high-concurrency matrix completed 200/200 requests at every point from C16 through C200. All 2,000 CuteDSL responses matched FLA exactly. No HTTP 409 response, timeout, OOM, kernel error, malformed output, or prefix-cache hit was observed with the API concurrency limit set to 256.

Static checks

The applicable checks passed:

• Ruff check
• Ruff format
• typos
• mypy for Python 3.10
• SPDX header validation
• forbidden-import validation
• torch.cuda API validation
• configuration validation

No documentation update is required because this PR only changes the finite-precision convergence of an existing internal kernel and adds a regression test.

AI assistance was used during failure analysis, regression-test development, and preparation of this PR description. The author reviewed the code changes and independently validated all reported test results.