[None][opt] attn kernel epilogue fuse RopeQuant#14947
Conversation
yunruis
changed the title
yunruis
changed the title
|
/bot run |
|
PR_Github #52016 [ run ] triggered by Bot. Commit: |
|
PR_Github #52016 [ run ] completed with state
|
|
/bot run |
|
PR_Github #52212 [ run ] triggered by Bot. Commit: |
|
PR_Github #52212 [ run ] completed with state
|
yunruis
force-pushed
the
user/yunruis/dsv4_attn_epilogue_fusion
branch
from
5231ac5 to
af4076e
Compare
|
/bot run |
|
PR_Github #52666 [ run ] triggered by Bot. Commit: |
|
PR_Github #52666 [ run ] completed with state |
|
/bot run |
|
/bot run |
|
PR_Github #52681 [ run ] triggered by Bot. Commit: |
|
PR_Github #52681 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #52855 [ run ] triggered by Bot. Commit: |
|
PR_Github #52855 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #52959 [ run ] triggered by Bot. Commit: |
|
PR_Github #52959 [ run ] completed with state |
yunruis
changed the title
|
/bot run --disable-fail-fast |
|
PR_Github #57427 [ run ] triggered by Bot. Commit: |
|
PR_Github #57427 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #57478 [ run ] triggered by Bot. Commit: |
|
PR_Github #57478 [ run ] completed with state |
Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com>
yunruis
force-pushed
the
user/yunruis/dsv4_attn_epilogue_fusion
branch
from
2bd7d24 to
d04b08e
Compare
|
/bot run --disable-fail-fast |
|
all comments has been resolved from my side. If you agree please mark it as resolved. Thanks @heyuhhh @mingyangHao |
|
PR_Github #57737 [ run ] triggered by Bot. Commit: |
|
PR_Github #57737 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #57757 [ run ] triggered by Bot. Commit: |
|
PR_Github #57757 [ run ] completed with state |
|
All comments has been resolved. |
|
✅ LFS objects already in storage (3240 files) — no sync needed. These LFS-tracked files are already present in this repository's LFS storage:
|
Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com>
Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com>
… (NVIDIA#16036) Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com>
6 participants
1. Summary
This PR integrates DeepSeek-V4 FMHA epilogue fusion into the TensorRT-LLM PyTorch MLA path.
Before this change, the DeepSeek-V4 attention path materialized the FMHA latent output and then launched a separate inverse-RoPE + FP8 quantization kernel before the DSv4 O projection BGEMM:
With epilogue fusion enabled, trtllm-gen FMHA performs inverse RoPE and 1x128 E4M3 quantization in the FMHA correction epilogue, and directly writes the physical layout consumed by the DSv4 O projection BGEMM:
The optimization removes the standalone rope-quant kernel launch and avoids the intermediate latent write/read between FMHA and quantization. The O projection itself is not fused into FMHA; it still runs as the existing DSv4 FP8 BGEMM, now consuming
(fp8_o, fp32_scale)produced by FMHA.2. Applicable Scenarios
The feature is intentionally gated to the validated DeepSeek-V4 Pro configuration:
num_heads == num_heads_tp == 128.The runtime kill switch is:
By default, the feature is allowed and the runtime gate decides whether the current batch can use the fused path.
3. Perf Gain
The kernel-level comparison should be reported as:
4. Major Changes
4.1 New DSv4 Epilogue Fusion Inputs and Outputs
New runtime inputs/outputs are added for the fused path:
output/oPtroutput_sf/oSfPtrdsv4_inv_rope_cos_sin_cacheenable_dsv4_epilogue_fusionscaleBufMoutput_sf.size(2)scaleBufMis an input scalar, not an output tensor. It describes the padded token stride ofoutput_sfand is propagated to trtllm-gen so the kernel writes FP32 scales with the same layout allocated by TensorRT-LLM.The fused output layout is:
4.2 Python Model Forward Changes
The DeepSeek-V4 MLA forward path now has two possible attention outputs:
The O projection call site remains unified in
MLA.forward(). The fused path changes the input type to_deepseek_v4_o_proj()but does not move the O projection into FMHA.For
register_to_config, the MLA custom op contract is extended so custom-op execution can receive DSv4 fused buffers explicitly:trtllm::create_mla_outputscreates the legacy output and, when the runtime gate passes, the DSv4(fp8_o, output_sf)buffers.trtllm::mla_custom_op_inplaceaccepts optionaldsv4_output,dsv4_output_sf, andenable_dsv4_epilogue_fusion.forward_impl_with_deepseek_v4(..., dsv4_epilogue_output=...).This keeps the custom-op path functionally aligned with the eager Python path and avoids dropping the fused epilogue buffers before the O projection.