[https://nvbugs/6108995][fix] Fix workspace size calculation for fmha…

[richardhuo-nv]

…_bmm1_scale_size with FP8ContextMLA (#13880)

Summary by CodeRabbit

• Bug Fixes
  • Fixed memory allocation for 8-bit floating-point context operations to properly support multiple attention computation variants, improving stability and resource efficiency in specialized inference workloads.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually 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.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: b3614c90-2a3a-488a-b6d1-496f376d505b

📥 Commits

Reviewing files that changed from the base of the PR and between ca0aad2 and 8d5a8ba.

📒 Files selected for processing (1)

• cpp/tensorrt_llm/common/attentionOp.cpp

---

📝 Walkthrough

Walkthrough

Workspace allocation for FMHA scale buffers is expanded to support both FP8 Context FMHA and FP8 Context MLA configurations. Previously allocated only for FMHA, the scale buffers are now reserved when either FMHA or MLA is enabled.

Changes

FP8 Context MLA Workspace Buffer Allocation

Layer / File(s)	Summary
FMHA scale buffer workspace allocation for FP8 Context MLA cpp/tensorrt_llm/common/attentionOp.cpp	Scale buffer sizes (fmha_bmm1_scale_size, fmha_bmm2_scale_size) are allocated in AttentionOp::getWorkspaceSizeForContext when either mFP8ContextFMHA or mFP8ContextMLA is enabled, replacing the prior FMHA-only condition.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#14067: Both PRs adjust TensorRT-LLM attention workspace sizing—main PR expands FP8 Context FMHA/MLA scale buffer allocations, and the retrieved PR addresses related attention workspace_size configuration issues.

Suggested reviewers

• litaotju
• pengbowang-nv

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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