---
id: pr-tensorrt-llm-6538
repo: NVIDIA/TensorRT-LLM
pr: 6538
title: '[None][feat] Use Separate QKV Input Layout for Context MLA'
author: zhhuang-nv
date: '2025-08-19'
url: https://github.com/NVIDIA/TensorRT-LLM/pull/6538
source_category: upstream-code
architectures:
- sm90
tags:
- arch_pipeline
- attention_kv
- paged MQA
techniques:
- async-copy
- cache-policy
- kernel-fusion
- pipeline-stages
hardware_features:
- hopper
- tma
kernel_types:
- attention
- mla
- runtime
- softmax
languages:
- cuda-cpp
captured_at: '2026-05-16'
status: merged
merge_sha: 7e135d2ea7d989d6f45855bbed27804477ab9286
inclusion_reason: Preserve page/KV/layout and plan-run contracts; profile memory traffic separately from math utilization.
changed_paths:
- cpp/kernels/fmha_v2/fmha_test.py
- cpp/kernels/fmha_v2/setup.py
- cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h
- cpp/kernels/fmha_v2/src/fmha/kernel_traits.h
- cpp/kernels/fmha_v2/src/fused_multihead_attention.cpp
- cpp/tensorrt_llm/common/attentionOp.cpp
- cpp/tensorrt_llm/common/attentionOp.h
- cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_cubin.h
- cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_k_v_128_alibi_tma_ws_sm90.cubin.cpp
- cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_k_v_128_softcapping_tma_ws_sm90.cubin.cpp
- cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_k_v_128_tma_ws_sm90.cubin.cpp
- cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp
- cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp
- cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp
- cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp
- cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp
artifact_dir: evidence/pull-bundles/tensorrt-llm/gh-6538
head_ref_oid: 1c74fe8dcd6de7de11faab2141fb2efeca333ea2
---
# [None][feat] Use Separate QKV Input Layout for Context MLA

- Repository: `NVIDIA/TensorRT-LLM`
- PR: [6538](https://github.com/NVIDIA/TensorRT-LLM/pull/6538)
- Decision: `Preserve page/KV/layout and plan-run contracts; profile memory traffic separately from math utilization.`
- Reason: Preserve page/KV/layout and plan-run contracts; profile memory traffic separately from math utilization.

## Changed Paths

- `cpp/kernels/fmha_v2/fmha_test.py`
- `cpp/kernels/fmha_v2/setup.py`
- `cpp/kernels/fmha_v2/src/fmha/gmem_tile_qkv_packed.h`
- `cpp/kernels/fmha_v2/src/fmha/kernel_traits.h`
- `cpp/kernels/fmha_v2/src/fused_multihead_attention.cpp`
- `cpp/tensorrt_llm/common/attentionOp.cpp`
- `cpp/tensorrt_llm/common/attentionOp.h`
- `cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_cubin.h`
- `cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_k_v_128_alibi_tma_ws_sm90.cubin.cpp`
- `cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_k_v_128_softcapping_tma_ws_sm90.cubin.cpp`
- `cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_k_v_128_tma_ws_sm90.cubin.cpp`
- `cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_softmax_tma_ws_sm90.cubin.cpp`
- `cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_kv_128_tma_ws_sm90.cubin.cpp`
- `cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_alibi_tma_ws_sm90.cubin.cpp`
- `cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_softcapping_tma_ws_sm90.cubin.cpp`
- `cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/cubin/fmha_v2_flash_attention_bf16_64_128_S_q_paged_kv_128_tma_ws_sm90.cubin.cpp`

## Evidence Bundle

`evidence/pull-bundles/tensorrt-llm/gh-6538`
