SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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