DSA cuDNN real-kernel tests (packed THD + CP) return wrong results on GB200/SM100; masked by `flaky`/`flaky_in_dev` markers
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Summary
Two real-kernel tests added in #5099 (`tests/unit_tests/transformer/experimental_attention_variant/test_dsa_native_parity.py`) fail **numerically** on GB200 with the cuDNN DSA backend:
- `test_cudnn_indexer_topk_single_packed_cp_real_kernel_uses_bottom_right_alignment` — returned top-k indices are all wrong (do not match the bottom-right-aligned reference).
- `test_cudnn_full_fusion_real_kernel_packed_cp_varlen_matches_reference` — output vs reference cosine similarity **0.365** (reference path and cuDNN fused path disagree).
Both are the packed-THD + context-parallel (CP>1) path, which is exactly what the GB200 GLM-5.2 SFT recipe uses (`dsa_kernel_backend=cudnn`, CP4, packed sequences, `cp_comm_type=allgather`).
On `core_r0.19.0` and `main` both tests carry `@pytest.mark.flaky` / `@pytest.mark.flaky_in_dev` with comments citing an import error (`cutlass.cute.core has no attribute ThrMma`) and "a CUDA error in CI (deterministic, not truly flaky)". What we see is different and arguably worse: the kernels run and produce wrong numbers. The markers hide this in CI.
## Environment
- `nvcr.io/nvidia/nemo:26.08` (aarch64): Megatron-Core 0.19 (16ad357), TE 2.17.1, cudnn-frontend linked to libcudnn 9.21.1, torch 2.13.0a0+nv26.6 (built vs cuDNN 9.23.0 — separate mismatch, reported to NeMo), flash_mla 1.0.0 (sm_100), tilelang present
- GB200 NVL72, CC 10.0, driver 580.126.20, 4 GPU/node
## Reproduction (single node, 4 GPUs)
```bash
cd /opt/Megatron-Bridge/3rdparty/Megatron-LM
pytest -x -q tests/unit_tests/transformer/experimental_attention_variant/test_dsa_native_parity.py -k "cudnn" -p no:cacheprovider
```
Result: 78 passed, 2 failed (the two above). `-k tilelang` fails 8 tests as well (we don't use that backend). The TP-consistency test fails only because world_size==1.
Failure excerpts:
```
# test_cudnn_indexer_topk_single_packed_cp_real_kernel_uses_bottom_right_alignment (test_dsa_native_parity.py:853)
> torch.testing.assert_close(sorted_indices.cpu(),
torch.tensor([[[1, 2], [2, 3], [11, 12], [12, 13]]], dtype=torch.int32), rtol=0, atol=0)
E AssertionError: Tensor-likes are not equal!
E Mismatched elements: 8 / 8 (100.0%)
E Greatest absolute difference: 13 at index (0, 3, 0)
# test_cudnn_full_fusion_real_kernel_packed_cp_varlen_matches_reference (dsa_native_parity_utils.py:412)
> assert cosine > 1 - eps, f"{prefix}cosine_sim={cosine:.6f}"
E AssertionError: cosine_sim=0.365310
```
Both tensors are finite (the `isfinite` asserts before the cosine check pass), i.e. the kernels run to completion with wrong values.
Control for the cuDNN-version hypothesis: the image's torch was built against cuDNN 9.23.0 while libcudnn 9.21.1 is installed (separate packaging issue). Installing `nvidia-cudnn-cu13` 9.23 in a scratch layer and putting it first on `LD_LIBRARY_PATH` re-runs the two tests with the **same two failures** (`2 failed, 1 passed`), so the cuDNN runtime version is not the cause.
## Impact / what we checked
- Training-shape control: GLM-5.2 SFT at 32 nodes (EP32, CP4, seq 8192, packed) with `dsa_kernel_backend=cudnn` vs `=none` (reference) — per-step lm loss within ≤0.7% over 10 steps and the same 100-step plateau (0.646 vs 0.615 mean over steps 51–100). So at our shapes the fused path tracks the reference, but the unit tests say some packed+CP configurations are wrong. We could not determine which shapes are affected.
- Related: #4878 (CP for DSAttention) is still open although #5099 landed CP via THD; this issue is about the correctness of that THD+CP path on SM100.
## Ask
1. Confirm whether the SM100 real-kernel path for packed+CP is expected to pass; if yes, this is a kernel/cudnn-frontend bug on GB200.
2. Replace the `flaky` markers with a hardware-specific skip + tracking issue so the failure is visible.
Contributor guide
Research direction
Start with tests/unit_tests/transformer/experimental_attention_variant/test_dsa_native_parity.py and dsa_native_parity_utils.py, then run the provided pytest command on GB200/SM100. Investigate the two packed-THD plus context-parallel cuDNN failures and the existing flaky markers. Done means the failure is correctly tracked and no longer hidden by misleading flaky markers, using a justified hardware-specific skip if needed.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, testing-qa
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100