[TRTLLM-12807][feat] add CuteDSL FP8/FP16 MLA decode attention fmha lib - #17266
Closed
brnguyen2 wants to merge 29 commits into
Closed
[TRTLLM-12807][feat] add CuteDSL FP8/FP16 MLA decode attention fmha lib#17266brnguyen2 wants to merge 29 commits into
brnguyen2 wants to merge 29 commits into
Conversation
Add a CuteDSL attention backend that intercepts the MLA decode-only path
and dispatches to Blackwell CuTe DSL kernels (FP8 e4m3 / FP16), falling
back to TrtllmAttention for all other paths.
- attention_backend/cute_dsl.py: CuteDslAttention(TrtllmAttention)
- cute_dsl_kernels/blackwell/attention/mla: decode fp8/fp16 kernels + helpers
- custom_ops: register cute_dsl_mla_decode_{fp8,fp16}_blackwell ops
- utils/__init__: wire up CUTEDSL backend selection
- tests: CUTEDSL coverage in test_attention_mla + dedicated decode test
Signed-off-by: haow <haow@nvidia.com>
…ayer paged-KV, multi-layer test - attention_backend/cute_dsl.py: feed the scaled quant_q_buffer (not q.to(fp8)) and fold the fp8 dequant + de-folded log2(e) softmax scale into the bmm1/bmm2 scales; resolve the per-layer KV pool / page-table for layer_idx>0; gate the non-causal fast path off CUSTOM masks and speculative decoding; add a once-per-process ENGAGED log (symmetric to the fallback warning) confirming the decode kernel actually ran. - custom_ops/cute_dsl_custom_ops.py: add BF16 to the decode runner and the fp16 op with dtype dispatch + validation; pass workspace=None for split_kv==1 to avoid a zero-sized-buffer global write. - cute_dsl_kernels/.../mla_decode_fp8.py: allow a widened BF16 attention output. - cute_dsl_kernels/.../mla_decode_fp16.py: accept BF16 input/output in can_implement. - tests/.../test_cute_dsl_mla_decode.py: parametrize over num_layers to exercise the layer_idx>0 paged-KV path (reproduces the real DeepSeek-V3 E2E case). Signed-off-by: haow <haow@nvidia.com>
…fp8 scale Normalize the per-layer paged-KV pool view for the v1 KVCacheManager (interleaved single-pool layout) to a packed combined-slot view and fold the layer offset into the page table, so the CuTe DSL paged TMA addresses the correct layer's memory. No-op for the v2 manager. Cache the static fp8 dequant scales on the per-layer backend instance and read them eagerly during warmup, avoiding an illegal .item() device->host sync under CUDA graph capture. Add a long decode-only test that crosses paged-KV block boundaries mid-decode (short prompt, 64 steps), parametrized over v1/v2 KV cache managers, reproducing the DeepSeek-V3 E2E degeneration path. Signed-off-by: haow <haow@nvidia.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> (cherry picked from commit b893d83)
- Support seq_len_q>1 (linear-chain MTP / spec-decode) via the kernel's implicit causal mask in both the fp8 and fp16/bf16 decode kernels and the FMHA-library integration. - Gate CuteDslMlaFmha._is_supported_with_reason on the kernel's own can_implement under the default launch tiler, so the library only engages CuteDSL for problems the kernel can actually serve (the AutoTuner-less direct op path otherwise bypassed can_implement). - layer_wise_benchmarks runner: select CuteDSL via TLLM_FMHA_LIBS rather than the now-dead CUTEDSL attn_backend (always request TRTLLM). - test_attention_mla: revert the MLA_TEST_BACKEND env hook to the upstream hard-coded "TRTLLM" form. Signed-off-by: haow <haow@nvidia.com>
…shinfer NVIDIA#2743/NVIDIA#3309) - is_persistent = not is_var_seq: variable-seq paged decode now uses non-persistent tile scheduling instead of the hardcoded persistent path. - split-KV: fold-aware simplified heuristic (sq_eff via compute_fold_sq_ratio), CUDA-graph-safe split selection. - mla_helpers: non-persistent work-tile coordinate fix (b/s divmod + divisor), previously dormant because the path was never taken under hardcoded persistent. - moe_scheduler: attention-DP MoE-dispatch all-gather deadlock fix -- apply the empty-chunk substitution consistently across all ranks so the variable-size all-gather 'sizes' stay identical on every rank. - tests: cute_dsl_mla_decode updated (incl. fold_sq / small-batch split_kv). Signed-off-by: haow <haow@nvidia.com>
…et, cache-key fix, TRTLLM-Gen mixed warmup - fmha/cute_dsl.py: offset the per-layer page table by num_contexts so a mixed context+generation step reads the right generation pages (no-op for pure-decode); drop dev-only dump/param-log blocks - custom_ops/cute_dsl_custom_ops.py: drop split_kv from the kernel cache key (split_kv is now a dynamic cutlass.Int32, not baked into the compiled grid) - pyexecutor/model_engine.py: warm up the TRTLLM-Gen mixed context+generation FMHA variant so it does not JIT-compile at runtime when cute_dsl_mla declines a mixed step Signed-off-by: haow <haow@nvidia.com>
… whitelist gate, op schema cleanup - Make is_persistent the 4th AutoTuner tactic element (enumerated via get_is_persistent_candidates) instead of a fixed heuristic; batch-based heuristic remains only as the fallback default_tactic - Gate CuteDSL selection behind a perf whitelist of (num_heads, seq_len_q) combinations - Trim the custom-op schema from 19 to 14 params and drop integration-side debug env vars - Add autotune-warmup unit test coverage (test_cute_dsl_mla_decode) Signed-off-by: haow <haow@nvidia.com>
Signed-off-by: haow <haow@nvidia.com>
…rsistent, gate-aware test asserts - Strip stale design-history comments in fmha/cute_dsl.py and cute_dsl_custom_ops.py - Remove default_is_persistent heuristic; the untuned default tactic is always non-persistent (the AutoTuner picks the variant when it runs) - Revert the skip_context_assert test plumbing (test_attention_mla.py back to main); the decode test asserts context results too - Make the decode test's dispatch-count assertion follow the perf allowlist: admitted shapes must dispatch CuteDSL exactly layers*steps times, rejected shapes exactly 0 (fallback serves them) - Move the autotuned test geometry to num_heads=16 (allowlisted at seq_len_q=2) so the is_persistent tuning path is actually exercised Signed-off-by: haow <haow@nvidia.com>
…mment typos Signed-off-by: haow <haow@nvidia.com>
Signed-off-by: haow <haow@nvidia.com>
Signed-off-by: haow <haow@nvidia.com>
…nit params Signed-off-by: haow <haow@nvidia.com>
Signed-off-by: haow <haow@nvidia.com>
…mup, max-batch plumbing Signed-off-by: haow <haow@nvidia.com>
Signed-off-by: haow <haow@nvidia.com>
…spec-decode special-case Split _is_perf_favorable by kernel input dtype (the measured win regions differ): fp8 KV admits a (num_heads, seq_len_q) shape only above its critical batch size; bf16/fp16 KV admits only num_heads==16. Move the fp8 batch-floor table into the function as a local (_PERF_MIN_BATCH_FP8), and pass the resolved kernel dtype into the gate. Drop the (128, 1) + spec-decode rejection: the historical ADP+MTP3 ~-12% E2E regression was an autotuner/JIT miss on the untuned (128, sq=1) draft-step shape, not an intrinsic MTP loss. With the sq-axis autotuner warmup now covering that shape, a gate-off A/B on DeepSeek-V3 (ADP=1, MTP_draft_len=3, fp8 KV, 8xB200) measured +2.2% vs the flashinfer baseline, so the special-case now rejects a real win. Remove the TLLM_CUTE_DSL_DISABLE_PERF_GATE env escape hatch. Signed-off-by: haow <haow@nvidia.com>
Signed-off-by: haow <haow@nvidia.com>
…totuner free dim, trim comments Signed-off-by: haow <haow@nvidia.com>
…ce, trim gate Gate changes in `CuteDslMlaFmha`: - `is_supported()` now rejects `helix_position_offsets`; the kernel has no Helix position handling and would silently compute wrong results. - Collapse the sparse rejection to the module-level `attn.sparse_params` signal instead of also re-checking the per-forward sparse/topk index tensors, which are derived from it. - Drop the `kv_lora_rank` / `qk_rope_head_dim` positivity checks and the `num_tokens % num_generations` divisibility check, all implied by the MLA decode path that reaches this gate. - Drop the unused `predicted_tokens_per_seq` argument of the perf gate. Workspace: - Factor the required-size computation into `_required_workspace_size()` and reuse it from both `prepare_workspace()` and the forward path. - Slice `params.workspace` down to what this kernel actually owns before handing it to the op. The shared attention workspace is sized from `max_num_tokens` and can reach several GiB; the AutoTuner rebuilds every input with a dynamic dim as a float32 `torch.rand` tensor (4x the int8 byte count) while profiling, which could OOM. Docs and style: - `ATTENTION_DEVELOPER_GUIDE.md`: include `msa_sparse_gqa` in the documented default `TLLM_FMHA_LIBS` order. - `registry.py`: add the missing blank line before `init_fmha_libs()`. Signed-off-by: haow <haow@nvidia.com>
- Reject split_kv outside [1, 32] in can_implement() for both the FP8 and FP16 decode kernels; outside that range the kernel is slower than the reference path. - Raise the (num_heads=128, seq_len_q=1) perf-gate batch floor from 8 to 64. - Drop the duplicate sparse_params check in can_run(); sparse layers are already rejected in is_available(). Document why mixed context+generation batches stay disabled. - Widen the MLA unit-test context batches (max_num_contexts 10 -> 64, plus an explicit [10] * 64 case). - Drop the two attention __init__.py entries from the legacy lint allowlists, which are no longer legacy files. Signed-off-by: haow <haow@nvidia.com>
…_mla.py The module holds the MLA-only decode FMHA library, and the FMHA registry already exposes it under the key "cute_dsl_mla"; the file name now matches. This also disambiguates it from the unrelated CuTe DSL modules (custom_ops/cute_dsl_custom_ops.py, cute_dsl_utils.py, and the VisualGen attention_backend/cute_dsl package). Pure rename: the registry key, the TLLM_FMHA_LIBS token and the class name are unchanged, so no configuration or script needs updating. Imports in fmha/__init__.py and fmha/registry.py plus two path references in a comment and in ATTENTION_DEVELOPER_GUIDE.md follow the new name. Verified on B200: import resolves to the new module, the registry still returns CuteDslMlaFmha first for TLLM_FMHA_LIBS, and tests/unittest/_torch/attention/test_attention_mla.py passes 96/96. Signed-off-by: haow <haow@nvidia.com>
The perf gate was skipped wholesale while the AutoTuner was in tuning mode, so shapes the gate rejects on dtype, num_heads or seq_len_q grounds were still profiled and cached. Only the batch-size floor needs to be lifted during tuning (autotuner warmup issues gen requests at a single batch size, which the floor would reject, keeping the shape from ever being tuned). _is_perf_favorable now accepts batch_size=None to evaluate only the batch-size-independent conditions, and the caller passes None while tuning. Signed-off-by: haow <haow@nvidia.com>
…lowlist The measured fp8-KV win region for num_heads=128 only holds at seq_len_q 1 and 2; remove the (128,4) and (128,8) entries so those shapes fall back to the next FMHA library instead of being admitted above a batch floor. Signed-off-by: haow <haow@nvidia.com>
When the AutoTuner returns its -1 sentinel (cache miss at serving time), the op falls back to default_tactic, which derived split_kv from the raw runtime batch size. Tuning only ever profiles (and cute.compiles) the split_kv derived from each power-of-2 tuning bucket, so a raw-batch fallback almost always names a never-compiled kernel variant and JIT-compiles it inside the serving loop. Round the batch down to its tuning bucket (the same last_positive_power_of_2 mapping the tuning config uses) before deriving split_kv: a fallback on a tuned runner now reuses an already-compiled kernel, and on an untuned runner the number of distinct fallback variants is bounded by the bucket count instead of one per distinct batch size. The is_persistent choice is unchanged: its threshold (64) is a power of two, so rounding down to a power of two never crosses it. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
test_attention_mla runs with the autotuner off, so the CuTe DSL MLA decode op only ever exercises its default_tactic (-1) branch. Add a tuning-mode test on the fp8-KV DeepSeek decode geometry that asserts: - a tuning-mode pass profiles the op and both tactic elements the tuner owns (split_kv and both is_persistent candidates are compiled), and - a subsequent serving-mode pass reuses the tuned kernels with no new runtime cute.compile (which would stall the serving loop), while matching the reference output. The l0_b200 list already collects unittest/_torch/attention as a directory, so the new test runs in pre-merge B200 CI without a test-list change. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Collaborator
Author
|
/bot run |
A disaggregated generation server runs decode-only batches, so the decode-only CuTe DSL MLA lib takes essentially every forward there, yet no disagg test covered it and the only off switch (TLLM_FMHA_LIBS) is unset in every checked-in disagg config. Add one smoke: DeepSeek-V3-Lite bf16 on a ctxTP1+genTP2 cluster (gen TP2 yields the 16 heads/rank the bf16 path admits at any batch size), asserting client output and the lib's kernel-compile marker in a generation-worker log so a silent fallback to the next FMHA lib fails the test. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Per review discussion: CuTe DSL prefill FMHA/MLA kernels are planned for this same library, so use the general name now instead of renaming later. Registry key cute_dsl_mla becomes cute_dsl (TLLM_FMHA_LIBS spelling changes accordingly), the class becomes CuteDslFmha, and the file moves back to fmha/cute_dsl.py to keep matching the registered name. The kernel-specific custom op names (trtllm::cute_dsl_mla_decode_*) are unchanged. No dispatch behavior change. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Collaborator
|
PR_Github #63856 [ run ] triggered by Bot. Commit: |
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #63862 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63856 [ run ] completed with state |
brnguyen2
changed the title
3 tasks
Collaborator
|
PR_Github #63862 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #63911 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #63911 [ run ] completed with state
|
Collaborator
|
Default-on and decode-only means this takes nearly every attention call on a disaggregated generation server, and no stage on that hardware exercises it. Can we get coverage in before it's the default? |
Collaborator
Author
|
#15138 was merged — closing this mirror. |
brnguyen2
added a commit
to brnguyen2/TensorRT-LLM
that referenced
this pull request
Aug 5, 2026
Squash of the KimiLinear model integration on top of the kernel PRs (NVIDIA#17190, NVIDIA#17054, NVIDIA#17266, NVIDIA#17225): - KimiLinear model (modeling_kimi_k3) and KimiLinearConfig registration - Kimi K3 support modules: KDA mixer, K3 MoE, K3 MLA, fused attention-residual wrapper - K3 MLA module refactored onto the general MLA path (TRTLLM-14811) - fused_moe: SiTu activation and communication_method support - SiTu in the ActType_TrtllmGen python enum (python mirror of the C++ enum from the MoE kernel drop) - kda_decode: accept an optional out tensor (to be folded into NVIDIA#17054) - KDA kernel/runtime unit tests Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
brnguyen2
added a commit
to brnguyen2/TensorRT-LLM
that referenced
this pull request
Aug 5, 2026
Squash of the KimiLinear model integration on top of the kernel PRs (NVIDIA#17190, NVIDIA#17054, NVIDIA#17266, NVIDIA#17225): - KimiLinear model (modeling_kimi_k3) and KimiLinearConfig registration - Kimi K3 support modules: KDA mixer, K3 MoE, K3 MLA, fused attention-residual wrapper - K3 MLA module refactored onto the general MLA path (TRTLLM-14811) - fused_moe: SiTu activation and communication_method support - SiTu in the ActType_TrtllmGen python enum (python mirror of the C++ enum from the MoE kernel drop) - kda_decode: accept an optional out tensor (to be folded into NVIDIA#17054) - KDA kernel/runtime unit tests Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
brnguyen2
added a commit
to brnguyen2/TensorRT-LLM
that referenced
this pull request
Aug 5, 2026
Squash of the KimiLinear model integration on top of the kernel PRs (NVIDIA#17190, NVIDIA#17054, NVIDIA#17266, NVIDIA#17225): - KimiLinear model (modeling_kimi_k3) and KimiLinearConfig registration - Kimi K3 support modules: KDA mixer, K3 MoE, K3 MLA, fused attention-residual wrapper - K3 MLA module refactored onto the general MLA path (TRTLLM-14811) - fused_moe: SiTu activation and communication_method support - SiTu in the ActType_TrtllmGen python enum (python mirror of the C++ enum from the MoE kernel drop) - kda_decode: accept an optional out tensor (to be folded into NVIDIA#17054) - KDA kernel/runtime unit tests Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Mirror of #15138 (TRTLLM-12807): that PR's branch tip plus the fixes for its four open review threads, pushed to a fork so CI can keep running while the branch owner is unavailable. This draft tracks #15138 and will be closed when that PR is updated, merged, or closed. Not intended for review.
Commits on top of the #15138 tip (each thread's in-PR reply links back here):
[None][fix]CuteDSL MLA decode: bucket the fallback tactic's batch size — the AutoTuner-1fallback now reuses tuning-bucket kernel variants instead of JIT-compiling a raw-batchsplit_kvvariant in the serving loop.[None][test]cover the AutoTuner tuning path intest_attention_mla.py(tuning pass profiles both tactic elements; serving pass asserts no new runtimecute.compile).[None][test]disagg decode-only smoke (ctxTP1+genTP2 DeepSeek-V3-Lite bf16) asserting the lib actually takes the generation-server decode forwards via its kernel-compile log marker.[None][chore]renameCuteDslMlaFmha/cute_dsl_mlatoCuteDslFmha/cute_dsl— kept last so it can be dropped if the MLA-specific name should stay.All verified green on B200 (SM100): the new autotune test, the pre-existing cute_dsl-dispatching
test_attention_mlacases, and the new disagg smoke.