[jiaganc]

@coderabbitai summary

Description

This PR removes redundant recurrent-state and metadata movement from Kimi K3 KDA prefill.

• Make the optimized KDA prefill path read and write the recurrent-state pool directly with native int32 state indices, eliminating full-state gather, transpose, and scatter copies.
• Share indexed recurrent-state row reset logic between KDA and GDN under the Mamba module.
• Prepare KDA chunk indices and sequence properties once per executor step with Kimi K3-specific attention metadata.
• Keep Mamba2Metadata as the lazy default for generic attention metadata to avoid an import cycle, while selecting the Kimi K3 metadata subclass in the executor.
• Consolidate correctness coverage around module parity and fallback parity, and register the new metadata test in the B200 test list.

This is a stacked draft targeting codex/kda-refactor-part-1. The branch contains one squashed commit beyond that base.

Test Coverage

• Built all 3,878 TensorRT-LLM targets for CUDA architecture 100-real on B200.
• Ran every added or modified runnable unit-test file in this branch on B200 GPU 0:
  • test_kda_prefill_op.py
  • test_kimi_k3_mamba_metadata.py
  • test_kimi_kda_fused_verify_parity.py
  • test_kimi_kda_verify_parity.py
  • test_gdn_kernel_optimizations.py
  • test_recurrent_state_cache.py
• Result: 81 passed, 0 failed, 0 skipped in 899.58 seconds.
• Repository pre-commit hooks passed for the squashed commit.

PR Checklist

Please review the following before submitting this PR:

• PR description clearly explains what and why.

• PR follows the TRT-LLM coding guidelines to the best of my knowledge.

• Test cases cover the new code paths.

• No public API or dependency changes are introduced.

• Documentation, CODEOWNERS, and architecture diagram updates are not required for this internal optimization.

• 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, comment /bot help.