diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 3fde9bd1bf76..91364b67d840 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -677,6 +677,7 @@ def __init__( layer_idx: Optional[int] = None, is_sparse_attention_layer: bool = False, disable_index_value: bool = False, + aux_stream: Optional[torch.cuda.Stream] = None, ): config = model_config.pretrained_config self.pretrained_config = config @@ -708,6 +709,10 @@ def __init__( getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) ) + # Dtype of the q/k/v activations fed into norm and RoPE. KV-cache + # quantization changes cache storage only, so this stays the compute dtype. + self.attn_activation_dtype = config.torch_dtype + # Per-head Gemma RMSNorm — one set of weights shared across heads. self.q_norm = RMSNorm( hidden_size=self.head_dim_value, @@ -722,6 +727,12 @@ def __init__( use_gemma=self.use_gemma_norm, ) + # Stream and events used to overlap independent projection/norm/RoPE + # work. The stream is shared across layers via aux_stream_dict, matching + # DeepSeekV3; a private stream is used when no stream is supplied. + self.aux_stream = aux_stream if aux_stream is not None else torch.cuda.Stream() + self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] + self.is_sparse_attention_layer = bool(is_sparse_attention_layer) self.disable_index_value = bool(disable_index_value) if self.is_sparse_attention_layer: @@ -781,8 +792,23 @@ def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, """ q_shape = q.shape k_shape = k.shape - q = self.q_norm(q.reshape(-1, self.head_dim_value)).reshape(q_shape) - k = self.k_norm(k.reshape(-1, self.head_dim_value)).reshape(k_shape) + + def _q_norm(): + return self.q_norm(q.reshape(-1, self.head_dim_value)).reshape(q_shape) + + def _k_norm(): + return self.k_norm(k.reshape(-1, self.head_dim_value)).reshape(k_shape) + + # The q-norm and k-norm are independent, so overlap them on the aux + # stream. Multi-stream is disabled under torch.compile. + q, k = maybe_execute_in_parallel( + _q_norm, + _k_norm, + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + disable_on_compile=True, + ) return q, k def apply_index_qk_norm( @@ -809,26 +835,94 @@ def apply_index_qk_norm( ) idx_q_shape = idx_q.shape idx_k_shape = idx_k.shape - idx_q = self.index_q_norm(idx_q.reshape(-1, self.sparse_index_dim)).reshape(idx_q_shape) - idx_k = self.index_k_norm(idx_k.reshape(-1, self.sparse_index_dim)).reshape(idx_k_shape) + + def _idx_q_norm(): + return self.index_q_norm(idx_q.reshape(-1, self.sparse_index_dim)).reshape(idx_q_shape) + + def _idx_k_norm(): + return self.index_k_norm(idx_k.reshape(-1, self.sparse_index_dim)).reshape(idx_k_shape) + + idx_q, idx_k = maybe_execute_in_parallel( + _idx_q_norm, + _idx_k_norm, + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + disable_on_compile=True, + ) return idx_q, idx_k - def apply_rope( + def _fused_qk_norm_rope( self, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - position_ids: torch.Tensor, - ): - """Run per-head QK norm before partial RoPE. + qkv: torch.Tensor, + position_ids: Optional[torch.Tensor], + *, + num_heads_q: int, + num_heads_k: int, + num_heads_v: int, + head_dim: int, + q_norm: RMSNorm, + k_norm: RMSNorm, + ) -> Optional[torch.Tensor]: + """Fuse per-head Gemma RMSNorm and partial RoPE into one kernel. + + Runs torch.ops.trtllm.fused_qk_norm_rope over the packed + [Q heads, K heads, optional V heads] layout and returns the mutated + tensor, which aliases qkv when qkv is already contiguous. The kernel + norms the full head_dim and rotates only the leading rotary_dim + channels, matching M3's whole-head norm with front partial RoPE, and + leaves the V heads untouched. + + Returns None with qkv unmodified when the kernel does not apply, so the + caller runs norm and RoPE separately: non-bf16 activations (the kernel + is bf16-only), missing position_ids, or no rotary embedding. + """ + if position_ids is None or qkv.dtype != torch.bfloat16: + return None + if ( + self.rotary_emb is None + or self.pos_embd_params is None + or self.pos_embd_params.rope is None + ): + return None + + # Partial-RoPE dim comes from RopeParams (M3 rotates 64 of 128). + rotary_dim = int(self.pos_embd_params.rope.dim) + # The kernel assumes a contiguous [num_tokens, total_heads * head_dim]. + qkv = qkv.contiguous() + torch.ops.trtllm.fused_qk_norm_rope( + qkv, + num_heads_q, + num_heads_k, + num_heads_v, + head_dim, + rotary_dim, + q_norm.variance_epsilon, + q_norm.weight, + k_norm.weight, + self.pos_embd_params.rope.theta, + self.pos_embd_params.is_neox, + position_ids.reshape(-1).contiguous().to(torch.int32), + 1.0, # factor: no YARN (M3 has no rope_scaling) + 0.0, # low + 0.0, # high + 1.0, # attention_factor + True, # is_qk_norm + self.use_gemma_norm, # use_gemma + False, # use_mrope + 0, # mrope_section1 + 0, # mrope_section2 + ) + return qkv + + def _expect_fused_qk_norm_rope(self, position_ids: Optional[torch.Tensor]) -> bool: + """Whether the fused kernel is expected to run instead of the fallback. - The base ``Attention.apply_rope`` consumes split q/k/v. We split, - apply per-head QK norm, then defer to the base partial-RoPE - implementation (driven by ``RopeParams.dim < head_dim``). + Every M3 config keeps bf16 attention activations, so with position_ids + present a fallback means the fused kernel silently stopped applying. + The forward paths assert on this. """ - q, k, v = self.split_qkv(q, k, v) - q, k = self.apply_qk_norm(q, k) - return super().apply_rope(q, k, v, position_ids) + return self.attn_activation_dtype == torch.bfloat16 and position_ids is not None def forward( self, @@ -877,9 +971,8 @@ def _dense_forward( Steps: 1. Project Q/K/V via fused ``qkv_proj``. - 2. Apply per-head Gemma RMSNorm to Q/K (same as - :meth:`_sparse_forward` step 2 minus the index branch). - 3. Apply partial RoPE. + 2-3. Apply per-head Gemma RMSNorm and partial RoPE to Q/K, fused + into one kernel by :meth:`_fused_qk_norm_rope` when it applies. 4. Pull the paged main K/V cache from the M3 cache manager. 5. Read the pre-built :class:`MiniMaxM3SparseAttentionMetadata` from ``attn_metadata.minimax_m3``. Production code paths @@ -904,16 +997,35 @@ def _dense_forward( "attn_metadata; received None." ) - # 1. Projections (no index branch). + # Projections (no index branch). qkv = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - # 2. Per-head Gemma RMSNorm on Q/K (no index norm). - q, k = self.apply_qk_norm(q, k) - - # 3. Partial RoPE on Q/K (no index branch). - if self.rotary_emb is not None and position_ids is not None: - q, k = self.rotary_emb(position_ids, [q, k]) + # Per-head Gemma RMSNorm and partial RoPE on Q/K. + fused_qkv = self._fused_qk_norm_rope( + qkv, + position_ids, + num_heads_q=self.num_heads, + num_heads_k=self.num_key_value_heads, + num_heads_v=self.num_key_value_heads, + head_dim=self.head_dim, + q_norm=self.q_norm, + k_norm=self.k_norm, + ) + if fused_qkv is not None: + q, k, v = fused_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + # Match the contiguity of the separate path; V stays a column-slice view. + q, k = q.contiguous(), k.contiguous() + else: + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 dense attention (layer {self.layer_idx}) expected the " + f"fused QK-norm+RoPE kernel (bf16 activations, head_dim=" + f"{self.head_dim}) but fell back to the separate path; qkv dtype " + f"is {qkv.dtype} (expected {self.attn_activation_dtype})." + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.apply_qk_norm(q, k) + if self.rotary_emb is not None and position_ids is not None: + q, k = self.rotary_emb(position_ids, [q, k]) # Keep token-wise projections and the output projection visible to # torch.compile. Only the metadata/cache-dependent attention core is @@ -1189,9 +1301,10 @@ def _sparse_forward( Steps: 1. Project ``hidden_states`` to Q/K/V (fused ``qkv_proj``) plus index Q (per-head) and index K (single replicated). - 2. Apply per-head Gemma RMSNorm to both branches. - 3. Apply partial RoPE (``rotary_dim`` channels of ``head_dim``) - to both branches. + 2-3. Apply per-head Gemma RMSNorm and partial RoPE to the main and + index branches, each fused into one kernel by + :meth:`_fused_qk_norm_rope` when it applies. The index branch + passes num_heads_v=0 because it carries no value heads. 4. Pull paged main K/V cache (reshaped to flat-slot view) and paged side index-K cache from the :class:`MiniMaxM3KVCacheManagerV2`. @@ -1222,22 +1335,72 @@ def _sparse_forward( f"MiniMax-M3 sparse forward (layer {self.layer_idx}) requires " "attn_metadata; received None." ) - # 1. Projections. - qkv = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - idx_qk = self.index_qk_proj(hidden_states) - idx_q, idx_k = idx_qk.split([self.index_q_size, self.index_k_size], dim=-1) - # 2. Per-head Gemma RMSNorm on both branches. - q, k = self.apply_qk_norm(q, k) - idx_q, idx_k = self.apply_index_qk_norm(idx_q, idx_k) - - # 3. Partial RoPE on both branches. The base ``Attention`` - # constructor created ``self.rotary_emb`` for the configured - # partial ``rotary_dim`` because ``rope_fusion=False``. - if self.rotary_emb is not None and position_ids is not None: - q, k = self.rotary_emb(position_ids, [q, k]) - idx_q, idx_k = self.rotary_emb(position_ids, [idx_q, idx_k]) + # Project, norm, and apply RoPE for the main and index branches. Both + # read only hidden_states and write disjoint outputs, so they overlap on + # the aux stream and join before the attention core. + def _main_norm_rope(): + qkv = self.qkv_proj(hidden_states) + fused_qkv = self._fused_qk_norm_rope( + qkv, + position_ids, + num_heads_q=self.num_heads, + num_heads_k=self.num_key_value_heads, + num_heads_v=self.num_key_value_heads, + head_dim=self.head_dim, + q_norm=self.q_norm, + k_norm=self.k_norm, + ) + if fused_qkv is not None: + q, k, v = fused_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + return q.contiguous(), k.contiguous(), v + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 sparse attention (layer {self.layer_idx}) expected the " + f"fused QK-norm+RoPE kernel (bf16 activations, head_dim=" + f"{self.head_dim}) but fell back to the separate path; qkv dtype " + f"is {qkv.dtype} (expected {self.attn_activation_dtype})." + ) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.apply_qk_norm(q, k) + if self.rotary_emb is not None and position_ids is not None: + q, k = self.rotary_emb(position_ids, [q, k]) + return q, k, v + + def _index_norm_rope(): + idx_qk = self.index_qk_proj(hidden_states) + fused_idx = self._fused_qk_norm_rope( + idx_qk, + position_ids, + num_heads_q=self.sparse_num_index_heads, + num_heads_k=1, + num_heads_v=0, + head_dim=self.sparse_index_dim, + q_norm=self.index_q_norm, + k_norm=self.index_k_norm, + ) + if fused_idx is not None: + idx_q, idx_k = fused_idx.split([self.index_q_size, self.index_k_size], dim=-1) + return idx_q.contiguous(), idx_k.contiguous() + assert not self._expect_fused_qk_norm_rope(position_ids), ( + f"MiniMax-M3 sparse index branch (layer {self.layer_idx}) expected the " + f"fused QK-norm+RoPE kernel (bf16 activations, index_dim=" + f"{self.sparse_index_dim}) but fell back to the separate path; idx " + f"dtype is {idx_qk.dtype} (expected {self.attn_activation_dtype})." + ) + idx_q, idx_k = idx_qk.split([self.index_q_size, self.index_k_size], dim=-1) + idx_q, idx_k = self.apply_index_qk_norm(idx_q, idx_k) + if self.rotary_emb is not None and position_ids is not None: + idx_q, idx_k = self.rotary_emb(position_ids, [idx_q, idx_k]) + return idx_q, idx_k + + (q, k, v), (idx_q, idx_k) = maybe_execute_in_parallel( + _main_norm_rope, + _index_norm_rope, + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + disable_on_compile=True, + ) o = self._forward_attention_core(q, k, v, idx_q, idx_k, attn_metadata) return self.o_proj(o) @@ -1374,6 +1537,7 @@ def __init__( layer_idx=layer_idx, is_sparse_attention_layer=is_sparse, disable_index_value=disable_index_value, + aux_stream=aux_stream_dict[AuxStreamType.Attention], ) _, moe_layer_ids = get_moe_layer_ids(config) @@ -1448,10 +1612,11 @@ def __init__(self, model_config: "ModelConfig[PretrainedConfig]"): model_config.pretrained_config.torch_dtype = torch.bfloat16 config = model_config.pretrained_config self.vocab_size = config.vocab_size - # Two aux streams: one for MoE shared/routed parallel execution, - # one for MoE chunking overlap inside the fused MoE kernel. - # Matches the DeepSeekV3 convention. + # Aux streams shared across layers, matching the DeepSeekV3 convention: + # one for attention branch overlap, one for MoE shared/routed parallel + # execution, and one for MoE chunking overlap inside the fused MoE kernel. self.aux_stream_dict = { + AuxStreamType.Attention: torch.cuda.Stream(), AuxStreamType.MoeShared: torch.cuda.Stream(), AuxStreamType.MoeChunkingOverlap: torch.cuda.Stream(), } diff --git a/tests/unittest/_torch/models/test_minimax_m3.py b/tests/unittest/_torch/models/test_minimax_m3.py index ecf64099d77d..4e45181aaf52 100644 --- a/tests/unittest/_torch/models/test_minimax_m3.py +++ b/tests/unittest/_torch/models/test_minimax_m3.py @@ -600,6 +600,242 @@ def test_minimax_m3_attention_dense_apply_index_qk_norm_raises(): attn.apply_index_qk_norm(idx_q, idx_k) +def _make_fused_qk_norm_rope_test_config(): + """Return (text_config, ModelConfig) with a kernel-supported head_dim. + + fused_qk_norm_rope only compiles for head_dim in {64, 128, 256}, so the + head_dim=32 config from _make_attention_test_config cannot drive the fused + path. This variant keeps the real M3 head_dim=128, sparse_index_dim=128 and + rotary_dim=64 geometry with few heads so the tensors stay small. + """ + n_layers = 4 + sparse_cfg = { + "use_sparse_attention": True, + "sparse_index_dim": 128, + "sparse_num_index_heads": 2, + "sparse_topk_blocks": 4, + "sparse_block_size": 16, + "sparse_init_block": 0, + "sparse_local_block": 1, + "sparse_score_type": "max", + "sparse_disable_index_value": [0, 1, 1, 1], + "sparse_attention_freq": [0, 1, 1, 1], + } + text_cfg = _wrap_dict_as_config( + { + "hidden_size": 512, + "intermediate_size": 128, + "num_hidden_layers": n_layers, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 128, + "vocab_size": 256, + "max_position_embeddings": 256, + "rms_norm_eps": 1e-6, + "use_gemma_norm": True, + "rope_theta": 5000000.0, + "rotary_dim": 64, + "partial_rotary_factor": 0.5, + "qk_norm_type": "per_head", + "use_qk_norm": True, + "sparse_attention_config": sparse_cfg, + "torch_dtype": torch.bfloat16, + } + ) + model_cfg = ModelConfig( + pretrained_config=text_cfg, + mapping=Mapping(), + skip_create_weights_in_init=True, + ) + return text_cfg, model_cfg + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 fused QK-norm+RoPE needs CUDA") +def test_minimax_m3_fused_qk_norm_rope_main_matches_separate(): + """Fused main-branch helper matches separate norm plus partial RoPE. + + Covers the helper's wiring against the path it replaces: partial rotary dim, + theta and neox flag from RopeParams, Gemma scaling, per-head norm weights, + and the norm epsilon. + """ + _, model_cfg = _make_fused_qk_norm_rope_test_config() + attn = MiniMaxM3Attention( + model_config=model_cfg, + layer_idx=0, + is_sparse_attention_layer=False, + ) + device = torch.device("cuda") + dtype = torch.bfloat16 + head_dim = attn.head_dim + + torch.manual_seed(0) + attn.q_norm.weight = torch.nn.Parameter(torch.randn(head_dim, dtype=dtype, device=device) * 0.2) + attn.k_norm.weight = torch.nn.Parameter(torch.randn(head_dim, dtype=dtype, device=device) * 0.2) + + seq = 6 + qkv = torch.randn(seq, attn.q_size + 2 * attn.kv_size, dtype=dtype, device=device) + position_ids = torch.arange(seq, dtype=torch.int32, device=device) + 3 + + # Fused path. + fused = attn._fused_qk_norm_rope( + qkv.clone(), + position_ids, + num_heads_q=attn.num_heads, + num_heads_k=attn.num_key_value_heads, + num_heads_v=attn.num_key_value_heads, + head_dim=head_dim, + q_norm=attn.q_norm, + k_norm=attn.k_norm, + ) + assert fused is not None, "bf16 qkv + position_ids must take the fused path" + q_f, k_f, v_f = fused.split([attn.q_size, attn.kv_size, attn.kv_size], dim=-1) + + # Separate fallback path. + q_s, k_s, v_s = qkv.split([attn.q_size, attn.kv_size, attn.kv_size], dim=-1) + q_s, k_s = attn.apply_qk_norm(q_s, k_s) + q_s, k_s = attn.rotary_emb(position_ids, [q_s, k_s]) + + torch.testing.assert_close(q_f.contiguous(), q_s.contiguous(), rtol=5e-2, atol=1e-1) + torch.testing.assert_close(k_f.contiguous(), k_s.contiguous(), rtol=5e-2, atol=1e-1) + # V is untouched by both paths. + torch.testing.assert_close(v_f.contiguous(), v_s.contiguous(), rtol=0, atol=0) + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 fused QK-norm+RoPE needs CUDA") +def test_minimax_m3_fused_qk_norm_rope_index_matches_separate(): + """Fused index-branch helper matches the separate path. + + The index branch norms and rotates the concatenated idx_q (per-head) and + idx_k (single replicated head) with num_heads_v=0, then splits back. + """ + _, model_cfg = _make_fused_qk_norm_rope_test_config() + attn = MiniMaxM3Attention( + model_config=model_cfg, + layer_idx=3, + is_sparse_attention_layer=True, + disable_index_value=True, + ) + device = torch.device("cuda") + dtype = torch.bfloat16 + sparse_index_dim = attn.sparse_index_dim + num_index_heads = attn.sparse_num_index_heads + + torch.manual_seed(1) + attn.index_q_norm.weight = torch.nn.Parameter( + torch.randn(sparse_index_dim, dtype=dtype, device=device) * 0.3 + ) + attn.index_k_norm.weight = torch.nn.Parameter( + torch.randn(sparse_index_dim, dtype=dtype, device=device) * 0.3 + ) + + seq = 5 + idx_q = torch.randn(seq, num_index_heads * sparse_index_dim, dtype=dtype, device=device) + idx_k = torch.randn(seq, sparse_index_dim, dtype=dtype, device=device) + position_ids = torch.arange(seq, dtype=torch.int32, device=device) + 7 + + # Fused path over concatenated [idx_q, idx_k]. + fused = attn._fused_qk_norm_rope( + torch.cat([idx_q, idx_k], dim=-1), + position_ids, + num_heads_q=num_index_heads, + num_heads_k=1, + num_heads_v=0, + head_dim=sparse_index_dim, + q_norm=attn.index_q_norm, + k_norm=attn.index_k_norm, + ) + assert fused is not None + iq_f, ik_f = fused.split([num_index_heads * sparse_index_dim, sparse_index_dim], dim=-1) + + # Separate fallback path. + iq_s, ik_s = attn.apply_index_qk_norm(idx_q, idx_k) + iq_s, ik_s = attn.rotary_emb(position_ids, [iq_s, ik_s]) + + torch.testing.assert_close(iq_f.contiguous(), iq_s.contiguous(), rtol=5e-2, atol=1e-1) + torch.testing.assert_close(ik_f.contiguous(), ik_s.contiguous(), rtol=5e-2, atol=1e-1) + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 fused QK-norm+RoPE needs CUDA") +def test_minimax_m3_fused_qk_norm_rope_fallbacks(): + """The fused helper returns None (fallback) for non-bf16 or no position_ids.""" + _, model_cfg = _make_fused_qk_norm_rope_test_config() + attn = MiniMaxM3Attention( + model_config=model_cfg, + layer_idx=0, + is_sparse_attention_layer=False, + ) + device = torch.device("cuda") + seq = 3 + total = attn.q_size + 2 * attn.kv_size + position_ids = torch.arange(seq, dtype=torch.int32, device=device) + + # No position_ids means RoPE cannot run, so fall back. + qkv_bf16 = torch.randn(seq, total, dtype=torch.bfloat16, device=device) + assert ( + attn._fused_qk_norm_rope( + qkv_bf16, + None, + num_heads_q=attn.num_heads, + num_heads_k=attn.num_key_value_heads, + num_heads_v=attn.num_key_value_heads, + head_dim=attn.head_dim, + q_norm=attn.q_norm, + k_norm=attn.k_norm, + ) + is None + ) + + # Non-bf16 activations hit the bf16-only guard, so fall back. + qkv_fp16 = torch.randn(seq, total, dtype=torch.float16, device=device) + assert ( + attn._fused_qk_norm_rope( + qkv_fp16, + position_ids, + num_heads_q=attn.num_heads, + num_heads_k=attn.num_key_value_heads, + num_heads_v=attn.num_key_value_heads, + head_dim=attn.head_dim, + q_norm=attn.q_norm, + k_norm=attn.k_norm, + ) + is None + ) + + +@pytest.mark.gpu +@pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 attention construction needs CUDA") +def test_minimax_m3_expect_fused_qk_norm_rope_predicate(): + """bf16 activations plus position_ids require the fused kernel. + + M3 keeps bf16 attention activations under every quantization flavor, so a + runtime fallback there trips the forward assertions. Non-bf16 activations or + missing position_ids relax the expectation. + """ + _, model_cfg = _make_fused_qk_norm_rope_test_config() + attn = MiniMaxM3Attention( + model_config=model_cfg, + layer_idx=3, + is_sparse_attention_layer=True, + disable_index_value=True, + ) + device = torch.device("cuda") + position_ids = torch.arange(4, dtype=torch.int32, device=device) + + # bf16 activations require fusion. + assert attn.attn_activation_dtype == torch.bfloat16 + assert attn._expect_fused_qk_norm_rope(position_ids) is True + + # No position_ids means RoPE cannot run, so a fallback is allowed. + assert attn._expect_fused_qk_norm_rope(None) is False + + # Non-bf16 activations allow a fallback with no assertion. + attn.attn_activation_dtype = torch.float16 + assert attn._expect_fused_qk_norm_rope(position_ids) is False + + @pytest.mark.gpu @pytest.mark.skipif(not _has_cuda(), reason="MiniMax-M3 attention construction needs CUDA") def test_minimax_m3_attention_real_config_index_branch_shapes():