SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

Loosened trtllm_ragged_attention_deepseek shape assertion (#3064)

<!-- .github/pull_request_template.md -->

## ๐Ÿ“Œ Description

`trtllm_ragged_attention_deepseek` kernel is available for GLM-5 but is
currently blocked by the assertion.
Loosened assertion and added unit test to verify. Cutlass does **not**
support this shape, so a naive implementation is used to test
correctness.
Also added missing checks in cutlass code to error out on unsupported
shapes.

## ๐Ÿ” Related Issues

<!-- Link any related issues here -->

## ๐Ÿš€ 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](https://pre-commit.com/).

## ๐Ÿงช Tests

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added GLMโ€‘5 head-dimension configuration support for attention prefill
(256,256,256).

* **Improvements**
* Validation expanded to accept the GLMโ€‘5 dimension set and now reports
actual dimensions alongside supported tuples on failure.
* Failure paths now return errors immediately and emit clearer
diagnostics when unsupported mask/head-dimension combinations occur.

* **Tests**
* Added CUDA-only ragged-attention tests with a float32 reference to
verify GLMโ€‘5 prefill behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Claude <noreply@anthropic.com>
N
nvjullin committed
1aa32d03ef93e5d75de221db965b3e07cfc4e1c5
Parent: c5c089b
Committed by GitHub <noreply@github.com> on 5/7/2026, 10:26:55 PM