unittest: Add head dim 256 test cases and mark as xfail

[bkryu]

📌 Description

Adding unit test for head_dim=256 cases for trtllm-gen decode and marking them as xfail.

Renames test_trtllm_batch_decode to _test_trtllm_batch_decode as a base function. Test functions now call _test_trtllm_batch_decode with a group of parameter combinations:

• test_trtllm_batch_decode --> 1632 existing parameter combinations
• test_trtllm_batch_decode_bs1 --> 1 xfail case with batch size 1
• test_trtllm_batch_decode_head_dim_256 --> 40 xfail cases with head_dim=256.
• test_trtllm_batch_decode_long_sequence_length --> 48 cases of long seqlen.

🔍 Related Issues

#1993

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

Summary by CodeRabbit

• Tests
  • Refactored batch-decode tests into a parameterized wrapper to improve coverage across head-dimension variants.
  • Added new scenarios including head_dim=256 and long-sequence-length cases (head_dim=256 marked as expected failures).
  • Updated single-batch tests to include head-dimension parameterization.
  • Removed direct execution blocks to standardize on pytest-based test runs.

[coderabbitai]

Walkthrough

Tests updated to parameterize batch-decode tests with a new head_dim parameter, add head_dim=256 variants (marked xfail), refactor common helper _test_trtllm_batch_decode, and remove the if __name__ == "__main__" invocation block for pytest execution.

Changes

Cohort / File(s)	Summary
Test Parameterization & New Head-Dim Variants tests/attention/test_trtllm_gen_attention.py	Added head_dim parameter to test_trtllm_batch_decode and test_trtllm_batch_decode_bs1; introduced _test_trtllm_batch_decode helper; added test_trtllm_batch_decode_head_dim_256 and long-sequence/head_dim=256 variants (marked xfail); expanded parameterized test combinations; removed __main__ block.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

• Review points:
  • Ensure head_dim is passed through all helper and test call paths.
  • Confirm xfail markers are correctly applied to head_dim=256 variants.
  • Verify removal of __main__ doesn't affect local debug patterns.

Poem

🐰 I hopped through tests with nimble pace,
I passed a head_dim into the race,
Two-five-six waves a cheeky sigh,
Parametric hops reach for the sky,
A rabbit cheers for cleaner test-space.

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 14.29% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The PR title "unittest: Add head dim 256 test cases and mark as xfail" directly corresponds to the main changes in the changeset. The raw summary confirms that new test variants for head_dim=256 were introduced and marked as xfail, which is exactly what the title describes. The title is concise, clear, and specific enough that a teammate reviewing the git history would understand the primary change without needing to read the full diff. While the title doesn't mention secondary changes like the test refactoring or removal of main blocks, this is appropriate and expected per the guidelines.
Description Check	✅ Passed	The pull request description follows the required template structure and includes all mandatory sections. The 📌 Description section clearly explains the changes (adding unit tests for head_dim=256 cases marked as xfail, refactoring test functions with parametrization), the 🔍 Related Issues section links to issue #1993, and the 🚀 Pull Request Checklist is complete with all three pre-commit checks and both test items marked as completed. The Reviewer Notes section is left blank, which is acceptable since the template explicitly marks it as optional. The description provides specific, actionable information about the test reorganization and new test variants created.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[yzh119]

It's unusual practice (but it's totally okay to do so and it's not introduced in this PR) to call one test_* function instead another one, all top-level functions with prefix test_ will be treated as standalone unittests.

Can we create a function _test_trtllm_batch_decode as the common body of these unittests, instead of calling another top-level test_trtllm_batch_decode function?

[bkryu]

Hi @yzh119, I think this PR is a good opportunity to do make this change.

I have:

• renamed test_trtllm_batch_decode to _test_trtllm_batch_decode as a base function.
• There are test functions that call _test_trtllm_batch_decode with a group of parameter combinations:
  • test_trtllm_batch_decode --> 1632 existing parameter combinations
  • test_trtllm_batch_decode_bs1 --> 1 xfail case with batch size 1
  • test_trtllm_batch_decode_head_dim_256 --> 40 xfail cases with head_dim=256.
  • test_trtllm_batch_decode_long_sequence_length --> 48 cases of long seqlen.

[bkryu]

The long seqlen was added because I saw #1968 and tested what happens if try testing long seqlens. We start to see failures starting from 4k

[yzh119]

LGTM

[yzh119]

/bot run

[flashinfer-bot]

GitLab MR !96 has been created, and the CI pipeline #37485772 is currently running. I'll report back once the pipeline job completes.

[flashinfer-bot]

[SUCCESS] Pipeline #37485772: 13/17 passed