[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 22 commits into
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
| 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 |
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.
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.
There was a problem hiding this comment.
Discussed offline. We don't change the sync files (mla_decode.* + mla_helper.py) unless necessary fix.
So still need to add the three cutedsl kernel files into legacy-files.txt. Then call python3 scripts/legacy_utils.py check-configs to update .pre-commit-config.yaml, pyproject.toml and ruff-legacy.toml
| 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.
| 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.
There was a problem hiding this comment.
Discussed offline. The functionality is not affected even split_kv > MAX_SPLITS. But will add checking to make sure splitkv is in [1,32] for performance.
| # fallback (cache not warmed), supply the default 4-tuple so split_kv + | ||
| # is_persistent still come from a length-4 tactic. | ||
| if not (isinstance(best_tactic, tuple) and len(best_tactic) == 4): | ||
| best_tactic = runner.default_tactic(int(q_latent.shape[-1])) |
There was a problem hiding this comment.
I think we do not have default coverage for -1 during warmup as the default (needs confirmation). Then the fallback may cause another runtime cute.compile. This should not be a big deal because normally we do not expect fallback to happen.
There was a problem hiding this comment.
The unit test without autotuner will get -1 .
| mutates_args=("o", "workspace"), | ||
| device_types="cuda", | ||
| ) | ||
| def cute_dsl_mla_decode_fp8_blackwell( |
There was a problem hiding this comment.
So we can directly rely on the current attention unit tests to cover the new custom ops. Can the new tuning path be covered correctly?
There was a problem hiding this comment.
The current attention unit test could cover the new custom ops, but doesn't cover autotune.
There was a problem hiding this comment.
@yihwang-nv is adding DSA unit tests for attention backends in #16309, if we add unittest/_torch/attention/test_attention_backends.py to tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml, will it cover autotune?
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #61851 [ run ] triggered by Bot. Commit: |
|
PR_Github #61851 [ run ] completed with state
|
…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>
ab37c6d to
462c0d2
Compare
|
/bot run --disable-fail-fast --add-multi-gpu-test |
There was a problem hiding this comment.
♻️ Duplicate comments (1)
tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py (1)
386-425: 🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick winPage IDs still not scaled by
layers_in_poolin the packed view.After
kv_poolis re-viewed as packed combined slots (line 394-398), logical pagepfor layer slotlmaps top * layers_in_pool + lin the combined index space. Lines 424-425 only addlayer_in_pool, so for anyp > 0the kernel indexes the wrong page/layer wheneverlayers_in_pool > 1(interleaved KV cache manager layout). This was flagged on an earlier revision and remains unresolved.🐛 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🤖 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 386 - 425, Update the page-table transformation after the packed KV cache view in the surrounding attention setup: when layers_in_pool > 1, scale every logical page ID in page_table by layers_in_pool, then add layer_in_pool to select the current layer slot. Preserve the existing unmodified page table behavior when layers_in_pool equals 1.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py (2)
157-160: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated tiler constants risk drift between eligibility check and runner.
mma_qk_tiler_mn/mma_pv_tiler_mnare hardcoded here to mirrorCuteDSLNVMlaDecodeBlackwellRunner's actual launch config incute_dsl_custom_ops.py. The comment flags the sync requirement, but nothing enforces it — if the runner's defaults change, this eligibility check silently goes stale andcan_implementvalidates against the wrong tiler.Consider exposing the default tiler as a shared class attribute/classmethod on
CuteDSLNVMlaDecodeBlackwellRunnerand importing it here instead of re-declaring the literals.🤖 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 157 - 160, Replace the hardcoded default tilers in the eligibility check with a shared class attribute or classmethod exposed by CuteDSLNVMlaDecodeBlackwellRunner, and import/reuse that value in cute_dsl.py. Update the runner’s default configuration to be the single source of truth while preserving the existing (128, 128) and (128, 256) behavior.
45-91: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPublic interface methods lack docstrings.
is_available,is_supported,run_mla_generation, andprepare_workspaceare the public surface of this FMHA backend but have no docstrings documenting arguments/behavior (only inline comments on individual checks). Given the non-obvious eligibility rules and workspace-sizing contract, Google-style docstrings would materially help future maintainers.As per coding guidelines: "Prefer docstrings for external interfaces, use Google-style docstrings, document public function arguments, and include tensor dimensions and constrained dtypes where applicable."
Also applies to: 189-205, 485-505, 529-551
🤖 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 45 - 91, The public methods is_available, is_supported, run_mla_generation, and prepare_workspace lack interface documentation. Add concise Google-style docstrings to each method describing arguments, return or operational behavior, eligibility constraints, workspace-sizing requirements, and relevant tensor dimensions and constrained dtypes, while leaving the implementation unchanged.Source: Coding guidelines
🤖 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.
Duplicate comments:
In `@tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py`:
- Around line 386-425: Update the page-table transformation after the packed KV
cache view in the surrounding attention setup: when layers_in_pool > 1, scale
every logical page ID in page_table by layers_in_pool, then add layer_in_pool to
select the current layer slot. Preserve the existing unmodified page table
behavior when layers_in_pool equals 1.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.py`:
- Around line 157-160: Replace the hardcoded default tilers in the eligibility
check with a shared class attribute or classmethod exposed by
CuteDSLNVMlaDecodeBlackwellRunner, and import/reuse that value in cute_dsl.py.
Update the runner’s default configuration to be the single source of truth while
preserving the existing (128, 128) and (128, 256) behavior.
- Around line 45-91: The public methods is_available, is_supported,
run_mla_generation, and prepare_workspace lack interface documentation. Add
concise Google-style docstrings to each method describing arguments, return or
operational behavior, eligibility constraints, workspace-sizing requirements,
and relevant tensor dimensions and constrained dtypes, while leaving the
implementation unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0644b885-7858-4af0-992c-b59fecf1dfef
📒 Files selected for processing (3)
tensorrt_llm/_torch/attention_backend/fmha/cute_dsl.pytensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.pytensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md
🚧 Files skipped from review as they are similar to previous changes (2)
- tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.md
- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
|
PR_Github #62126 [ run ] triggered by Bot. Commit: |
… split_kv - Perf gate: _PERF_MIN_BATCH_FP8[(128, 1)] 8 -> 64. The measured crossover for the unsharded 128-head fp8 decode shape is at batch 64; the old floor let small per-rank batches through, where the default TRTLLM kernels are faster. - Drop the duplicate sparse_params check in _is_supported_with_reason. is_available already rejects sparse layers, so the second check was dead code. - can_implement (both fp8 and fp16 kernels): reject split_kv outside [1, 32], the range the kernel is tuned for. - test_attention_mla: add a 64-request context list and raise max_num_contexts to 64 so the fp8 decode gate is actually exercised. Verified on B200: 96/96 pass and 12 of the parametrized cells now reach the CuTe DSL decode kernel (fp8, num_heads=128, batch 50/64, seq_len_q 1 and 4, v1 and v2 KV cache). - Legacy lint: drop the two CuTe DSL attention __init__.py from legacy-files.txt and regenerate the managed blocks in ruff-legacy.toml / pyproject.toml / .pre-commit-config.yaml. Signed-off-by: haow <haow@nvidia.com>
6 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.