Skip to content

[None][feat] Enable CuTe DSL MLA with Helix #18131

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
mingyangHao merged 3 commits into from
Aug 28, 2026
Merged

[None][feat] Enable CuTe DSL MLA with Helix #18131

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e95fcf2
[None][feat] Enable CuTe DSL MLA with Helix
mingyangHao Aug 24, 2026
2b4091d
[None][perf] Fuse CuTe DSL MLA softmax stats output
mingyangHao Aug 26, 2026
43e2ee1
[None][fix] Preserve optional MLA mutation argument
mingyangHao Aug 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 18 additions & 2 deletions tensorrt_llm/_torch/attention_backend/fmha/cute_dsl_mla.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -287,8 +287,6 @@ def _is_supported_with_reason(
# configuration after K3's 96-to-128 head padding was removed.
if (meta.num_contexts != 0 and attn.num_heads != 96) or meta.num_generations <= 0:
return False, "CuTe DSL MLA FMHA only supports decode-only batches."
if meta.helix_position_offsets is not None:
return False, "CuTe DSL MLA FMHA does not support Helix parallelism."
if meta.beam_width != 1:
return False, f"Beam search is not supported, got beam_width={meta.beam_width}."
# Linear-chain MTP / spec-decode (seq_len_q > 1) IS supported: the
Expand All @@ -305,6 +303,23 @@ def _is_supported_with_reason(
return False, "CuTe DSL MLA FMHA does not support custom/tree speculative masks."
seq_len_q = q.shape[0] // meta.num_generations
batch_size = meta.num_generations
if meta.helix_position_offsets is not None:
if seq_len_q != 1:
return False, "CuTe DSL MLA FMHA only supports single-token decode with Helix."
Comment thread
mingyangHao marked this conversation as resolved.
softmax_stats = fwd.softmax_stats_tensor
if softmax_stats is None:
return False, "CuTe DSL MLA FMHA requires softmax_stats_tensor with Helix."
expected_stats_shape = (q.shape[0], attn.num_heads, 2)
if softmax_stats.shape != expected_stats_shape:
return False, (
"CuTe DSL MLA FMHA requires Helix softmax stats with shape "
f"{expected_stats_shape}, got {tuple(softmax_stats.shape)}."
)
if softmax_stats.dtype != torch.float32 or not softmax_stats.is_contiguous():
return False, (
"CuTe DSL MLA FMHA requires contiguous float32 Helix softmax stats, "
f"got dtype={softmax_stats.dtype}, contiguous={softmax_stats.is_contiguous()}."
)

from tensorrt_llm._torch.autotuner import AutoTuner

Expand Down Expand Up @@ -492,6 +507,7 @@ def _run_mla_decode(
output_scale,
# Max batch size for the AutoTuner to profile.
int(meta.max_num_requests),
params.fwd.softmax_stats_tensor,
)

def run_mla_generation(
Expand Down
124 changes: 92 additions & 32 deletions tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9081,6 +9081,7 @@ def __init__(
seq_len_q: int,
page_size: int,
max_batch_size: int = 0,
emit_softmax_stats: bool = False,
):
super().__init__()
kernel_class = self.__class__._KERNEL_CLASS_BY_DTYPE.get(in_dtype)
Expand All @@ -9095,18 +9096,20 @@ def __init__(
self.seq_len_q = seq_len_q
self.page_size = page_size
self.max_batch_size = max_batch_size
self.emit_softmax_stats = emit_softmax_stats

def unique_id(self):
# seq_len_q is part of the id: each decode variant (the MTP
# target step's sq = 1 + draft_len, the draft steps' sq = 1)
# constructs its own runner and is tuned independently during
# the autotuner warmup's generation forward.
return (
base_id = (
self.in_dtype,
self.num_heads,
self.seq_len_q,
self.page_size,
)
return base_id + (True, ) if self.emit_softmax_stats else base_id

@classmethod
def _get_max_active_blocks(cls) -> int:
Expand Down Expand Up @@ -9367,7 +9370,8 @@ def get_tuning_config(self) -> TuningConfig:
# 4 page_table: (max_blocks_per_sequence, B)
# 5 cache_seqs: (B,)
# 6 o: (H, D, S_q, B)
# 7 workspace: (workspace_size,)
# 7 workspace: (workspace_size,)
# 8 softmax_stats: (B * S_q, H, 2), optional

# cache_seqs (index 5) is the single free dynamic batch dim;
# every other batch-carrying dim is tied to it by a constraint.
Expand All @@ -9394,6 +9398,14 @@ def get_tuning_config(self) -> TuningConfig:
i, d, lambda shapes, _i=i, _d=d: shapes[_i][_d])
for (i, d) in static_size_dims)

stats_constraints = ()
if self.emit_softmax_stats:
stats_constraints = (ConstraintSpec(
8,
0,
lambda shapes: shapes[5][0] * self.seq_len_q,
), )

# The batch search space, fixed up-front by max_batch_size
# when the engine max is known.
batch_buckets = (get_last_power_of_2_num_tokens_buckets(
Expand All @@ -9406,7 +9418,8 @@ def get_tuning_config(self) -> TuningConfig:
batch_buckets,
last_positive_power_of_2,
), ),
constraint_specs=batch_constraints + static_constraints,
constraint_specs=(batch_constraints + static_constraints +
stats_constraints),
inputs_pre_hook=self._tuning_inputs_pre_hook,
)
return cache[key]
Expand Down Expand Up @@ -9450,6 +9463,9 @@ def forward(
inputs[6] (o): Output tensor of shape (H, D, S_q, B).
inputs[7] (workspace): Contiguous raw workspace with at least
the workspace_size returned by get_workspace_layout.
inputs[8] (softmax_stats): Optional contiguous float32 tensor
of shape (B * S_q, H, 2). The kernel writes an equivalent
softmax (max, sum) pair for Helix reduction.
tactic: Tuple containing (mma_qk_tiler_mn, mma_pv_tiler_mn,
split_kv, is_persistent).
**kwargs: Optional softmax_scale and output_scale values.
Expand All @@ -9459,7 +9475,7 @@ def forward(
tensor of shape (H, S_q, B) remains in the workspace.
"""
(q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, o,
workspace) = inputs
workspace, softmax_stats) = inputs
softmax_scale = float(kwargs.get("softmax_scale", 1.0))
output_scale = float(kwargs.get("output_scale", 1.0))

Expand Down Expand Up @@ -9489,6 +9505,22 @@ def forward(
# workspace = lse + split_kv_workspace
batch_size = cache_seqs.shape[0]
d_latent = q_latent.shape[1]
softmax_stats_kernel = None
if softmax_stats is not None:
expected_shape = (batch_size * seq_len_q, self.num_heads, 2)
if (softmax_stats.shape != expected_shape
or softmax_stats.dtype != torch.float32
or softmax_stats.device != o.device
or not softmax_stats.is_contiguous()):
raise RuntimeError(
"CuteDSLNVMlaDecodeBlackwellRunner requires contiguous "
"float32 softmax_stats on the output device with shape "
f"{expected_shape}, got shape={tuple(softmax_stats.shape)}, "
f"dtype={softmax_stats.dtype}, device={softmax_stats.device}, "
f"contiguous={softmax_stats.is_contiguous()}.")
softmax_stats_kernel = softmax_stats.view(
batch_size, seq_len_q, self.num_heads,
2).permute(2, 1, 0, 3)
max_batch_size = max(batch_size, self.max_batch_size)
(lse_offset, lse_size, split_kv_offset, split_kv_size,
required_workspace_size) = self.get_workspace_layout(
Expand Down Expand Up @@ -9556,8 +9588,8 @@ def forward(
num_heads=self.num_heads,
seq_len_q=seq_len_q,
fold_sq=fold_sq,
emit_softmax_stats=self.emit_softmax_stats,
)

q_latent_ct = cute.runtime.from_dlpack(
q_latent,
assumed_align=16).mark_layout_dynamic(leading_dim=1)
Expand All @@ -9579,6 +9611,10 @@ def forward(
divisibility=(128 // out_dtype.width))
lse_ct = cute.runtime.from_dlpack(
lse, assumed_align=16).mark_layout_dynamic(leading_dim=0)
softmax_stats_ct = (cute.runtime.from_dlpack(
softmax_stats_kernel, assumed_align=16).mark_layout_dynamic(
leading_dim=3) if softmax_stats_kernel is not None else
None)
use_workspace = split_kv > 1 and split_workspace.numel() > 0
workspace_ct = (cute.runtime.from_dlpack(
split_workspace, assumed_align=32).mark_layout_dynamic()
Expand All @@ -9588,36 +9624,49 @@ def forward(
# Variable split-KV (block_split_kvs) is not used on this path:
block_split_kvs_ct = None

CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache[cache_key] = \
compile_args = [
q_latent_ct,
q_rope_ct,
c_latent_ct,
c_rope_ct,
page_table_ct,
o_ct,
lse_ct,
]
compile_target = mla
if self.emit_softmax_stats:
compile_target = mla.run_with_softmax_stats
compile_args.append(softmax_stats_ct)
compile_args.extend([
workspace_ct,
split_kv,
cache_seqs_ct,
block_split_kvs_ct,
cutlass.Float32(softmax_scale),
cutlass.Float32(output_scale),
stream,
])
CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache[cache_key] = (
cute.compile(
mla,
q_latent_ct,
q_rope_ct,
c_latent_ct,
c_rope_ct,
page_table_ct,
o_ct,
lse_ct,
workspace_ct,
split_kv,
cache_seqs_ct,
block_split_kvs_ct,
cutlass.Float32(softmax_scale),
cutlass.Float32(output_scale),
stream,
compile_target,
*compile_args,
options="--opt-level 2",
)
))

compiled_mla = CuteDSLNVMlaDecodeBlackwellRunner.kernel_cache[
cache_key]
compiled_mla(
runtime_args = [
q_latent,
q_rope,
c_latent,
c_rope,
page_table,
o,
lse,
]
if self.emit_softmax_stats:
runtime_args.append(softmax_stats_kernel)
runtime_args.extend([
split_workspace if
(split_kv > 1 and split_workspace.numel() > 0) else None,
split_kv,
Expand All @@ -9626,12 +9675,13 @@ def forward(
softmax_scale,
output_scale,
stream,
)
])
compiled_mla(*runtime_args)
return o

@torch.library.custom_op(
"trtllm::cute_dsl_mla_decode_fp8_blackwell",
mutates_args=("o", "workspace"),
mutates_args=("o", "workspace", "softmax_stats"),
device_types="cuda",
)
def cute_dsl_mla_decode_fp8_blackwell(
Expand All @@ -9648,7 +9698,11 @@ def cute_dsl_mla_decode_fp8_blackwell(
page_size: int,
softmax_scale: float,
output_scale: float,
max_batch_size: int = 0,
# Keep the last two arguments required in the custom-op schema. PyTorch
# elides trailing default-valued arguments before its mutation fallback,
# while mutates_args retains their positional indices.
max_batch_size: int,
softmax_stats: Optional[torch.Tensor],
) -> None:
"""CuTe DSL FP8 MLA decode (Blackwell SM100/SM103).
"""
Expand All @@ -9665,10 +9719,11 @@ def cute_dsl_mla_decode_fp8_blackwell(
seq_len_q=seq_len_q,
page_size=page_size,
max_batch_size=max_batch_size,
emit_softmax_stats=softmax_stats is not None,
)
inputs = [
q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, o,
workspace
workspace, softmax_stats
]
tuner = AutoTuner.get()
_, best_tactic = tuner.choose_one(
Expand Down Expand Up @@ -9704,13 +9759,14 @@ def _(
page_size: int,
softmax_scale: float,
output_scale: float,
max_batch_size: int = 0,
max_batch_size: int,
softmax_stats: Optional[torch.Tensor],
) -> None:
return None

@torch.library.custom_op(
"trtllm::cute_dsl_mla_decode_fp16_blackwell",
mutates_args=("o", "workspace"),
mutates_args=("o", "workspace", "softmax_stats"),
device_types="cuda",
)
def cute_dsl_mla_decode_fp16_blackwell(
Expand All @@ -9727,7 +9783,9 @@ def cute_dsl_mla_decode_fp16_blackwell(
page_size: int,
softmax_scale: float,
output_scale: float,
max_batch_size: int = 0,
# See the FP8 op above: these must remain required schema arguments.
max_batch_size: int,
softmax_stats: Optional[torch.Tensor],
) -> None:
"""CuTe DSL FP16/BF16 MLA decode (Blackwell SM100/SM103).
"""
Expand Down Expand Up @@ -9761,10 +9819,11 @@ def cute_dsl_mla_decode_fp16_blackwell(
seq_len_q=seq_len_q,
page_size=page_size,
max_batch_size=max_batch_size,
emit_softmax_stats=softmax_stats is not None,
)
inputs = [
q_latent, q_rope, c_latent, c_rope, page_table, cache_seqs, o,
workspace
workspace, softmax_stats
]
tuner = AutoTuner.get()
_, best_tactic = tuner.choose_one(
Expand Down Expand Up @@ -9800,6 +9859,7 @@ def _(
page_size: int,
softmax_scale: float,
output_scale: float,
max_batch_size: int = 0,
max_batch_size: int,
softmax_stats: Optional[torch.Tensor],
) -> None:
return None
Loading
Loading