[None][feat] add CuteDSL FP8/FP16 MLA decode attention fmha lib#15138
[None][feat] add CuteDSL FP8/FP16 MLA decode attention fmha lib#15138haow-nv wants to merge 21 commits into
Conversation
0596464 to
e7b9ea6
Compare
haow-nv
changed the title
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
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, LSE, and workspace
Possibly related PRs
Suggested labels: 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
| if layers_in_pool > 1: | ||
| page_table = page_table + layer_in_pool |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick win
Scale page IDs when folding in the interleaved layer slot.
After kv_pool is re-viewed as packed combined slots, logical page p for layer slot l maps to p * layers_in_pool + l; adding only layer_in_pool can read KV pages from the wrong layer/page.
🐛 Proposed fix
page_table = page_table_layer[meta.num_contexts:].transpose(0, 1).to(torch.int32)
if layers_in_pool > 1:
- page_table = page_table + layer_in_pool
+ page_table = page_table * layers_in_pool + layer_in_pool📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if layers_in_pool > 1: | |
| page_table = page_table + layer_in_pool | |
| page_table = page_table_layer[meta.num_contexts:].transpose(0, 1).to(torch.int32) | |
| if layers_in_pool > 1: | |
| page_table = page_table * layers_in_pool + layer_in_pool |
🤖 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/attention_backend/fmha/cute_dsl.py` around lines 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.
| if os.environ.get("TLLM_CUTE_DSL_VAR_SPLIT_KV", "1") != "0": | ||
| max_active_blocks = self._get_max_active_blocks() | ||
| # Host upper bound on KV length: per-sequence page capacity * page | ||
| # size (page_table is [pages_per_seq, batch], a fixed shape under | ||
| # CUDA-graph capture). Yields the grid's split_kv max on the host. | ||
| k_max = page_table.shape[0] * page_size | ||
| max_splits = max(1, (k_max + self._CUTE_DSL_QK_TILE_K - 1) // self._CUTE_DSL_QK_TILE_K) | ||
| split_kv = self._split_kv_from_max_splits( | ||
| max_splits, batch_size, seq_len_q, max_active_blocks | ||
| ) | ||
| if split_kv > 1: | ||
| is_var_split_kv = True | ||
| block_split_kvs = self._compute_block_split_kvs( | ||
| cache_seqs_base, batch_size, seq_len_q, max_active_blocks, split_kv | ||
| ) | ||
| # get_workspace_size = B*H*S*split_kv*(D+1)*acc_width//8; fold | ||
| # cancels (H_eff*S_eff == H*S), acc=fp32 (width 32 -> //8 = 4). | ||
| ws_bytes = batch_size * num_heads * seq_len_q * split_kv * (d_latent + 1) * 4 | ||
| workspace = torch.empty(ws_bytes, dtype=torch.int8, device=q.device) | ||
| else: | ||
| split_kv = 1 | ||
|
|
||
| softmax_scale = float(1.0 / (math.sqrt(qk_nope_head_dim + d_rope) * attn.q_scaling)) | ||
| output_scale = 1.0 | ||
| if kernel_dtype == torch.float8_e4m3fn: | ||
| if params.fwd.mla_bmm1_scale is None or params.fwd.mla_bmm2_scale is None: | ||
| raise RuntimeError("FP8 CuTe DSL MLA decode requires MLA FP8 scales.") | ||
| cached = getattr(self, "_cute_dsl_fp8_scale", None) | ||
| if cached is None: | ||
| if torch.cuda.is_current_stream_capturing(): | ||
| raise RuntimeError( | ||
| "CuTe DSL MLA FMHA: fp8 decode scale was not cached for " | ||
| f"layer {attn.layer_idx} before CUDA graph capture." | ||
| ) | ||
| softmax_scale = float(params.fwd.mla_bmm1_scale[1].item()) / _LOG2_E | ||
| output_scale = float(params.fwd.mla_bmm2_scale[0].item()) | ||
| self._cute_dsl_fp8_scale = (softmax_scale, output_scale) | ||
| else: | ||
| softmax_scale, output_scale = cached | ||
| if ( | ||
| os.environ.get("TLLM_CUTE_DSL_SCALE_DUMP") | ||
| and not torch.cuda.is_current_stream_capturing() | ||
| ): | ||
| live_softmax_scale = float(params.fwd.mla_bmm1_scale[1].item()) / _LOG2_E | ||
| live_output_scale = float(params.fwd.mla_bmm2_scale[0].item()) | ||
| print( | ||
| "[CUTEDSL_SCALE] layer=%d cached=(%.8f,%.8f) " | ||
| "live=(%.8f,%.8f) drift=%s" | ||
| % ( | ||
| attn.layer_idx, | ||
| softmax_scale, | ||
| output_scale, | ||
| live_softmax_scale, | ||
| live_output_scale, | ||
| abs(live_softmax_scale - softmax_scale) > 1e-6 | ||
| or abs(live_output_scale - output_scale) > 1e-6, | ||
| ), | ||
| flush=True, | ||
| ) | ||
|
|
||
| out_kernel_dtype = torch.bfloat16 if kernel_dtype == torch.float8_e4m3fn else kernel_dtype | ||
| output_view = output.view(batch_size, seq_len_q, num_heads, d_latent) | ||
|
|
||
| # Single fused decode over all ``seq_len_q`` query tokens. For | ||
| # multi-query (MTP / linear spec-decode) the kernel applies the causal | ||
| # mask internally: query token ``t`` attends to KV positions | ||
| # ``[0, K - (seq_len_q - 1) + t)``. ``cache_seqs_base`` already counts | ||
| # every freshly-appended token of this step (K), so token ``t``'s bound | ||
| # equals ``cache_seqs_base - (seq_len_q - 1) + t`` -- exactly the | ||
| # per-query trim the previous one-token-at-a-time loop applied. For | ||
| # ``seq_len_q == 1`` this reduces to a plain decode. | ||
| q_latent = q_view[..., :d_latent].permute(2, 3, 1, 0) | ||
| q_rope = q_view[..., d_latent:].permute(2, 3, 1, 0) | ||
|
|
||
| # When the kernel output dtype matches the module output dtype (the | ||
| # common fp8-KV case: both bf16), have the kernel write straight into | ||
| # ``output`` instead of a temp buffer that is then D2D-copied back. That | ||
| # copy was ~1.7us/call and, at small batch where the decode win is only | ||
| # a few us, ate the win (MLA module went flat/slightly slower despite a | ||
| # faster decode kernel). ``output_view`` is a contiguous | ||
| # [B, S_q, H, d_latent] view, so its permute(2,3,1,0) is byte-identical | ||
| # in layout to a fresh contiguous o_storage's -- the op's compact-shape | ||
| # marking still holds. Only fall back to the temp+copy on a dtype | ||
| # mismatch (the kernel can only emit out_kernel_dtype). | ||
| write_output_direct = output.dtype == out_kernel_dtype | ||
| if write_output_direct: | ||
| o_storage = output_view | ||
| else: | ||
| o_storage = torch.empty( | ||
| (batch_size, seq_len_q, num_heads, d_latent), | ||
| dtype=out_kernel_dtype, | ||
| device=q.device, | ||
| ) | ||
| o_kernel = o_storage.permute(2, 3, 1, 0) | ||
| lse_storage = torch.empty( | ||
| (batch_size, seq_len_q, num_heads), | ||
| dtype=torch.float32, | ||
| device=q.device, | ||
| ) | ||
| lse = lse_storage.permute(2, 1, 0) | ||
|
|
||
| # The decode path uses variable-seq mode by default (real serving has | ||
| # unequal per-request KV lengths). Experiment toggle: set | ||
| # TLLM_CUTE_DSL_VAR_SEQ=0 to force the fixed-length path -- only valid | ||
| # when every sequence shares one KV length (e.g. profiling/microbench). | ||
| is_var_seq = os.environ.get("TLLM_CUTE_DSL_VAR_SEQ", "1") != "0" |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Disable variable split-KV when fixed-sequence mode is forced.
The kernel rejects is_var_split_kv=True with is_var_seq=False; with TLLM_CUTE_DSL_VAR_SEQ=0, the current default split-KV path can still set that invalid combination.
🐛 Proposed fix
+ # The decode path uses variable-seq mode by default (real serving has
+ # unequal per-request KV lengths). Experiment toggle: set
+ # TLLM_CUTE_DSL_VAR_SEQ=0 to force the fixed-length path -- only valid
+ # when every sequence shares one KV length (e.g. profiling/microbench).
+ is_var_seq = os.environ.get("TLLM_CUTE_DSL_VAR_SEQ", "1") != "0"
+
if os.environ.get("TLLM_CUTE_DSL_VAR_SPLIT_KV", "1") != "0":
+ if not is_var_seq:
+ split_kv = 1
+ else:
- max_active_blocks = self._get_max_active_blocks()
+ max_active_blocks = self._get_max_active_blocks()
...
-
- # The decode path uses variable-seq mode by default (real serving has
- # unequal per-request KV lengths). Experiment toggle: set
- # TLLM_CUTE_DSL_VAR_SEQ=0 to force the fixed-length path -- only valid
- # when every sequence shares one KV length (e.g. profiling/microbench).
- is_var_seq = os.environ.get("TLLM_CUTE_DSL_VAR_SEQ", "1") != "0"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if os.environ.get("TLLM_CUTE_DSL_VAR_SPLIT_KV", "1") != "0": | |
| max_active_blocks = self._get_max_active_blocks() | |
| # Host upper bound on KV length: per-sequence page capacity * page | |
| # size (page_table is [pages_per_seq, batch], a fixed shape under | |
| # CUDA-graph capture). Yields the grid's split_kv max on the host. | |
| k_max = page_table.shape[0] * page_size | |
| max_splits = max(1, (k_max + self._CUTE_DSL_QK_TILE_K - 1) // self._CUTE_DSL_QK_TILE_K) | |
| split_kv = self._split_kv_from_max_splits( | |
| max_splits, batch_size, seq_len_q, max_active_blocks | |
| ) | |
| if split_kv > 1: | |
| is_var_split_kv = True | |
| block_split_kvs = self._compute_block_split_kvs( | |
| cache_seqs_base, batch_size, seq_len_q, max_active_blocks, split_kv | |
| ) | |
| # get_workspace_size = B*H*S*split_kv*(D+1)*acc_width//8; fold | |
| # cancels (H_eff*S_eff == H*S), acc=fp32 (width 32 -> //8 = 4). | |
| ws_bytes = batch_size * num_heads * seq_len_q * split_kv * (d_latent + 1) * 4 | |
| workspace = torch.empty(ws_bytes, dtype=torch.int8, device=q.device) | |
| else: | |
| split_kv = 1 | |
| softmax_scale = float(1.0 / (math.sqrt(qk_nope_head_dim + d_rope) * attn.q_scaling)) | |
| output_scale = 1.0 | |
| if kernel_dtype == torch.float8_e4m3fn: | |
| if params.fwd.mla_bmm1_scale is None or params.fwd.mla_bmm2_scale is None: | |
| raise RuntimeError("FP8 CuTe DSL MLA decode requires MLA FP8 scales.") | |
| cached = getattr(self, "_cute_dsl_fp8_scale", None) | |
| if cached is None: | |
| if torch.cuda.is_current_stream_capturing(): | |
| raise RuntimeError( | |
| "CuTe DSL MLA FMHA: fp8 decode scale was not cached for " | |
| f"layer {attn.layer_idx} before CUDA graph capture." | |
| ) | |
| softmax_scale = float(params.fwd.mla_bmm1_scale[1].item()) / _LOG2_E | |
| output_scale = float(params.fwd.mla_bmm2_scale[0].item()) | |
| self._cute_dsl_fp8_scale = (softmax_scale, output_scale) | |
| else: | |
| softmax_scale, output_scale = cached | |
| if ( | |
| os.environ.get("TLLM_CUTE_DSL_SCALE_DUMP") | |
| and not torch.cuda.is_current_stream_capturing() | |
| ): | |
| live_softmax_scale = float(params.fwd.mla_bmm1_scale[1].item()) / _LOG2_E | |
| live_output_scale = float(params.fwd.mla_bmm2_scale[0].item()) | |
| print( | |
| "[CUTEDSL_SCALE] layer=%d cached=(%.8f,%.8f) " | |
| "live=(%.8f,%.8f) drift=%s" | |
| % ( | |
| attn.layer_idx, | |
| softmax_scale, | |
| output_scale, | |
| live_softmax_scale, | |
| live_output_scale, | |
| abs(live_softmax_scale - softmax_scale) > 1e-6 | |
| or abs(live_output_scale - output_scale) > 1e-6, | |
| ), | |
| flush=True, | |
| ) | |
| out_kernel_dtype = torch.bfloat16 if kernel_dtype == torch.float8_e4m3fn else kernel_dtype | |
| output_view = output.view(batch_size, seq_len_q, num_heads, d_latent) | |
| # Single fused decode over all ``seq_len_q`` query tokens. For | |
| # multi-query (MTP / linear spec-decode) the kernel applies the causal | |
| # mask internally: query token ``t`` attends to KV positions | |
| # ``[0, K - (seq_len_q - 1) + t)``. ``cache_seqs_base`` already counts | |
| # every freshly-appended token of this step (K), so token ``t``'s bound | |
| # equals ``cache_seqs_base - (seq_len_q - 1) + t`` -- exactly the | |
| # per-query trim the previous one-token-at-a-time loop applied. For | |
| # ``seq_len_q == 1`` this reduces to a plain decode. | |
| q_latent = q_view[..., :d_latent].permute(2, 3, 1, 0) | |
| q_rope = q_view[..., d_latent:].permute(2, 3, 1, 0) | |
| # When the kernel output dtype matches the module output dtype (the | |
| # common fp8-KV case: both bf16), have the kernel write straight into | |
| # ``output`` instead of a temp buffer that is then D2D-copied back. That | |
| # copy was ~1.7us/call and, at small batch where the decode win is only | |
| # a few us, ate the win (MLA module went flat/slightly slower despite a | |
| # faster decode kernel). ``output_view`` is a contiguous | |
| # [B, S_q, H, d_latent] view, so its permute(2,3,1,0) is byte-identical | |
| # in layout to a fresh contiguous o_storage's -- the op's compact-shape | |
| # marking still holds. Only fall back to the temp+copy on a dtype | |
| # mismatch (the kernel can only emit out_kernel_dtype). | |
| write_output_direct = output.dtype == out_kernel_dtype | |
| if write_output_direct: | |
| o_storage = output_view | |
| else: | |
| o_storage = torch.empty( | |
| (batch_size, seq_len_q, num_heads, d_latent), | |
| dtype=out_kernel_dtype, | |
| device=q.device, | |
| ) | |
| o_kernel = o_storage.permute(2, 3, 1, 0) | |
| lse_storage = torch.empty( | |
| (batch_size, seq_len_q, num_heads), | |
| dtype=torch.float32, | |
| device=q.device, | |
| ) | |
| lse = lse_storage.permute(2, 1, 0) | |
| # The decode path uses variable-seq mode by default (real serving has | |
| # unequal per-request KV lengths). Experiment toggle: set | |
| # TLLM_CUTE_DSL_VAR_SEQ=0 to force the fixed-length path -- only valid | |
| # when every sequence shares one KV length (e.g. profiling/microbench). | |
| is_var_seq = os.environ.get("TLLM_CUTE_DSL_VAR_SEQ", "1") != "0" | |
| # The decode path uses variable-seq mode by default (real serving has | |
| # unequal per-request KV lengths). Experiment toggle: set | |
| # TLLM_CUTE_DSL_VAR_SEQ=0 to force the fixed-length path -- only valid | |
| # when every sequence shares one KV length (e.g. profiling/microbench). | |
| is_var_seq = os.environ.get("TLLM_CUTE_DSL_VAR_SEQ", "1") != "0" | |
| if os.environ.get("TLLM_CUTE_DSL_VAR_SPLIT_KV", "1") != "0": | |
| if not is_var_seq: | |
| split_kv = 1 | |
| else: | |
| max_active_blocks = self._get_max_active_blocks() | |
| # Host upper bound on KV length: per-sequence page capacity * page | |
| # size (page_table is [pages_per_seq, batch], a fixed shape under | |
| # CUDA-graph capture). Yields the grid's split_kv max on the host. | |
| k_max = page_table.shape[0] * page_size | |
| max_splits = max(1, (k_max + self._CUTE_DSL_QK_TILE_K - 1) // self._CUTE_DSL_QK_TILE_K) | |
| split_kv = self._split_kv_from_max_splits( | |
| max_splits, batch_size, seq_len_q, max_active_blocks | |
| ) | |
| if split_kv > 1: | |
| is_var_split_kv = True | |
| block_split_kvs = self._compute_block_split_kvs( | |
| cache_seqs_base, batch_size, seq_len_q, max_active_blocks, split_kv | |
| ) | |
| # get_workspace_size = B*H*S*split_kv*(D+1)*acc_width//8; fold | |
| # cancels (H_eff*S_eff == H*S), acc=fp32 (width 32 -> //8 = 4). | |
| ws_bytes = batch_size * num_heads * seq_len_q * split_kv * (d_latent + 1) * 4 | |
| workspace = torch.empty(ws_bytes, dtype=torch.int8, device=q.device) | |
| else: | |
| split_kv = 1 | |
| softmax_scale = float(1.0 / (math.sqrt(qk_nope_head_dim + d_rope) * attn.q_scaling)) | |
| output_scale = 1.0 | |
| if kernel_dtype == torch.float8_e4m3fn: | |
| if params.fwd.mla_bmm1_scale is None or params.fwd.mla_bmm2_scale is None: | |
| raise RuntimeError("FP8 CuTe DSL MLA decode requires MLA FP8 scales.") | |
| cached = getattr(self, "_cute_dsl_fp8_scale", None) | |
| if cached is None: | |
| if torch.cuda.is_current_stream_capturing(): | |
| raise RuntimeError( | |
| "CuTe DSL MLA FMHA: fp8 decode scale was not cached for " | |
| f"layer {attn.layer_idx} before CUDA graph capture." | |
| ) | |
| softmax_scale = float(params.fwd.mla_bmm1_scale[1].item()) / _LOG2_E | |
| output_scale = float(params.fwd.mla_bmm2_scale[0].item()) | |
| self._cute_dsl_fp8_scale = (softmax_scale, output_scale) | |
| else: | |
| softmax_scale, output_scale = cached | |
| if ( | |
| os.environ.get("TLLM_CUTE_DSL_SCALE_DUMP") | |
| and not torch.cuda.is_current_stream_capturing() | |
| ): | |
| live_softmax_scale = float(params.fwd.mla_bmm1_scale[1].item()) / _LOG2_E | |
| live_output_scale = float(params.fwd.mla_bmm2_scale[0].item()) | |
| print( | |
| "[CUTEDSL_SCALE] layer=%d cached=(%.8f,%.8f) " | |
| "live=(%.8f,%.8f) drift=%s" | |
| % ( | |
| attn.layer_idx, | |
| softmax_scale, | |
| output_scale, | |
| live_softmax_scale, | |
| live_output_scale, | |
| abs(live_softmax_scale - softmax_scale) > 1e-6 | |
| or abs(live_output_scale - output_scale) > 1e-6, | |
| ), | |
| flush=True, | |
| ) | |
| out_kernel_dtype = torch.bfloat16 if kernel_dtype == torch.float8_e4m3fn else kernel_dtype | |
| output_view = output.view(batch_size, seq_len_q, num_heads, d_latent) | |
| # Single fused decode over all ``seq_len_q`` query tokens. For | |
| # multi-query (MTP / linear spec-decode) the kernel applies the causal | |
| # mask internally: query token ``t`` attends to KV positions | |
| # ``[0, K - (seq_len_q - 1) + t)``. ``cache_seqs_base`` already counts | |
| # every freshly-appended token of this step (K), so token ``t``'s bound | |
| # equals ``cache_seqs_base - (seq_len_q - 1) + t`` -- exactly the | |
| # per-query trim the previous one-token-at-a-time loop applied. For | |
| # ``seq_len_q == 1`` this reduces to a plain decode. | |
| q_latent = q_view[..., :d_latent].permute(2, 3, 1, 0) | |
| q_rope = q_view[..., d_latent:].permute(2, 3, 1, 0) | |
| # When the kernel output dtype matches the module output dtype (the | |
| # common fp8-KV case: both bf16), have the kernel write straight into | |
| # ``output`` instead of a temp buffer that is then D2D-copied back. That | |
| # copy was ~1.7us/call and, at small batch where the decode win is only | |
| # a few us, ate the win (MLA module went flat/slightly slower despite a | |
| # faster decode kernel). ``output_view`` is a contiguous | |
| # [B, S_q, H, d_latent] view, so its permute(2,3,1,0) is byte-identical | |
| # in layout to a fresh contiguous o_storage's -- the op's compact-shape | |
| # marking still holds. Only fall back to the temp+copy on a dtype | |
| # mismatch (the kernel can only emit out_kernel_dtype). | |
| write_output_direct = output.dtype == out_kernel_dtype | |
| if write_output_direct: | |
| o_storage = output_view | |
| else: | |
| o_storage = torch.empty( | |
| (batch_size, seq_len_q, num_heads, d_latent), | |
| dtype=out_kernel_dtype, | |
| device=q.device, | |
| ) | |
| o_kernel = o_storage.permute(2, 3, 1, 0) | |
| lse_storage = torch.empty( | |
| (batch_size, seq_len_q, num_heads), | |
| dtype=torch.float32, | |
| device=q.device, | |
| ) | |
| lse = lse_storage.permute(2, 1, 0) |
🤖 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/attention_backend/fmha/cute_dsl.py` around lines 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.
| # kernel name -> (activation dtype, kv cache dtype) | ||
| # | ||
| # NOTE: the float16 instance of the FP16 decode op | ||
| # (``cute_dsl_mla_decode_fp16_blackwell`` with dtype=float16 / fp16 KV cache) | ||
| # aborts the process (SIGABRT) on SM100 in this environment, so that exact | ||
| # dtype is excluded until the crash is root-caused. The bf16 instance uses the | ||
| # same op and is covered below for DeepSeek-V3 bf16 runs. | ||
| _KERNEL_DTYPES = { | ||
| "fp8": (torch.bfloat16, torch.float8_e4m3fn), | ||
| "bf16": (torch.bfloat16, torch.bfloat16), | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Don't mask the FP16 crash while the runtime still advertises FP16 support.
This test drops the torch.float16 case, but CuteDslMlaFmha._get_kernel_dtype() still routes real FP16 decode batches into cute_dsl_mla_decode_fp16_blackwell. So the suite is currently hiding a supported production path that is documented here as aborting the process. Either gate FP16 out in tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py until it is fixed, or keep an explicit xfail-style regression that proves the unsupported behavior. Based on path instructions, coverage in tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py is insufficient for the advertised FP16 path and needs follow-up outside this PR.
🤖 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 `@tests/unittest/_torch/attention/test_cute_dsl_mla_decode.py` around lines 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.
Source: Path instructions
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 |
|
PR_Github #61265 [ run ] triggered by Bot. Commit: |
|
PR_Github #60959 [ run ] completed with state |
|
PR_Github #61268 [ run ] triggered by Bot. Commit: |
|
PR_Github #61265 [ run ] completed with state |
|
PR_Github #61268 [ run ] completed with state
|
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #61502 [ run ] triggered by Bot. Commit: |
|
/bot run --add-multi-gpu-test |
|
PR_Github #61509 [ run ] triggered by Bot. Commit: |
|
PR_Github #61502 [ run ] completed with state |
| if (not self.is_draft_model and self.guided_decoder is None | ||
| and not self.mapping.has_pp()): | ||
| # Add generation request to warmup the autotuner cache. | ||
| warmup_configs.append((1 + self.max_total_draft_tokens, 1)) |
There was a problem hiding this comment.
The context-only case uses (curr_max_num_tokens, 0), to avoid violating curr_max_num_tokens, we can use (curr_max_num_tokens, 1) here.
There was a problem hiding this comment.
make the file name match with registered name, how about using cute_dsl_mla.py?
There was a problem hiding this comment.
We need to integrate cutedsl fmha/mla prefill kernel latter. Shall we update the fmha lib name from "CuteDslMlaFmha" to "CuteDslFmha" here?
There was a problem hiding this comment.
If we want to use a single fmha lib to integrate both mla and non-mla kernels, we can use CuteDslFmha and cute_dsl as registration name.
| MAX_SPLITS = 256 | ||
|
|
||
|
|
||
| def ceil_div(a: int, b: int) -> int: |
There was a problem hiding this comment.
Please use the ceil_div in tensorrt_llm/math_utils.py instead of re-refining it.
There was a problem hiding this comment.
For these 5 new files, we can apply new ruff fules for them. We can revert all changes in .pre-commit-config.yaml, legacy-files.txt, pyproject.toml and ruff-legacy.toml.
There was a problem hiding this comment.
Thanks. Will revert all changes.
| if attn.qk_nope_head_dim is None or attn.qk_nope_head_dim <= 0: | ||
| logger.debug("CuTe DSL MLA FMHA is unavailable: qk_nope_head_dim must be positive.") | ||
| return False | ||
| if attn.kv_lora_rank != 512 or attn.qk_rope_head_dim != 64: |
There was a problem hiding this comment.
This is stricter condition, don't need to check attn.kv_lora_rank <= 0 and attn.qk_rope_head_dim <= 0 before
| ) -> tuple[bool, str]: | ||
| if fwd.attention_input_type != AttentionInputType.generation_only: | ||
| return False, "CuTe DSL MLA FMHA only supports generation-only attention." | ||
| if meta.num_contexts != 0 or meta.num_generations <= 0: |
There was a problem hiding this comment.
This condition seems to be equivalent to fwd.attention_input_type != AttentionInputType.generation_only, can we remove it?
| or meta.num_sparse_topk > 0 | ||
| or attn.sparse_params is not None | ||
| ): | ||
| return False, "CuTe DSL MLA FMHA does not support sparse attention." |
There was a problem hiding this comment.
We have already checked sparse_params in is_available, is it necessary to check sparse attention again here?
| or getattr(meta, "is_spec_dec_dynamic_tree", False) | ||
| ): | ||
| return False, "CuTe DSL MLA FMHA does not support custom/tree speculative masks." | ||
| if q.shape[0] % meta.num_generations != 0: |
There was a problem hiding this comment.
We have checked gen-only batch before, this condition is true for gen-only batch, should we check it in runtime again?
| f"num_generations ({meta.num_generations}).", | ||
| ) | ||
| seq_len_q = q.shape[0] // meta.num_generations | ||
| if seq_len_q < 1: |
There was a problem hiding this comment.
is_supported is not for runtime assertion, it is for quick dispatch, we should avoid adding too much obvious conditions here. This condition must be false for correct q shape in gen-only case, can we remove it?
| num_heads: int, | ||
| batch_size: int, | ||
| seq_len_q: int, | ||
| predicted_tokens_per_seq: int, |
There was a problem hiding this comment.
predicted_tokens_per_seq is unused.
| return False, reason | ||
| if meta.kv_cache_block_offsets is None: | ||
| return False, "Paged KV block offsets are required." | ||
| if meta.kv_cache_manager is None: |
There was a problem hiding this comment.
This condition is enough to exclude no-cache case, kv_cache_block_offsets/pool_mapping checks are unnecessary.
| pool_mapping = meta.host_kv_cache_pool_mapping | ||
| if pool_mapping is None: | ||
| return False, "KV cache pool mapping is required." | ||
| if fwd.latent_cache is None: |
There was a problem hiding this comment.
latent_cache is not None for MLA case, which is checked, can we remove this condition?
| tokens_per_block = meta.tokens_per_block | ||
| if tokens_per_block is None: | ||
| tokens_per_block = getattr(meta.kv_cache_manager, "tokens_per_block", 0) | ||
| if tokens_per_block <= 1: |
There was a problem hiding this comment.
Don't need to check the correctness of tokens_per_block here
| f"has_fp8_kv_cache={getattr(attn, 'has_fp8_kv_cache', False)}.", | ||
| ) | ||
| if kernel_dtype == torch.float8_e4m3fn and ( | ||
| fwd.quant_q_buffer is None or fwd.mla_bmm1_scale is None or fwd.mla_bmm2_scale is None |
There was a problem hiding this comment.
MLA module has ensured it, don't need to check here.
|
PR_Github #61509 [ run ] completed with state
|
| `TLLM_FMHA_LIBS=-flashinfer_trtllm_gen` to force the fallback path. Each FMHA | ||
| library exposes `is_available()` for module/static environment checks and | ||
| `is_supported()` for per-forward request checks. | ||
| `cute_dsl_mla,flashinfer_trtllm_gen,fallback`; use `TLLM_FMHA_LIBS=fallback` |
There was a problem hiding this comment.
Seems that we forget to add msa_sparse_gqa in ATTENTION_DEVELOPER_GUIDE.md, can you also add it?
There was a problem hiding this comment.
Sure. Will add it.
| """ | ||
| if split_kv == 1: | ||
| return 0 | ||
| return B * H * S * split_kv * (D + 1) * acc_dtype.width // 8 |
There was a problem hiding this comment.
Pad the split-KV workspace to the 128-row MMA store
Both kernel variants accept every H <= 128, but size and lay out the split-KV partials with logical H. For the production H=16, S_q=1, split_kv>1 path, the physical MMA store writes 128 rows: padded head 16 starts immediately after the H-sized acc_o, so the store aliases acc_lse and can continue beyond the raw workspace. The host-side maximum-workspace formula also uses H, so this can corrupt reduction inputs or trigger an illegal memory access.
Suggested fix: use workspace_H = max(H, 128) in both host sizing paths and the equivalent cutlass.max in both initialize_workspace() implementations, then derive the acc_o/acc_lse shapes and strides from workspace_H. Until that is fixed, reject split-KV when the post-fold head count is below 128.
| for i in range(local_split_kv): | ||
| for j in cutlass.range_constexpr(elements_per_thread): | ||
| element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps | ||
| rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] |
There was a problem hiding this comment.
Bound split-KV before indexing the fixed-size reduction table
smem_lse_scale has exactly MAX_SPLITS entries, but this loop runs to runtime local_split_kv; neither kernel's can_implement() bounds split_kv. For example, seq_len_k=32896, split_kv=257 produces 257 K tiles with the default 128-column tile, and iteration 256 reads smem_lse_scale[256] out of bounds. The FP16 reduction has the same issue.
Suggested fix: reject split_kv < 1 and split_kv > MAX_SPLITS in both can_implement() methods. For variable split-KV, also validate every block_split_kvs[b] against min(split_kv, MAX_SPLITS), or consume scales in bounded chunks.
| meta: "TrtllmAttentionMetadata", | ||
| fwd: AttentionForwardArgs, | ||
| ) -> tuple[bool, str]: | ||
| if fwd.attention_input_type != AttentionInputType.generation_only: |
There was a problem hiding this comment.
Reject Helix until the CuTe path exports softmax statistics
This support gate does not reject Helix. MLA._attn_forward_gen() allocates an uninitialized softmax_stats tensor and later consumes it in _helix_post_process(), while the CuTe path ignores softmax_stats_tensor and keeps its LSE private in its workspace. Helix therefore combines partial outputs with uninitialized statistics and can silently return incorrect tokens.
Suggested fix: mirror the sibling FlashInfer guard:
if meta.helix_position_offsets is not None:
return False, "CuTe DSL MLA FMHA does not support Helix parallelism."Supporting Helix instead requires exporting the correct per-token statistics and honoring inactive-rank semantics.
5 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
cute_dsl_mlaFMHA backend for Blackwell SM100/SM103 (CuTe DSL required) to dispatch eligible DeepSeek MLA generation-phase requests to FP8 or FP16/BF16 CuTe DSL decode kernels.CuteDslMlaFmhawith strong request gating/rejection (decode-only, geometry/paging metadata requirements, unsupported masks/speculation/beam/custom cases), kernel feasibility checks, KV page-table prep, and CuTe layout normalization.CuteDSLNVMlaDecodeBlackwellRunner) with:mla_helpers.py) to support persistent/non-persistent scheduling and correct tile indexing.cute_dsl_mlabecomes available by default unless overridden byTLLM_FMHA_LIBS, and expanded the attention developer guide accordingly.sizesvectors consistent whensplit_chunkyields 0-token chunks on some ranks.Review notes / risk areas
QA Engineer Review
No test changes.