fix(trace): repair TGV and XQA MLA reference tests (#3365)
## Summary - construct `tgv_gemm_sm100` trace `b` with column-major strides, matching the kernel contract - add a TGV test layout assertion and avoid converting CUDA sync failures into skips - keep the XQA MLA beam dimension in the reference path and remove the invalid 64-query-head positive case Fixes #3352 Fixes #3353 ## Testing - reproduced the old TGV bad-layout failure on B200 with `torch 2.9.1+cu129`: `GPUassert: an illegal instruction was encountered` - `python -m compileall flashinfer/trace/templates/page.py tests/trace/test_tgv_gemm_sm100_reference_correctness.py tests/trace/test_xqa_mla_reference_correctness.py` - `FLASHINFER_WORKSPACE_BASE=/tmp/flashinfer-pr-cache /tmp/flashinfer-cu129/bin/python -m pytest tests/trace/test_tgv_gemm_sm100_reference_correctness.py -q -rs` -> 2 passed - on SM120, reproduced `upstream/main` XQA MLA failure: `output with shape [2, 1, 128, 512]` does not match broadcast shape `[2, 2, 128, 512]` - on SM120, fix branch: `CUDA_VISIBLE_DEVICES=0 FLASHINFER_WORKSPACE_BASE=/tmp/flashinfer-sm120-fix-cache python -m pytest tests/trace/test_xqa_mla_reference_correctness.py -q -rs --tb=short` -> 2 passed Co-authored-by: Avery Huang <averyh@nvidia.com>
E
eigen committed
fc9b2f64c15e9d2e9b2ecec47a9803317b63abfa
Parent: 9035311
Committed by GitHub <noreply@github.com>
on 5/20/2026, 5:11:46 PM