Skip to content

[None][feat] Support FP8 KV cache for non-MLA PrimsTS attention #19396

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuxianq merged 13 commits into from
Nov 10, 2026
+659 −147
Merged

[None][feat] Support FP8 KV cache for non-MLA PrimsTS attention #19396

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
011eb24
[None][feat] Support FP8 KV cache for non-MLA PrimsTS attention
yuxianq Sep 18, 2026
2ce050a
[None][fix] Adapt FP8 decode integration to refreshed PrimTS APIs
yuxianq Oct 13, 2026
ae3b728
[None][fix] Separate persistent FP8 decode scale specializations
yuxianq Oct 29, 2026
77aaec7
[None][fix] Enable non-MLA PrimsTS attention on Rubin
yuxianq Oct 30, 2026
c8cafb6
[None][fix] Simplify PrimsTS FP8 scale arguments
yuxianq Oct 30, 2026
1fefcd7
[None][fix] Consolidate PrimsTS decode compile identity and reducer a…
yuxianq Oct 30, 2026
c124075
[None][fix] Preserve MLA architecture routing after rebase
yuxianq Oct 30, 2026
622ccf4
[None][fix] Bind FP8 attention scales through the upstream PrimTS APIs
yuxianq Nov 7, 2026
b146f59
[None][fix] Address PrimTS FP8 tail safety and review coverage
yuxianq Nov 8, 2026
724734c
[None][fix] Repair PrimTS accuracy selectors and model-test imports
yuxianq Nov 8, 2026
a9bdfe4
[None][fix] Share PrimTS architecture support after Rubin qualification
yuxianq Nov 8, 2026
17b8fe0
[None][fix] Zero FP8 KV tails in fused preprocessing
yuxianq Nov 8, 2026
fc41ebd
[None][test] Reduce PrimTS FP8 GPU regression matrix
yuxianq Nov 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 17 additions & 4 deletions cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,12 @@ inline __device__ void zeroPagedVCacheTail(QKVPreprocessingParams<T, KVCacheBuff
{
// Context owns the last physical page exclusively. Clear unused V rows because paged FMHA loads full pages and
// masked zero probabilities can still propagate NaNs through BMM2.
if constexpr (!GEN_PHASE && std::is_same_v<KVCacheBuffer, KVBlockArray> && std::is_same_v<T, TCache>)
#ifdef ENABLE_FP8
constexpr bool kFp8Cache = std::is_same_v<TCache, __nv_fp8_e4m3>;
#else
constexpr bool kFp8Cache = false;
#endif
if constexpr (!GEN_PHASE && std::is_same_v<KVCacheBuffer, KVBlockArray> && (std::is_same_v<T, TCache> || kFp8Cache))
{
if (params.generation_phase || !isUniqueTokenWriter || cacheSeqLen <= 0 || tokenIdxInKvCache != cacheSeqLen - 1)
{
Expand All @@ -210,12 +215,20 @@ inline __device__ void zeroPagedVCacheTail(QKVPreprocessingParams<T, KVCacheBuff
}

auto* vBlock = reinterpret_cast<TCache*>(params.kv_cache_buffer.getVBlockPtr(batchIdx, cacheSeqLen - 1));
auto* vBlockVec = reinterpret_cast<TVec*>(vBlock);
TVec const zero{};
for (int tokenInBlock = validTokensInLastBlock; tokenInBlock < tokensPerBlock; ++tokenInBlock)
{
int const vIdx = kvHeadIdx * tokensPerBlock * vecsPerHead + tokenInBlock * vecsPerHead + channelIdx;
vBlockVec[vIdx] = zero;
if constexpr (std::is_same_v<T, TCache>)
{
reinterpret_cast<TVec*>(vBlock)[vIdx] = TVec{};
}
else
{
// Match the packed FP8 stores used for the valid V tokens.
using CacheVec = typename mmha::packed_type<TCache, mmha::num_elems<TVec>::value>::type;
static_assert(sizeof(CacheVec) == sizeof(TVec) / sizeof(T));
reinterpret_cast<CacheVec*>(vBlock)[vIdx] = CacheVec{};
}
}
}
}
Expand Down
12 changes: 6 additions & 6 deletions cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -407,8 +407,8 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor
qkvParams.generation_phase = false;
qkvParams.q_only_input = qOnlyInput;
qkvParams.multi_processor_count = static_cast<int>(multi_processor_count);
qkvParams.rotary_vision_start = 0;
qkvParams.rotary_vision_length = 0;
qkvParams.rotary_vision_start = -1;
qkvParams.rotary_vision_length = -1;

switch (qkvDtype)
{
Expand Down Expand Up @@ -559,8 +559,8 @@ void trtllmGenContextPostprocess(torch::Tensor qkv_input, torch::Tensor workspac
qkvParams.quantized_fp8_output = fp8_context_fmha;
qkvParams.generation_phase = false;
qkvParams.multi_processor_count = static_cast<int>(multi_processor_count);
qkvParams.rotary_vision_start = 0;
qkvParams.rotary_vision_length = 0;
qkvParams.rotary_vision_start = -1;
qkvParams.rotary_vision_length = -1;

switch (qkvDtype)
{
Expand Down Expand Up @@ -745,8 +745,8 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace,
qkvParams.generation_phase = true;
qkvParams.q_only_input = qOnlyInput;
qkvParams.multi_processor_count = static_cast<int>(multi_processor_count);
qkvParams.rotary_vision_start = 0;
qkvParams.rotary_vision_length = 0;
qkvParams.rotary_vision_start = -1;
qkvParams.rotary_vision_length = -1;

switch (qkvDtype)
{
Expand Down
2 changes: 1 addition & 1 deletion tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -286,7 +286,7 @@ def _page_size_exclusive_to_cute_dsl(meta: "TrtllmAttentionMetadata") -> bool:
if tokens_per_block in getattr(kv_cache_manager, "trtllm_gen_extra_tokens_per_block", ()):
return False
return not (
get_sm_version() in PrimsTSFmha.SUPPORTED_SMS
get_sm_version() in PrimsTSFmha.SUPPORTED_MLA_SMS
and tokens_per_block in PrimsTSFmha.SUPPORTED_PAGE_SIZES
and getattr(kv_cache_manager, "dtype", None) in (DataType.FP8, DataType.BF16)
)
Expand Down
Loading
Loading