[https://nvbugs/6215793][fix] Thread ctx_total_kv_len through getWorkspaceSize→getWorkspaceSizeForContext…

[tensorrt-cicd]

Summary

• Root cause: FP8 context-MLA K/V workspace buffers sized by max_num_tokens but the quantize kernel writes total_kv_len rows (includes cached tokens under KV-cache reuse), overflowing mid-workspace and corrupting downstream slices.
• Fix: Thread ctx_total_kv_len through getWorkspaceSize→getWorkspaceSizeForContext and size K/V buffers as max(mChunkPrefillBufferBatchSize*max_num_tokens, total_kv_len)*dims.
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6215793

Summary by CodeRabbit

• Bug Fixes

  • Improved memory workspace allocation for attention operations to correctly account for key-value cache sizing during context processing.

• Tests

  • Removed test waiver for a previously failing accuracy test case.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7d217308-9c51-46e4-bd3c-9d692b1a9533

📥 Commits

Reviewing files that changed from the base of the PR and between fdcdcb3 and 9affc8b.

📒 Files selected for processing (4)

• cpp/tensorrt_llm/common/attentionOp.cpp
• cpp/tensorrt_llm/common/attentionOp.h
• cpp/tensorrt_llm/thop/attentionOp.cpp
• tests/integration/test_lists/waives.txt

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt

---

📝 Walkthrough

Walkthrough

This PR adds total_kv_len to attention workspace sizing interfaces, propagates it through runner calls, updates fp8 context MLA K/V sizing to use total KV bounds, and removes one integration test waiver entry.

Changes

Workspace sizing update

Layer / File(s)	Summary
API extension and call-chain propagation cpp/tensorrt_llm/common/attentionOp.h, cpp/tensorrt_llm/common/attentionOp.cpp, cpp/tensorrt_llm/thop/attentionOp.cpp	getWorkspaceSizeForContext and RunnerBase::getWorkspaceSize now accept total_kv_len, and attention() passes ctx_total_kv_len through Runner<T> into context workspace sizing.
FP8 context MLA K/V sizing logic cpp/tensorrt_llm/common/attentionOp.cpp	In the separate Q/K/V non-sparse MLA path, fp8 K/V row sizing now uses max(chunk-prefill-bound, total_kv_len) before computing fp8 K/V buffer sizes.

Integration waiver list change

Layer / File(s)	Summary
Waiver entry removal tests/integration/test_lists/waives.txt	Removed the waiver entry for accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[throughput_tp4].

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Suggested reviewers

• marinayanov
• xinhe-nv
• EmmaQiaoCh

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the bug fix (threading ctx_total_kv_len) and references the NVBugs ID (#6215793) following repository conventions.
Description check	✅ Passed	The PR description includes all key required sections: root cause explanation, the fix applied, test coverage confirmation, and bug link to nvbugs/6215793.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[trtllm-agent]

NVBug 6215793 was closed as Bug - Fixed. The repair-bot root cause matches the merged resolution #14852, so this PR is being auto-closed. If you believe this is wrong, reopen the PR — the bot will never auto-close a human-reopened PR for this cleanup generation.