[None][feat] support FP8 KV cache in PrimTS MLA decode

[yuxianq]

@coderabbitai summary

Description

Enable dense FP8 E4M3 MLA decode in PrimsTSFmha on SM100/SM103 using the existing vendored kernels.

• Consume quant_q_buffer and the BMM scales already produced by MLA preprocessing; keep the module-facing query and output BF16 and avoid re-quantizing Q.
• Size workspace and plan wrappers using FP8 Q/KV dtypes. Read the model-static scales once during eager warmup, using BMM1's regular scale (index 0), and reuse them without device-to-host scalar reads during decoding or CUDA graph capture.
• Validate preprocessing buffer contracts and reject unsupported quantization/output combinations. Preserve the existing generation-only, single-token, dense-MLA restrictions, FMHA priorities, and selection-cache keys.
• Fall back for FP8 MLA with <=64 local heads when per-sequence page-table capacity (kv_cache_block_offsets.shape[-1] * tokens_per_block) is below 128 tokens, which the vendored 1-CTA planner does not support. This is the wrapper plan's maximum KV length, not the live request length or the total shared KV-cache pool capacity.

FP8 standard attention, sparse/speculative MLA, quantized output, and vendor changes are not included.

Test Coverage

• Adapter and FMHA-manager unit suites: 176 passed. Coverage includes FP8 eligibility, malformed preprocessing buffers, quantization exclusions, workspace/plan dtypes, zero-copy query reinterpretation, non-unit scales, cached-scale reuse, and capture without warmup.
• B200 FP8 GPU validation: 9 passed, covering eager execution and CUDA graph replay.
• Existing B200 BF16 MLA regression tests: 2 passed (real v1 backend and wrapper graph replay).
• Real-preprocessing tests cover 6/12/96 local heads, two layers, multiple decode steps, numerical agreement with the existing BF16 reference, and adapter CUDA graph capture/replay.
• Independent Torch-reference tests cover batch sizes 2/65, non-unit BMM scales, FP8 probability rounding, a full-precision relative-error bound, and replay after changing query/page-table/sequence-length contents.
• KV-cache v2 tests are included, but local execution is blocked before attention by stale C++ bindings: KVCacheManagerConfig does not accept main's reuse_match_backoff argument. They need a current build/CI run. Full-model and distributed TP/DP runs have not been performed; head-count tests are per-rank shape coverage.
• Ruff, git diff --check, and pre-commit checks passed.

The worktree-local CUTLASS installation adds its base_dsl directory to sys.path. Local tests pre-import the attention package and explicitly declare that third-party path through pytest's pythonpath option to satisfy the repository's import-path check.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why.

• PR follows the TRT-LLM coding guidelines.

• Tests cover the new code paths.

• No public API changes or new dependencies.

• No ownership or significant architecture changes.

• Documentation updated as needed.

• Reviewers are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.