[Bugfix] Add TRTLLM attention custom op and execution contract (#7214)
Signed-off-by: Rahul Steiger <rsteiger@aws-cmh-slurm-1-vscode-04.cm.cluster> Signed-off-by: Rahul Steiger <rsteiger@nvidia.com> Co-authored-by: Rahul Steiger <rsteiger@aws-cmh-slurm-1-vscode-04.cm.cluster>
This commit is contained in:
1 parent
b65f97bc48
commit
4934741df9
4 files changed
+609
-47
No files matched your search
@@ -131,3 +131,24 @@ and real-kernel schema/fake checks pass. The full suite must be rerun with the
|
||||
vLLM version required by upstream. NPU/ROCm device numerics and compilation remain
|
||||
unvalidated. Pre-commit passes with the CI hook skips. Upstream PyTorch/CUTLASS
|
||||
warnings remain.
|
||||
|
||||
## TRTLLM dense execution
|
||||
|
||||
`TRTLLM_ATTN` declares `SUPPORTED` / `CUSTOM_OP` for noncausal dense BF16 on
|
||||
B200/GB200 (SM100) and B300/GB300 (SM103), head dimension 128, and equal Q/K/V
|
||||
head counts, without parallel, paged-KV, piecewise, or HSDP boundaries. Pre-construction, SAGE, skip-softmax, packed,
|
||||
and other unverified paths remain `UNMIGRATED`. Resolution and dispatch share
|
||||
metadata validation; workspace mutation is explicit in the custom-op schema.
|
||||
|
||||
CPU contract tests cover fullgraph replay and schema/fake consistency using a
|
||||
substituted dispatcher. Real-kernel validation requires Blackwell with FlashInfer:
|
||||
|
||||
```bash
|
||||
python -m pytest tests/diffusion/attention/test_trtllm_attn.py \
|
||||
-k dense_contract_fullgraph_matches_sdpa -q -rs
|
||||
```
|
||||
|
||||
These hardware tests compare eager/compiled output with FP32 SDPA and check
|
||||
schema/fake agreement. The TRTLLM and contract suite passed on Blackwell before the final architecture
|
||||
restriction was added: 61 passed, no skips (19 dependency deprecation warnings).
|
||||
The SM100/SM103 restriction is additionally covered by CPU contract tests.
|
||||
Reference in new issue
Block a user