[None][feat] add CuteDSL FP8/FP16 MLA decode attention fmha lib - #15138
Conversation
0596464 to
e7b9ea6
Compare
haow-nv
changed the title
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds a Blackwell CuTe DSL MLA decode backend with static tile scheduling, FP8 and FP16/BF16 custom operations, FMHA registration, runtime gating, paged-KV execution, warmup integration, and lint configuration. Also fixes cross-rank empty-chunk handling in the MoE scheduler. ChangesCuTe DSL MLA Decode Backend
MoE Scheduler Fix
Estimated code review effort: 5 (Critical) | ~90 minutes Sequence Diagram(s)sequenceDiagram
participant FMHARegistry
participant CuteDslMlaFmha
participant MLADecodeCustomOp
participant CuteDSLNVMlaDecodeBlackwellRunner
participant CompiledMLAKernel
FMHARegistry->>CuteDslMlaFmha: select cute_dsl_mla backend
CuteDslMlaFmha->>CuteDslMlaFmha: validate support and prepare paged KV
CuteDslMlaFmha->>MLADecodeCustomOp: invoke FP8 or FP16/BF16 decode
MLADecodeCustomOp->>CuteDSLNVMlaDecodeBlackwellRunner: select tactic and launch
CuteDSLNVMlaDecodeBlackwellRunner->>CompiledMLAKernel: compile-cache and execute
CompiledMLAKernel->>CuteDslMlaFmha: write output and workspace
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py (1)
7666-7857: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: de-duplicate the FP8/FP16 dispatch bodies.
After the per-dtype
in_dtypeselection, both ops run an identical block (SM gate, runner construction,inputslist,choose_one,runner(...)), and the tworegister_fakestubs are byte-identical. Folding the shared portion into one helper keeps the FP8 and FP16 paths from diverging during future maintenance.♻️ Sketch of a shared dispatch helper
def _run_cute_dsl_mla_decode_blackwell( op_name, in_dtype, q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, block_split_kvs, o, lse, workspace, num_heads, seq_len_q, page_size, is_persistent, is_var_seq, is_var_split_kv, split_kv, softmax_scale, output_scale, ): runner = CuteDSLNVMlaDecodeBlackwellRunner( in_dtype=in_dtype, num_heads=num_heads, seq_len_q=seq_len_q, page_size=page_size, is_persistent=is_persistent, is_var_seq=is_var_seq, is_var_split_kv=is_var_split_kv, ) inputs = [q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, block_split_kvs, o, lse, workspace] _, best_tactic = AutoTuner.get().choose_one( op_name, [runner], runner.get_tuning_config(), inputs) runner(inputs, tactic=best_tactic, split_kv=split_kv, softmax_scale=softmax_scale, output_scale=output_scale)Each
@torch.library.custom_opthen keeps only its SM check + dtype resolution and delegates the rest. Theregister_fakebodies (return None) can share a single function passed to both registrations.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py` around lines 7666 - 7857, The FP8 and FP16 Blackwell MLA decode ops currently duplicate the same SM gate, runner setup, tuning, and execution flow in cute_dsl_mla_decode_fp8_blackwell and cute_dsl_mla_decode_fp16_blackwell, and the two register_fake stubs are identical. Extract the shared dispatch logic into a helper such as _run_cute_dsl_mla_decode_blackwell that takes op_name and in_dtype, then have each custom_op only handle its dtype validation/resolution before delegating. Reuse a single fake implementation for both `@torch.library.register_fake` registrations so the two paths stay aligned and easier to maintain.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py`:
- Around line 522-523: The page-table adjustment in the `cute_dsl.py` path that
folds in the interleaved layer slot is using only `layer_in_pool`, which can
point to the wrong KV page after `kv_pool` is re-viewed as packed combined
slots. Update the logic around `page_table` so the layer contribution is scaled
by `layers_in_pool` (matching the logical mapping used by the packed view) and
keep the behavior in the same `layers_in_pool > 1` branch.
- Around line 539-644: The split-KV selection in cute_dsl.py can still enable
is_var_split_kv even when fixed-sequence mode is forced, which creates an
invalid kernel configuration. In the decode path around
_split_kv_from_max_splits, _compute_block_split_kvs, and the is_var_seq toggle,
gate the variable split-KV logic on is_var_seq being true, and force split_kv
back to 1 when TLLM_CUTE_DSL_VAR_SEQ=0 so is_var_split_kv can never be set in
fixed-sequence mode. Ensure the workspace allocation and downstream kernel
arguments follow the adjusted split_kv value.
In `@tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py`:
- Around line 82-92: The FP16 decode path is still advertised by
CuteDslMlaFmha._get_kernel_dtype() and routes to
cute_dsl_mla_decode_fp16_blackwell, but the test matrix in
test_cute_dsl_mla_decode.py now omits torch.float16 and hides the crash. Update
the runtime in tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py to gate
FP16 out until it is fixed, or add an explicit xfail/regression test that
exercises the FP16 path and documents the aborting behavior. Keep the test
coverage aligned with the actual supported kernel dtypes so the unsupported FP16
path is not silently masked.
---
Nitpick comments:
In `@tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py`:
- Around line 7666-7857: The FP8 and FP16 Blackwell MLA decode ops currently
duplicate the same SM gate, runner setup, tuning, and execution flow in
cute_dsl_mla_decode_fp8_blackwell and cute_dsl_mla_decode_fp16_blackwell, and
the two register_fake stubs are identical. Extract the shared dispatch logic
into a helper such as _run_cute_dsl_mla_decode_blackwell that takes op_name and
in_dtype, then have each custom_op only handle its dtype validation/resolution
before delegating. Reuse a single fake implementation for both
`@torch.library.register_fake` registrations so the two paths stay aligned and
easier to maintain.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e1c74753-bc97-42cb-8418-7ace8aa2ef65
📒 Files selected for processing (19)
.pre-commit-config.yamllegacy-files.txtpyproject.tomlruff-legacy.tomltensorrt_llm/_torch/attention_backend/fmha/__init__.pytensorrt_llm/_torch/attention_backend/fmha/cute_dsl.pytensorrt_llm/_torch/attention_backend/fmha/registry.pytensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/__init__.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/__init__.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp16.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_helpers.pytensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/modules/fused_moe/moe_scheduler.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/tools/layer_wise_benchmarks/runner.pytests/unittest/_torch/attention/test_attention_mla.pytests/unittest/_torch/attention/test_cute_dsl_mla_decode.py
fc530d6 to
cf75689
Compare
|
/bot run --disable-fail-fast --add-multi-gpu-test |
bd40e10 to
5dceec8
Compare
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
/bot run |
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #57746 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast --add-multi-gpu-test |
…_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>
cf8dc0b to
b2f3b28
Compare
b2f3b28 to
faf5aba
Compare
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #63676 [ run ] triggered by Bot. Commit: |
yuanjingx87
left a comment
yuanjingx87
left a comment
There was a problem hiding this comment.
Approved on oss compliance perspective
This comment was marked as duplicate.
This comment was marked as duplicate.
Sorry, something went wrong.
|
PR_Github #63676 [ run ] completed with state
|
…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>
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #63737 [ run ] triggered by Bot. Commit: |
|
PR_Github #63737 [ run ] completed with state
|
|
/bot run --add-multi-gpu-test |
|
PR_Github #63922 [ run ] triggered by Bot. Commit: |
|
PR_Github #63922 [ run ] completed with state
|
|
/bot run --add-multi-gpu-test |
|
PR_Github #63955 [ run ] triggered by Bot. Commit: |
|
PR_Github #63955 [ run ] completed with state |
|
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? |
Hi @Shixiaowei02 , I contact you through slack. It seems that I don't find the correct people. Could you ping me? Do we have testcases for disaggregated generation server? I think the cutedsl mla is naturally covered by those cases? |
11 participants
Description
Adds CuteDSL MLA decode as an internal FMHA library of the
TRTLLMattention backend (not a separateattn_backend). The newCuteDslMlaFmhalibrary intercepts the generation-phase (decode) MLA portion of a batch and dispatches it to Blackwell CuTe DSL kernels; context-phase requests and anything the library rejects fall through to the next library in the ordered list (flashinfer_trtllm_gen, thenfallback).Selection & gating
cute_dsl_mlain the FMHA library registry. The default order iscute_dsl_mla,flashinfer_trtllm_gen,fallback; override withTLLM_FMHA_LIBS(exact list, e.g.TLLM_FMHA_LIBS=flashinfer_trtllm_gen,fallback, or deltas, e.g.TLLM_FMHA_LIBS=-cute_dsl_mla).is_available()(static): SM100/103 +nvidia-cutlass-dslinstalled + MLA withkv_lora_rank=512,qk_rope_head_dim=64,num_heads <= 128(DeepSeek geometry).can_implementcheck, plus a perf allowlist — the library only takes(num_heads, seq_len_q)in{(16, 2), (16, 4), (128, 1)}, the shapes where CuteDSL measured as an end-to-end win over TRTLLM-Gen in DeepSeek-V3 A/B runs (TP8 + MTP, and attention-DP without MTP). All other shapes fall back to TRTLLM-Gen with a debug log.Main pieces
attention_backend/fmha/cute_dsl.py:CuteDslMlaFmha(PhasedFmha)— generation-phase hook that prepares latent/rope Q views, paged-KV page table, and per-request cache lengths, then calls the CuteDSL decode op. FP8 KV cache routes to the FP8 kernel; FP16/BF16 activations route to the FP16/BF16 kernel.cute_dsl_kernels/blackwell/attention/mla/: vendored Blackwell CuTe DSL MLA decode kernels (FP8 and FP16/BF16) + helpers, adapted from the CUTLASS MLA reference and aligned with FlashInfer'smla_decode(persistent and split-KV modes,seq_len_q > 1causal masking with fold-sq for MTP).custom_ops/cute_dsl_custom_ops.py: registerstrtllm::cute_dsl_mla_decode_{fp8,fp16}_blackwelland the kernel runner (JIT compile/cache).split_kvandis_persistentare AutoTuner tactic elements chosen per shape (power-of-two split-KV candidate ladder, batch-bucketed tuning config);o/lse/workspaceare mutated in place, and the workspace is sized batch-independently so CUDA-graph capture across batch sizes stays valid.pyexecutor/model_engine.py: adds one mixed context+generation warmup request so TRTLLM-Gen kernels are JIT-warmed for the mixed batches thatcute_dsl_mladoes not take (it is decode-only).modules/fused_moe/moe_scheduler.py: attention-DP empty-chunk substitution fix — the substitution now updates every rank's slot in the collective sizes (emptiness is deterministic from the cross-rank-consistentall_rank_num_tokens_list), keeping the variable-size all-gather consistent across ranks. Previously each rank only patched its own slot, which deadlocked warmup with imbalanced attention-DP batches.modules/ATTENTION_DEVELOPER_GUIDE.mdand lint-exclusion lists (.pre-commit-config.yaml,legacy-files.txt,pyproject.toml,ruff-legacy.toml) updated for the new library and vendored kernel files.MTP (
seq_len_q > 1) and CUDA graphs are supported on the CuteDSL path.Test Coverage
Validated locally on B200 (SM100); both suites are gated on
get_sm_version() in (100, 103)andIS_CUTLASS_DSL_AVAILABLE.tests/unittest/_torch/attention/test_attention_mla.py(existing): full MLA module dispatch passes withcute_dsl_mlain the default library order.End-to-end: DeepSeek-V3 FP8 on 8x B200 (TP8/EP8, CUDA graphs on) runs green with and without attention-DP and MTP — no fallbacks, correct outputs, MTP acceptance rate identical to the TRTLLM-Gen baseline; throughput on allowlisted shapes is at parity or better, which is what the perf allowlist encodes.
PR Checklist
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Dev Engineer Review
CuteDslMlaFmha) under thecute_dsl_mlaregistry key and exported it from the FMHA package for dispatch of eligible DeepSeek MLA generation-phase decode requests to CuTe DSL kernels (FP8 and FP16/BF16).tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py, including CuTe/SM100/103 availability and tight geometry/mode constraints (decode-only, specific lora/head-dim expectations, head-count limits, and rejection of unsupported speculative/beam/tree/custom mask modes).CuteDSLNVMlaDecodeBlackwellRunner) intensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py:can_implement,tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/:mla_decode_fp8.py,mla_decode_fp16.py,mla_helpers.py(static tile scheduling helpers for persistent/non-persistent modes),__init__.pyexports.cute_dsl_mlain the default/ordering guidance (ATTENTION_DEVELOPER_GUIDE.md).ExternalCommMoEScheduler._forward_multiple_chunks(tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py) by ensuring all ranks use identicalsizesvectors whensplit_chunkyields 0-token chunks on some ranks.tensorrt_llm/_torch/pyexecutor/model_engine.pyto enable a mixed context+generation warmup only under the specified compatibility conditions..pre-commit-config.yaml,legacy-files.txt,pyproject.toml,ruff-legacy.toml).Review notes / risk areas
can_implement+ tiler selection), workspace sizing/slicing (including split-K paths), and persistent vs non-persistent static tile scheduling.QA Engineer Review
No test changes.