[None][perf] Enable configurable MLA skip correction for TRT-LLM Gen FMHA - #18118

Description

This PR exposes threshold-based skip correction for the TRT-LLM Gen MLA FMHA path through the TRT-LLM runtime and TorchLlmArgs. The optimization is opt-in and disabled by default, so existing inference behavior is unchanged unless users explicitly enable it.

The required TRT-LLM Gen kernel/export changes are already present on main through 6c1ce33e7fbb1730c3d85150e6ac73ebc1b76deb. This PR reuses those assets and contains only the TRT-LLM runtime, configuration, validation, and test integration needed to activate the feature.

Scope relative to main

This branch does not regenerate or modify the TRT-LLM Gen exported artifacts already available on main, including:

• FMHA cubins and cubin metadata
• FMHA static libraries
• trtllmGen_fmha_export generated interfaces
• generated cuda_ptx.h and kernelMetaInfo.h

The remaining changes under trtllmGenKernels/fmha are runtime glue: they pass the configured threshold into the existing generated kernels, validate it against the selected BMM2 dtype, and apply the correct softmax scale in the separate reduction path.

The PR is a single signed commit on top of the latest main.

User-facing configuration

• enable_skip_correction: enables the optimization; defaults to False.
• skip_correction_threshold: controls the base-2 row-max threshold; defaults to 8.0 and accepts values in (0, 32].
• The selected kernel's BMM2 dtype further limits the runtime threshold to 8 for E4M3, 15 for FP16, and 32 for BF16.
• Enabling the feature requires attn_backend="TRTLLM".
• When disabled, TRT-LLM forwards an effective threshold of 0.0, preserving the existing FMHA behavior.

Example configuration:

attn_backend: TRTLLM
enable_skip_correction: true
skip_correction_threshold: 8.0

trtllm-serve /path/to/DeepSeek-V4-Flash \
  --backend pytorch \
  --config skip-correction.yaml

Omit skip_correction_threshold to use the default value of 8.0.

Implementation

• Propagates the effective threshold through TorchLlmArgs, ModelConfig, MLA construction, FMHA backend selection, the Python/C++ attention operator, and the TRT-LLM Gen FMHA runner.
• Adds an FMHA capability flag and selects only a backend that supports skip correction; currently this is FallbackFmha.
• Validates the runtime threshold against the selected kernel's BMM2 dtype.
• Passes the threshold-derived softmax scale to the separate multi-CTA reduction path.
• Supports SM100 and SM103. Unsupported architectures disable the optimization with a one-time warning.
• Adds opt-in DeepSeek-V4-Flash GSM8K L0 coverage on B200 and B300.
• Updates the LLM-argument golden manifest and API stability reference for the additive configuration fields.

Impact and compatibility

• API impact: additive and backward compatible.
• Default runtime behavior: unchanged because skip correction is disabled by default.
• Performance impact: limited to explicitly enabled MLA workloads on supported Blackwell architectures.
• Numerical risk: bounded by architecture, FMHA-backend capability, and BMM2 dtype validation.
• Dependencies: no new runtime dependencies.
• Rubin support remains out of scope.

Test Coverage

Feature-path tests explicitly enable skip correction; existing tests continue to exercise the disabled/default behavior.

Added coverage

• tests/unittest/llmapi/test_llm_args.py
  • configuration defaults and overrides
  • invalid threshold rejection
  • TRTLLM attention-backend requirement
  • YAML configuration loading
• tests/unittest/_torch/attention/test_attention_mla.py
  • supported and unsupported SM behavior
  • non-positive runtime threshold disabling
  • FMHA capability and registry filtering
  • dense MLA correctness with BF16 and FP8 KV cache
• tests/unittest/_torch/attention/sparse/test_sparse_mla_forward.py
  • sparse MLA correctness with BF16 and FP8 KV cache
• tests/integration/defs/accuracy/test_llm_api_pytorch.py
  • DeepSeek-V4-Flash GSM8K with skip correction explicitly enabled
• B200 and B300 L0 test-list registration

Validation results

• Full SM100 build with warnings-as-errors: passed (--parallel 224).
• Incremental C++ build including unit-test binaries: passed (make -j224).
• LLM args and API stability: 379 passed.
• Dense MLA skip-correction tests: 22 passed.
• Sparse MLA skip-correction tests: 2 passed.
• FMHA page-index tests: 30 passed.
• LLM-argument manifest regeneration: no unexpected diff.
• DeepSeek-V4-Flash GSM8K on B200: 95.262 accuracy versus 95.110 reference.
• GitHub DCO, PR checks, base freshness, API compatibility, and pre-commit checks: passed.
• Pre-push confidentiality scan: passed.

CI Failure Analysis

The previous L0_Test-x86_64-Multi-GPU failure was:

accuracy/test_disaggregated_serving.py::TestNemotron3Super120B::test_ctx_dp2_gen_tp4

This case does not use the MLA path changed by this PR, and skip correction is disabled by default. Recent CI history also shows the same accuracy-boundary failure on unrelated PRs, while post-merge runs passed. It is therefore classified as an existing flaky accuracy case rather than a regression from this change.

This PR does not add or modify any waiver for that test.

PR Checklist

• Reuses the TRT-LLM Gen artifacts already merged to main.
• Contains one signed feature commit without generated cubin/library churn.
• New opt-in code paths have unit and E2E coverage.
• Existing configurations preserve disabled/default behavior.
• The LLM API change is additive and labeled api-compatible.
• No new dependencies, ownership changes, or test waivers are introduced.
• Rubin support is explicitly out of scope.