diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py index 7abd08bdd29b..d1599dce7679 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py @@ -1,3 +1,5 @@ +import math + import torch from torch import nn @@ -6,6 +8,86 @@ from tensorrt_llm._torch.models.modeling_utils import register_mapper from tensorrt_llm._torch.utils import split +# 2D block edge of FP8 block-scale tensors (weight_scale_inv), matching +# FP8BlockScalesLinearMethod. +_FP8_BLOCK_SIZE = 128 + + +def grouped_to_dense_in_proj_qkvz_perm(num_k_heads: int, head_k_dim: int, + num_v_heads: int, head_v_dim: int, + tp_size: int) -> torch.Tensor: + """Row permutation from grouped-interleaved in_proj_qkvz to dense [Q|K|V|Z]. + + HF GDN checkpoints pack in_proj_qkvz rows per key-head group as + [q_g | k_g | v_g | z_g]. The GDN mixer consumes the projection as plain + column slices (mixed_qkv = [Q|K|V], then z), which requires the dense + row order [all Q | all K | all V | all Z]. Rows are permuted per TP-rank + chunk so the column-parallel contiguous row split still hands each rank + its own heads. + """ + assert num_k_heads % tp_size == 0 and num_v_heads % tp_size == 0 + ng = num_k_heads // tp_size + ratio = num_v_heads // num_k_heads + dk, rdv = head_k_dim, ratio * head_v_dim + group_size = 2 * dk + 2 * rdv + base = torch.arange(ng).unsqueeze(1) * group_size + q = (base + torch.arange(dk)).flatten() + k = (base + dk + torch.arange(dk)).flatten() + v = (base + 2 * dk + torch.arange(rdv)).flatten() + z = (base + 2 * dk + rdv + torch.arange(rdv)).flatten() + rank_perm = torch.cat([q, k, v, z]) + rank_rows = ng * group_size + return torch.cat([rank_perm + rank * rank_rows for rank in range(tp_size)]) + + +def grouped_to_dense_in_proj_ba_perm(num_k_heads: int, num_v_heads: int, + tp_size: int) -> torch.Tensor: + """Row permutation from grouped-interleaved in_proj_ba to dense [b|a].""" + assert num_k_heads % tp_size == 0 and num_v_heads % tp_size == 0 + ng = num_k_heads // tp_size + ratio = num_v_heads // num_k_heads + base = torch.arange(ng).unsqueeze(1) * (2 * ratio) + b = (base + torch.arange(ratio)).flatten() + a = (base + ratio + torch.arange(ratio)).flatten() + rank_perm = torch.cat([b, a]) + rank_rows = ng * 2 * ratio + return torch.cat([rank_perm + rank * rank_rows for rank in range(tp_size)]) + + +def _rows_to_scale_block_perm(perm: torch.Tensor, block: int) -> torch.Tensor: + """Derive the permutation of `block`-row scale blocks implied by a row + permutation, requiring the row permutation to move whole aligned blocks.""" + if perm.numel() <= block: + # A single scale block covers every row: any within-block row + # permutation leaves the (one-row) scale tensor unchanged. + return torch.zeros(1, dtype=torch.long) + assert perm.numel() % block == 0, ( + f"row permutation of {perm.numel()} rows is not divisible by the " + f"scale block size {block}") + blocks = perm.view(-1, block) + firsts = blocks[:, :1] + assert torch.equal(blocks, firsts + torch.arange(block)) and \ + (firsts % block == 0).all(), ( + "in_proj row permutation must move whole aligned scale blocks to " + "permute block-scale tensors; head dims must be multiples of " + f"{block}") + return firsts.squeeze(1) // block + + +def _permute_rows(tensor: torch.Tensor, perm: torch.Tensor) -> torch.Tensor: + """Reorder dim-0 rows of a (possibly lazy safetensors) tensor. + + Goes through a uint8 view so packed/quantized dtypes (float8, packed FP4 + uint8) reorder identically to plain floats. + """ + t = tensor[...] if not isinstance(tensor, torch.Tensor) else tensor + squeeze = t.dim() == 1 + if squeeze: + t = t.unsqueeze(1) + u8 = t.contiguous().view(torch.uint8) + out = u8[perm].view(t.dtype) + return out.squeeze(1) if squeeze else out + @register_mapper("HF", "Qwen3NextForCausalLM") class Qwen3NextHfWeightMapper(Qwen2MoeHfWeightMapper): @@ -40,6 +122,41 @@ def _duplicate_kv_weights(self, module: nn.Module, new_name: str, return weights + def _permute_in_proj_to_dense(self, key: str, tensor, tp_size: int): + """Reorder one in_proj_qkvz / in_proj_ba tensor to the dense layout. + + Row tensors (weight of any dtype including packed FP4, per-row + scales, bias) are permuted directly; FP8 2D-block ``weight_scale_inv`` + tensors are permuted at scale-block granularity; scalar per-tensor + scales pass through unchanged. + """ + config = self.config.pretrained_config + if ".in_proj_qkvz." in key: + perm = grouped_to_dense_in_proj_qkvz_perm( + config.linear_num_key_heads, config.linear_key_head_dim, + config.linear_num_value_heads, config.linear_value_head_dim, + tp_size) + elif ".in_proj_ba." in key: + perm = grouped_to_dense_in_proj_ba_perm( + config.linear_num_key_heads, config.linear_num_value_heads, + tp_size) + else: + return tensor + + t = tensor[...] if not isinstance(tensor, torch.Tensor) else tensor + if t.dim() == 0 or t.numel() == 1: + return t + rows = perm.numel() + if t.shape[0] == rows: + return _permute_rows(t, perm) + if key.endswith("weight_scale_inv") and \ + t.shape[0] == math.ceil(rows / _FP8_BLOCK_SIZE): + return _permute_rows( + t, _rows_to_scale_block_perm(perm, _FP8_BLOCK_SIZE)) + raise ValueError( + f"Cannot map {key} with shape {tuple(t.shape)} onto the dense " + f"in_proj layout ({rows} rows expected)") + def preprocess_weights(self, weights: dict) -> dict: config = self.config.pretrained_config tp_size = self.config.mapping.tp_size @@ -86,9 +203,13 @@ def preprocess_weights(self, weights: dict) -> dict: w = w.to(torch.float32) new_weights[key] = w elif "in_proj" in key: - # Don't need to split in_proj weight based on the implementation of reference. - # Need to know the reason. - new_weights[key] = weights[name] + # in_proj stays unsplit here (the column-parallel Linear + # splits contiguous row chunks itself); rows are reordered + # from the checkpoint's grouped-interleaved layout to the + # dense per-rank [Q|K|V|Z] / [b|a] layout the GDN mixer + # slices at runtime. + new_weights[key] = self._permute_in_proj_to_dense( + key, weights[name], tp_size) elif "conv1d" in key: w = weights[name] # removing dim(1) because we are using Linear to store conv1d weights diff --git a/tensorrt_llm/_torch/modules/mamba/fuse_elementwise_ops.py b/tensorrt_llm/_torch/modules/mamba/fuse_elementwise_ops.py index 543488445a61..1499f606689a 100644 --- a/tensorrt_llm/_torch/modules/mamba/fuse_elementwise_ops.py +++ b/tensorrt_llm/_torch/modules/mamba/fuse_elementwise_ops.py @@ -24,7 +24,7 @@ def _extract_transpose_prefill_kernel( src_ptr, dst_ptr, num_prefill_tokens, - d_in_proj, + src_stride_seq, d_inner, conv_dim, BLOCK_SEQ: tl.constexpr, @@ -42,9 +42,11 @@ def _extract_transpose_prefill_kernel( conv_mask = conv_offsets < conv_dim mask = seq_mask[:, None] & conv_mask[None, :] - # Cast to int64 to avoid overflow: seq_offsets * d_in_proj can exceed INT32_MAX - # (e.g., 131071 * 22656 = 2,969,544,576 > 2,147,483,647) - src_offsets = seq_offsets[:, None].to(tl.int64) * d_in_proj + d_inner + conv_offsets[None, :] + # Cast to int64 to avoid overflow: seq_offsets * src_stride_seq can exceed + # INT32_MAX (e.g., 131071 * 22656 = 2,969,544,576 > 2,147,483,647) + src_offsets = ( + seq_offsets[:, None].to(tl.int64) * src_stride_seq + d_inner + conv_offsets[None, :] + ) data = tl.load(src_ptr + src_offsets, mask=mask, other=0.0) dst_offsets = conv_offsets[:, None] * num_prefill_tokens + seq_offsets[None, :] @@ -58,11 +60,13 @@ def extract_transpose_prefill_slice( width: int, ) -> torch.Tensor: """ - Extract and transpose a contiguous prefill slice for causal_conv1d_fn. + Extract and transpose a prefill slice for causal_conv1d_fn. - Input: src[num_tokens, num_cols] + Input: src[num_tokens, num_cols], rows contiguous (arbitrary row stride, + so column-slice views of a wider tensor work in place) Output: [width, num_prefill_tokens] """ + assert src.stride(1) == 1 out = torch.empty(width, num_prefill_tokens, dtype=src.dtype, device=src.device) BLOCK_SEQ, BLOCK_CONV = 32, 128 @@ -72,7 +76,7 @@ def extract_transpose_prefill_slice( src, out, num_prefill_tokens, - src.shape[1], + src.stride(0), start_col, width, BLOCK_SEQ, @@ -207,7 +211,7 @@ def _transpose_and_split_qkv_kernel( v_ptr, num_prefill, num_decode, - num_cols, + decode_stride_seq, q_dim: tl.constexpr, k_dim: tl.constexpr, v_dim: tl.constexpr, @@ -216,8 +220,9 @@ def _transpose_and_split_qkv_kernel( ): """Fused transpose-prefill + split-decode into contiguous q, k, v. - Reads prefill from transposed layout [D, T_p] and decode from - row-major layout [T_d, D], writes both into contiguous q/k/v outputs. + Reads prefill from transposed layout [D, T_p] and decode from row-major + layout [T_d, D] with contiguous rows of arbitrary stride, writes both + into contiguous q/k/v outputs. Grid: (num_seq_blocks_total, num_dim_blocks, 3) program_id(2): 0=Q, 1=K, 2=V """ @@ -251,7 +256,7 @@ def _transpose_and_split_qkv_kernel( # Decode: read from decode_ptr[T_d, D] row-major decode_row = seq_offsets - num_prefill - decode_indices = decode_row[:, None].to(tl.int64) * num_cols + ( + decode_indices = decode_row[:, None].to(tl.int64) * decode_stride_seq + ( src_col_offset + dim_offsets[None, :] ).to(tl.int64) decode_data = tl.load( @@ -281,9 +286,13 @@ def transpose_and_split_qkv( """ Fused transpose prefill [D, T_p] + split decode [T_d, D] into contiguous q, k, v. + The decode tensor may be a column-slice view of a wider tensor (contiguous + rows, arbitrary row stride). + Replaces separate transpose_copy_back + split_qkv_contiguous for mixed batches. """ - num_cols, num_prefill = prefill_t.shape + assert decode.stride(-1) == 1 + num_prefill = prefill_t.shape[1] num_decode = decode.shape[0] total_seq = num_prefill + num_decode @@ -303,7 +312,7 @@ def transpose_and_split_qkv( v_flat, num_prefill, num_decode, - num_cols, + decode.stride(0), q_dim, k_dim, v_dim, diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index b59d5a909708..52120a76bdf3 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -38,6 +38,7 @@ transpose_and_split_qkv, ) from .layernorm_gated import RMSNorm as RMSNormGated +from .layernorm_gated import rms_norm_gated_token_major from .mamba2_metadata import Mamba2Metadata @@ -116,117 +117,6 @@ def divide(numerator, denominator): return numerator // denominator -@triton.jit -def fused_qkvzba_split_reshape_cat_kernel( - mixed_qkv, - z, - b, - a, - mixed_qkvz, - mixed_ba, - NUM_HEADS_QK: tl.constexpr, - NUM_HEADS_V: tl.constexpr, - HEAD_QK: tl.constexpr, - HEAD_V: tl.constexpr, -): - i_bs, i_qk = tl.program_id(0), tl.program_id(1) - QKVZ_DIM_T: tl.constexpr = HEAD_QK * 2 + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V * 2 - BA_DIM_T: tl.constexpr = NUM_HEADS_V // NUM_HEADS_QK * 2 - QKV_DIM_T: tl.constexpr = HEAD_QK * 2 + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V - q_end: tl.constexpr = HEAD_QK - blk_q_ptr = ( - mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T + i_qk * QKVZ_DIM_T + tl.arange(0, q_end) - ) - k_end: tl.constexpr = q_end + HEAD_QK - blk_k_ptr = ( - mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T + i_qk * QKVZ_DIM_T + tl.arange(q_end, k_end) - ) - v_end: tl.constexpr = k_end + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V - blk_v_ptr = ( - mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T + i_qk * QKVZ_DIM_T + tl.arange(k_end, v_end) - ) - z_end: tl.constexpr = v_end + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V - blk_z_ptr = ( - mixed_qkvz + i_bs * NUM_HEADS_QK * QKVZ_DIM_T + i_qk * QKVZ_DIM_T + tl.arange(v_end, z_end) - ) - blk_q_st_ptr = ( - mixed_qkv + i_bs * NUM_HEADS_QK * QKV_DIM_T + i_qk * HEAD_QK + tl.arange(0, HEAD_QK) - ) - blk_k_st_ptr = ( - mixed_qkv - + i_bs * NUM_HEADS_QK * QKV_DIM_T - + NUM_HEADS_QK * HEAD_QK - + i_qk * HEAD_QK - + tl.arange(0, HEAD_QK) - ) - blk_v_st_ptr = ( - mixed_qkv - + i_bs * NUM_HEADS_QK * QKV_DIM_T - + NUM_HEADS_QK * HEAD_QK * 2 - + i_qk * HEAD_V * NUM_HEADS_V // NUM_HEADS_QK - + tl.arange(0, HEAD_V * NUM_HEADS_V // NUM_HEADS_QK) - ) - blk_z_st_ptr = ( - z - + i_bs * NUM_HEADS_V * HEAD_V - + i_qk * HEAD_V * NUM_HEADS_V // NUM_HEADS_QK - + tl.arange(0, HEAD_V * NUM_HEADS_V // NUM_HEADS_QK) - ) - tl.store(blk_q_st_ptr, tl.load(blk_q_ptr)) - tl.store(blk_k_st_ptr, tl.load(blk_k_ptr)) - tl.store(blk_v_st_ptr, tl.load(blk_v_ptr)) - tl.store(blk_z_st_ptr, tl.load(blk_z_ptr)) - b_end: tl.constexpr = NUM_HEADS_V // NUM_HEADS_QK - a_end: tl.constexpr = b_end + NUM_HEADS_V // NUM_HEADS_QK - for i in tl.static_range(b_end): - blk_b_ptr = mixed_ba + i_bs * NUM_HEADS_QK * BA_DIM_T + i_qk * BA_DIM_T + i - blk_b_st_ptr = b + i_bs * NUM_HEADS_V + i_qk * NUM_HEADS_V // NUM_HEADS_QK + i - tl.store(blk_b_st_ptr, tl.load(blk_b_ptr)) - for i in tl.static_range(b_end, a_end): - blk_a_ptr = mixed_ba + i_bs * NUM_HEADS_QK * BA_DIM_T + i_qk * BA_DIM_T + i - blk_a_st_ptr = a + i_bs * NUM_HEADS_V + i_qk * NUM_HEADS_V // NUM_HEADS_QK + (i - b_end) - tl.store(blk_a_st_ptr, tl.load(blk_a_ptr)) - - -def fused_qkvzba_split_reshape_cat( - mixed_qkvz, - mixed_ba, - num_heads_qk, - num_heads_v, - head_qk, - head_v, -): - batch, seq_len = mixed_qkvz.shape[0], 1 - qkv_dim_t = num_heads_qk * head_qk * 2 + num_heads_v * head_v - batch_seq = batch * seq_len - - # Directly allocate output tensors in their final shapes (no intermediate buffers) - mixed_qkv = torch.empty( - (batch_seq, qkv_dim_t), dtype=mixed_qkvz.dtype, device=mixed_qkvz.device - ) - z = torch.empty( - (batch_seq, num_heads_v, head_v), dtype=mixed_qkvz.dtype, device=mixed_qkvz.device - ) - b = torch.empty((batch_seq, num_heads_v), dtype=mixed_ba.dtype, device=mixed_ba.device) - a = torch.empty((batch_seq, num_heads_v), dtype=mixed_ba.dtype, device=mixed_ba.device) - grid = (batch * seq_len, num_heads_qk) - fused_qkvzba_split_reshape_cat_kernel[grid]( - mixed_qkv, - z, - b, - a, - mixed_qkvz, - mixed_ba, - num_heads_qk, - num_heads_v, - head_qk, - head_v, - num_warps=1, - num_stages=3, - ) - return mixed_qkv, z, b, a - - # g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) @triton.jit def fused_gdn_gating_kernel( @@ -234,23 +124,26 @@ def fused_gdn_gating_kernel( A_log, a, dt_bias, - seq_len, + stride_a_row, NUM_HEADS: tl.constexpr, beta: tl.constexpr, threshold: tl.constexpr, BLK_HEADS: tl.constexpr, ): - i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_d = tl.program_id(0), tl.program_id(1) head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS) - off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off + # a may be a row-strided view sliced out of the packed ba projection; + # g is always allocated packed. + off_a = i_b * stride_a_row + head_off + off_g = i_b * NUM_HEADS + head_off mask = head_off < NUM_HEADS blk_A_log = tl.load(A_log + head_off, mask=mask) - blk_a = tl.load(a + off, mask=mask) + blk_a = tl.load(a + off_a, mask=mask) blk_bias = tl.load(dt_bias + head_off, mask=mask) x = blk_a.to(tl.float32) + blk_bias.to(tl.float32) softplus_x = tl.where(beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x) blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x - tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask) + tl.store(g + off_g, blk_g.to(g.dtype.element_ty), mask=mask) def fused_gdn_gating( @@ -261,11 +154,10 @@ def fused_gdn_gating( threshold: float = 20.0, ) -> torch.Tensor: batch, num_heads = a.shape - seq_len = 1 - grid = (batch, seq_len, triton.cdiv(num_heads, 8)) - g = torch.empty_like(a, dtype=torch.float32) + grid = (batch, triton.cdiv(num_heads, 8)) + g = torch.empty(batch, num_heads, dtype=torch.float32, device=a.device) fused_gdn_gating_kernel[grid]( - g, A_log, a, dt_bias, seq_len, num_heads, beta, threshold, 8, num_warps=1 + g, A_log, a, dt_bias, a.stride(0), num_heads, beta, threshold, 8, num_warps=1 ) return g @@ -278,30 +170,35 @@ def fused_gdn_gating_with_sigmoid_kernel( a, dt_bias, b, - seq_len, + stride_a_row, + stride_b_row, NUM_HEADS: tl.constexpr, sp_beta: tl.constexpr, threshold: tl.constexpr, BLK_HEADS: tl.constexpr, ): """Fuse gdn_gating + sigmoid(b) into one kernel.""" - i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_d = tl.program_id(0), tl.program_id(1) head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS) - off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off + # a/b may be row-strided views sliced out of the packed ba projection; + # g/beta_out are always allocated packed. + off_a = i_b * stride_a_row + head_off + off_b = i_b * stride_b_row + head_off + off_out = i_b * NUM_HEADS + head_off mask = head_off < NUM_HEADS blk_A_log = tl.load(A_log + head_off, mask=mask) - blk_a = tl.load(a + off, mask=mask) + blk_a = tl.load(a + off_a, mask=mask) blk_bias = tl.load(dt_bias + head_off, mask=mask) x = blk_a.to(tl.float32) + blk_bias.to(tl.float32) softplus_x = tl.where( sp_beta * x <= threshold, (1 / sp_beta) * tl.log(1 + tl.exp(sp_beta * x)), x ) blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x - tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask) + tl.store(g + off_out, blk_g.to(g.dtype.element_ty), mask=mask) # sigmoid(b) - blk_b = tl.load(b + off, mask=mask) + blk_b = tl.load(b + off_b, mask=mask) blk_beta = tl.sigmoid(blk_b.to(tl.float32)) - tl.store(beta_out + off, blk_beta.to(beta_out.dtype.element_ty), mask=mask) + tl.store(beta_out + off_out, blk_beta.to(beta_out.dtype.element_ty), mask=mask) def fused_gdn_gating_with_sigmoid( @@ -314,14 +211,13 @@ def fused_gdn_gating_with_sigmoid( ) -> tuple[torch.Tensor, torch.Tensor]: """Fused GDN gating + sigmoid: compute g and beta in one kernel launch.""" batch, num_heads = a.shape - seq_len = 1 - grid = (batch, seq_len, triton.cdiv(num_heads, 8)) - g = torch.empty_like(a, dtype=torch.float32) + grid = (batch, triton.cdiv(num_heads, 8)) + g = torch.empty(batch, num_heads, dtype=torch.float32, device=a.device) # Allocate beta in fp32 since (1) the kernel already computes sigmoid in fp32 # and was previously casting back to b.dtype only to be re-cast to fp32 by the # FlashInfer GDN prefill wrapper, and (2) the Triton chunk_gated_delta_rule # path also accepts fp32 beta. Eliminates a redundant cast in the FI hot path. - beta_out = torch.empty_like(b, dtype=torch.float32) + beta_out = torch.empty(batch, num_heads, dtype=torch.float32, device=b.device) fused_gdn_gating_with_sigmoid_kernel[grid]( g, beta_out, @@ -329,7 +225,8 @@ def fused_gdn_gating_with_sigmoid( a, dt_bias, b, - seq_len, + a.stride(0), + b.stride(0), num_heads, sp_beta, threshold, @@ -380,6 +277,7 @@ def __init__( self.num_v_heads_per_tp = divide(self.num_v_heads, self.attn_tp_size) self.key_dim_per_tp = self.head_k_dim * self.num_k_heads_per_tp self.value_dim_per_tp = self.head_v_dim * self.num_v_heads_per_tp + self.conv_dim_per_tp = self.key_dim_per_tp * 2 + self.value_dim_per_tp self.conv_kernel_size = config.linear_conv_kernel_dim self.layer_idx = layer_idx @@ -492,66 +390,6 @@ def __init__( self.event_dict = {key: torch.cuda.Event() for key in [EventType.Main, EventType.Attention]} self.aux_stream = aux_stream - def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba): - """ - Derives `query`, `key` and `value` tensors from `mixed_qkvzba`. - """ - batch_size = mixed_qkvz.size(0) - num_k_heads_local = self.num_k_heads // self.attn_tp_size - num_v_heads_local = self.num_v_heads // self.attn_tp_size - heads_ratio = self.num_v_heads // self.num_k_heads - - # Reshape qkvz: [b, d] -> [b, ng, (2*hk + 2*np/ng*hv)] - qkvz_dim_per_head = self.head_k_dim * 2 + self.head_v_dim * heads_ratio * 2 - mixed_qkvz = mixed_qkvz.view(batch_size, num_k_heads_local, qkvz_dim_per_head) - - # Reshape ba: [b, d] -> [b, ng, 2*np/ng] - mixed_ba = mixed_ba.view(batch_size, num_k_heads_local, heads_ratio * 2) - - # Direct slicing instead of torch.split for better performance - # Compute split boundaries once - q_end = self.head_k_dim - k_end = q_end + self.head_k_dim - v_end = k_end + heads_ratio * self.head_v_dim - z_end = v_end + heads_ratio * self.head_v_dim - - # Slice qkvz components: [b, ng, dim] -> individual components - query = mixed_qkvz[..., :q_end] - key = mixed_qkvz[..., q_end:k_end] - - # When heads_ratio == 1, ng == num_v_heads_local, so view works directly. - # When heads_ratio > 1 (dense models), the last-dim slice is - # [b, ng, ratio*hv] and we need [b, ng*ratio, hv]. A plain view - # fails because the slice is not contiguous in the packed qkvz - # tensor. Adding .contiguous() before view is equivalent to - # reshape but makes the copy explicit and avoids a hidden perf - # drop. An alternative zero-copy path would require changing - # the packing layout, which is a larger refactor. - if heads_ratio == 1: - value = mixed_qkvz[..., k_end:v_end] - z = mixed_qkvz[..., v_end:z_end] - else: - value = ( - mixed_qkvz[..., k_end:v_end] - .contiguous() - .view(batch_size, num_v_heads_local, self.head_v_dim) - ) - z = ( - mixed_qkvz[..., v_end:z_end] - .contiguous() - .view(batch_size, num_v_heads_local, self.head_v_dim) - ) - - # Slice ba components: [b, ng, 2*np/ng] -> [b, np] each - if heads_ratio == 1: - b = mixed_ba[..., 0] - a = mixed_ba[..., 1] - else: - b = mixed_ba[..., :heads_ratio].contiguous().view(batch_size, num_v_heads_local) - a = mixed_ba[..., heads_ratio:].contiguous().view(batch_size, num_v_heads_local) - - return query, key, value, z, b, a - def _compute_tokenwise_inputs(self, hidden_states: torch.Tensor): def _compute_projected_states_qkvz(): return self.in_proj_qkvz(hidden_states) @@ -568,22 +406,19 @@ def _compute_projected_states_ba(): disable_on_compile=True, ) - # Use fused kernel when possible to avoid elementwise ops - if self.num_v_heads // self.num_k_heads in [1, 2, 4]: # and is_cuda_graph: - mixed_qkv, z, b, a = fused_qkvzba_split_reshape_cat( - projected_states_qkvz, - projected_states_ba, - triton.cdiv(self.num_k_heads, self.attn_tp_size), - triton.cdiv(self.num_v_heads, self.attn_tp_size), - self.head_k_dim, - self.head_v_dim, - ) - else: - query, key, value, z, b, a = self.fix_query_key_value_ordering( - projected_states_qkvz, projected_states_ba - ) - query, key, value = map(lambda x: x.reshape(x.shape[0], -1), (query, key, value)) - mixed_qkv = torch.cat((query, key, value), dim=-1) + # The weight mapper reorders in_proj rows into the dense per-rank + # layouts [Q|K|V|Z] and [b|a] (see grouped_to_dense_in_proj_qkvz_perm), + # so every component is a plain column slice of the projection — + # no split/reshape kernel. Downstream consumers (causal_conv1d, + # the GDN decode kernels, the gated norm) read these row-strided + # views in place. + num_tokens = projected_states_qkvz.shape[0] + mixed_qkv = projected_states_qkvz[:, : self.conv_dim_per_tp] + z = projected_states_qkvz[:, self.conv_dim_per_tp :].view( + num_tokens, self.num_v_heads_per_tp, self.head_v_dim + ) + b = projected_states_ba[:, : self.num_v_heads_per_tp] + a = projected_states_ba[:, self.num_v_heads_per_tp :] return mixed_qkv, z, a, b @@ -593,12 +428,14 @@ def _postprocess_gdn_output( z: torch.Tensor, all_reduce_params: Optional[AllReduceParams] = None, ): - z_shape_og = z.shape - attn_out = attn_out.reshape(-1, attn_out.shape[-1]) - z = z.reshape(-1, z.shape[-1]) - attn_out = self.norm(attn_out, z) - attn_out = attn_out.reshape(z_shape_og) - attn_out = attn_out.reshape(*attn_out.shape[:-2], -1) + # z is a [num_tokens, num_v_heads, head_v_dim] view of the in_proj + # output whose (heads, head_dim) block is contiguous per token; the + # gated norm reads it through its token stride instead of packing a + # copy. + attn_out = rms_norm_gated_token_major( + attn_out.reshape(-1, self.head_v_dim), z, self.norm.weight, self.norm.eps + ) + attn_out = attn_out.view(-1, self.value_dim_per_tp) return self.out_proj(attn_out, all_reduce_params=all_reduce_params) def forward_decode( diff --git a/tensorrt_llm/_torch/modules/mamba/layernorm_gated.py b/tensorrt_llm/_torch/modules/mamba/layernorm_gated.py index d968d8745e3d..1aaa276423c2 100644 --- a/tensorrt_llm/_torch/modules/mamba/layernorm_gated.py +++ b/tensorrt_llm/_torch/modules/mamba/layernorm_gated.py @@ -109,6 +109,119 @@ def _layer_norm_fwd_1pass_kernel( tl.store(Y + cols, y, mask=mask) +# Rows per program of the multi-row gated-RMSNorm kernel. At the GDN decode +# shape (thousands of 128-element rows) one row per CTA leaves the kernel +# launch-limited; 4 rows with 4 warps reproduces the single-row kernel's +# reduction order (bitwise-identical output) at ~2x the throughput. +_MULTIROW_ROWS = 4 +_MULTIROW_NUM_WARPS = 4 +_MULTIROW_MAX_N = 256 + + +@triton.jit +def _rms_norm_gated_fwd_multirow_kernel( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + Z, # pointer to the gate branch + Rstd, # pointer to the 1/std + stride_x_row, + stride_y_row, + stride_z_tok, + M, # number of rows in X + eps, + N: tl.constexpr, # row length; power of two, whole row per program + ROWS: tl.constexpr, # rows per program + HEADS_PER_TOK: tl.constexpr, +): + """rmsnorm(x) * silu(z), several short rows per program. + + Z is addressed token-major: row r reads z at + (r // HEADS_PER_TOK) * stride_z_tok + (r % HEADS_PER_TOK) * N. With + HEADS_PER_TOK == 1 and stride_z_tok == z's row stride this is a plain + [M, N] z; with HEADS_PER_TOK == heads it reads a [num_tokens, heads, N] + view whose (heads, N) block is contiguous per token, e.g. a column slice + of a wider projection. + """ + rows = tl.program_id(0) * ROWS + tl.arange(0, ROWS) + row_mask = rows < M + cols = tl.arange(0, N) + mask2d = row_mask[:, None] + x_off = rows[:, None].to(tl.int64) * stride_x_row + cols[None, :] + x = tl.load(X + x_off, mask=mask2d, other=0.0).to(tl.float32) + var = tl.sum(x * x, axis=1) / N + rstd = 1.0 / tl.sqrt(var + eps) + tl.store(Rstd + rows, rstd, mask=row_mask) + w = tl.load(W + cols).to(tl.float32) + y = x * rstd[:, None] * w[None, :] + tok = rows // HEADS_PER_TOK + head = rows % HEADS_PER_TOK + z_off = (tok[:, None].to(tl.int64) * stride_z_tok + head[:, None] * N + + cols[None, :]) + z = tl.load(Z + z_off, mask=mask2d, other=0.0).to(tl.float32) + y *= z * tl.sigmoid(z) + y_off = rows[:, None].to(tl.int64) * stride_y_row + cols[None, :] + tl.store(Y + y_off, y.to(Y.dtype.element_ty), mask=mask2d) + + +def _multirow_gated_rmsnorm_eligible(N, ngroups, bias, z, norm_before_gate, + is_rms_norm): + return (is_rms_norm and norm_before_gate and z is not None and bias is None + and ngroups == 1 and N <= _MULTIROW_MAX_N and (N & (N - 1)) == 0) + + +def rms_norm_gated_token_major(x, z, weight, eps, out=None): + """rmsnorm(x) * silu(z) with z read in place from a 3D token-major view. + + x: [num_tokens * heads, N] with contiguous rows. z: [num_tokens, heads, N] + whose (heads, N) block is contiguous per token and whose token stride is + arbitrary (e.g. a column slice of a wider per-token projection). Falls + back to the generic kernel on a packed copy of z when the shape is not + eligible for the multi-row kernel. + """ + M, N = x.shape + num_tokens, heads, n_z = z.shape + assert n_z == N and num_tokens * heads == M, ( + f"z shape {tuple(z.shape)} does not match x shape {tuple(x.shape)}") + weight = weight.contiguous() + eligible = (x.stride(-1) == 1 and z.stride(2) == 1 and z.stride(1) == N + and N <= _MULTIROW_MAX_N and (N & (N - 1)) == 0) + if not eligible: + y, _, _ = _layer_norm_fwd( + x, + weight, + None, + eps, + z=z.reshape(M, N), + out=out, + norm_before_gate=True, + is_rms_norm=True, + ) + return y + if out is None: + out = torch.empty_like(x) + rstd = torch.empty((M, ), dtype=torch.float32, device=x.device) + grid = (triton.cdiv(M, _MULTIROW_ROWS), ) + with torch.cuda.device(x.device.index): + _rms_norm_gated_fwd_multirow_kernel[grid]( + x, + out, + weight, + z, + rstd, + x.stride(0), + out.stride(0), + z.stride(0), + M, + eps, + N=N, + ROWS=_MULTIROW_ROWS, + HEADS_PER_TOK=heads, + num_warps=_MULTIROW_NUM_WARPS, + ) + return out + + def _layer_norm_fwd( x, weight, @@ -143,6 +256,27 @@ def _layer_norm_fwd( mean = (torch.empty((ngroups * M, ), dtype=torch.float32, device=x.device) if not is_rms_norm else None) rstd = torch.empty((ngroups * M, ), dtype=torch.float32, device=x.device) + if _multirow_gated_rmsnorm_eligible(group_size, ngroups, bias, z, + norm_before_gate, is_rms_norm): + grid = (triton.cdiv(M, _MULTIROW_ROWS), ) + with torch.cuda.device(x.device.index): + _rms_norm_gated_fwd_multirow_kernel[grid]( + x, + out, + weight, + z, + rstd, + x.stride(0), + out.stride(0), + z.stride(0), + M, + eps, + N=group_size, + ROWS=_MULTIROW_ROWS, + HEADS_PER_TOK=1, + num_warps=_MULTIROW_NUM_WARPS, + ) + return out, mean, rstd # Less than 64KB per feature: enqueue fused kernel MAX_FUSED_SIZE = 65536 // x.element_size() BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size)) diff --git a/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py b/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py index 56a47a8e3905..26fd4a65e3b6 100644 --- a/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py +++ b/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py @@ -12,7 +12,9 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for GDN kernel optimizations: fused gating+sigmoid, split_qkv, transpose_and_split.""" +"""Unit tests for GDN kernel optimizations: fused gating+sigmoid, split_qkv, +transpose_and_split, the dense in_proj row permutation, and the multi-row +gated RMSNorm.""" import pytest import torch @@ -34,6 +36,160 @@ def _ref_gdn_gating_with_sigmoid(A_log, a, dt_bias, b, beta=1.0, threshold=20.0) return g, beta_out +def _ref_split_grouped_qkvz(y, num_groups, head_k_dim, head_v_dim, heads_ratio): + """Split a grouped-interleaved in_proj_qkvz output ([q|k|v|z] per group).""" + bsz = y.shape[0] + g = y.view(bsz, num_groups, 2 * head_k_dim + 2 * heads_ratio * head_v_dim) + q = g[..., :head_k_dim].reshape(bsz, num_groups * head_k_dim) + k = g[..., head_k_dim : 2 * head_k_dim].reshape(bsz, num_groups * head_k_dim) + v = g[..., 2 * head_k_dim : 2 * head_k_dim + heads_ratio * head_v_dim].reshape( + bsz, num_groups * heads_ratio * head_v_dim + ) + z = g[..., 2 * head_k_dim + heads_ratio * head_v_dim :].reshape( + bsz, num_groups * heads_ratio, head_v_dim + ) + return q, k, v, z + + +def _ref_split_grouped_ba(y, num_groups, heads_ratio): + bsz = y.shape[0] + g = y.view(bsz, num_groups, 2 * heads_ratio) + b = g[..., :heads_ratio].reshape(bsz, num_groups * heads_ratio) + a = g[..., heads_ratio:].reshape(bsz, num_groups * heads_ratio) + return b, a + + +# ---- Tests for the dense in_proj row permutation ---- + + +@pytest.mark.parametrize( + "num_k_heads,num_v_heads,head_k_dim,head_v_dim", + [(16, 32, 128, 128), (16, 16, 128, 128), (8, 32, 64, 64)], +) +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_grouped_to_dense_in_proj_perm(num_k_heads, num_v_heads, head_k_dim, head_v_dim, tp_size): + """Permuting in_proj rows to the dense layout must make plain column + slices of each TP rank's projection output reproduce the grouped split.""" + from tensorrt_llm._torch.models.checkpoints.hf.qwen3_next_weight_mapper import ( + _permute_rows, + grouped_to_dense_in_proj_ba_perm, + grouped_to_dense_in_proj_qkvz_perm, + ) + + torch.manual_seed(42) + ratio = num_v_heads // num_k_heads + rows = 2 * num_k_heads * head_k_dim + 2 * num_v_heads * head_v_dim + w = torch.randn(rows, 32, dtype=torch.bfloat16) + w_ba = torch.randn(2 * num_v_heads, 32, dtype=torch.bfloat16) + x = torch.randn(4, 32, dtype=torch.float32) + + perm = grouped_to_dense_in_proj_qkvz_perm( + num_k_heads, head_k_dim, num_v_heads, head_v_dim, tp_size + ) + perm_ba = grouped_to_dense_in_proj_ba_perm(num_k_heads, num_v_heads, tp_size) + assert torch.equal(perm.sort().values, torch.arange(rows)) + w_dense = _permute_rows(w, perm) + w_ba_dense = _permute_rows(w_ba, perm_ba) + + for rank in range(tp_size): + ng = num_k_heads // tp_size + shard = rows // tp_size + y_grouped = x @ w[rank * shard : (rank + 1) * shard].T.float() + y_dense = x @ w_dense[rank * shard : (rank + 1) * shard].T.float() + q_ref, k_ref, v_ref, z_ref = _ref_split_grouped_qkvz( + y_grouped, ng, head_k_dim, head_v_dim, ratio + ) + k_end = 2 * ng * head_k_dim + qkv_dim = k_end + ng * ratio * head_v_dim + assert torch.equal(y_dense[:, : ng * head_k_dim], q_ref) + assert torch.equal(y_dense[:, ng * head_k_dim : k_end], k_ref) + assert torch.equal(y_dense[:, k_end:qkv_dim], v_ref) + assert torch.equal(y_dense[:, qkv_dim:].view(4, ng * ratio, head_v_dim), z_ref) + + shard_ba = 2 * num_v_heads // tp_size + yb_grouped = x @ w_ba[rank * shard_ba : (rank + 1) * shard_ba].T.float() + yb_dense = x @ w_ba_dense[rank * shard_ba : (rank + 1) * shard_ba].T.float() + b_ref, a_ref = _ref_split_grouped_ba(yb_grouped, ng, ratio) + assert torch.equal(yb_dense[:, : ng * ratio], b_ref) + assert torch.equal(yb_dense[:, ng * ratio :], a_ref) + + +def test_in_proj_perm_quantized_dtypes(): + """Row permutation must be dtype-agnostic (fp8, packed-uint8, 1-D) and + consistent with FP8 2D-block scale permutation.""" + from tensorrt_llm._torch.models.checkpoints.hf.qwen3_next_weight_mapper import ( + _permute_rows, + _rows_to_scale_block_perm, + grouped_to_dense_in_proj_qkvz_perm, + ) + + torch.manual_seed(42) + perm = grouped_to_dense_in_proj_qkvz_perm(16, 128, 32, 128, 1) + w_fp8 = torch.randn(12288, 64).to(torch.float8_e4m3fn) + assert torch.equal(_permute_rows(w_fp8, perm).float(), w_fp8.float()[perm]) + w_packed = torch.randint(0, 255, (12288, 1024), dtype=torch.uint8) + assert torch.equal(_permute_rows(w_packed, perm), w_packed[perm]) + w_1d = torch.randn(12288, dtype=torch.bfloat16) + assert torch.equal(_permute_rows(w_1d, perm), w_1d[perm]) + + # Expanding block scales to rows then permuting rows must equal permuting + # scale blocks then expanding. + scale = torch.randn(12288 // 128, 7) + block_perm = _rows_to_scale_block_perm(perm, 128) + lhs = scale.repeat_interleave(128, dim=0)[perm] + rhs = _permute_rows(scale, block_perm).repeat_interleave(128, dim=0) + assert torch.equal(lhs, rhs) + + +# ---- Tests for the multi-row gated RMSNorm ---- + + +def _ref_gated_rmsnorm(x, w, z, eps): + xf = x.float() + rstd = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) + y = xf * rstd * w.float() + zf = z.float() + return (y * zf * torch.sigmoid(zf)).to(x.dtype) + + +@skip_no_cuda +@pytest.mark.parametrize( + "num_tokens,heads,N", + [(8192, 32, 128), (7, 32, 128), (1, 1, 128), (333, 4, 64), (1024, 16, 256)], +) +def test_rms_norm_gated_token_major(num_tokens, heads, N): + """Token-major z (a column-slice view of a wider projection) must match + the reference on both the multi-row fast path and the generic fallback.""" + from tensorrt_llm._torch.modules.mamba.layernorm_gated import ( + _layer_norm_fwd, + rms_norm_gated_token_major, + ) + + torch.manual_seed(42) + device = torch.device("cuda") + M = num_tokens * heads + x = torch.randn(M, N, dtype=torch.bfloat16, device=device) + w = torch.rand(N, dtype=torch.bfloat16, device=device) + 0.5 + wide = torch.randn(num_tokens, heads * N + 512, dtype=torch.bfloat16, device=device) + z = wide[:, 512:].view(num_tokens, heads, N) + + y = rms_norm_gated_token_major(x, z, w, 1e-6) + ref = _ref_gated_rmsnorm(x, w, z.reshape(M, N), 1e-6) + torch.testing.assert_close(y, ref, rtol=1e-2, atol=1e-2) + + # The dense-z dispatch of the generic entry point must agree bitwise. + y_dense, _, _ = _layer_norm_fwd( + x, + w, + None, + 1e-6, + z=z.reshape(M, N).contiguous(), + norm_before_gate=True, + is_rms_norm=True, + ) + assert torch.equal(y, y_dense) + + def _ref_split_qkv_contiguous(mixed_qkv, q_dim, k_dim, v_dim): """Reference: torch.split + contiguous.""" q, k, v = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1) @@ -188,3 +344,63 @@ def test_transpose_and_split_qkv( torch.testing.assert_close(q_out.view(total_seq, -1), q_ref, rtol=0, atol=0) torch.testing.assert_close(k_out.view(total_seq, -1), k_ref, rtol=0, atol=0) torch.testing.assert_close(v_out.view(total_seq, -1), v_ref, rtol=0, atol=0) + + +# ---- Strided (column-slice view) inputs, as produced by the dense in_proj ---- + + +@skip_no_cuda +def test_gdn_gating_strided_views(): + """a/b sliced out of the packed ba projection must match packed inputs.""" + from tensorrt_llm._torch.modules.mamba.gdn_mixer import ( + fused_gdn_gating, + fused_gdn_gating_with_sigmoid, + ) + + torch.manual_seed(42) + device = torch.device("cuda") + ba = torch.randn(300, 64, dtype=torch.bfloat16, device=device) + b_view, a_view = ba[:, :32], ba[:, 32:] + + A_log = torch.randn(32, dtype=torch.float32, device=device) + dt_bias = torch.randn(32, dtype=torch.float32, device=device) + + g_v, beta_v = fused_gdn_gating_with_sigmoid(A_log, a_view, dt_bias, b_view) + g_c, beta_c = fused_gdn_gating_with_sigmoid( + A_log, a_view.contiguous(), dt_bias, b_view.contiguous() + ) + assert torch.equal(g_v, g_c) and torch.equal(beta_v, beta_c) + g_ref, beta_ref = _ref_gdn_gating_with_sigmoid(A_log, a_view, dt_bias, b_view) + torch.testing.assert_close(g_v, g_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(beta_v, beta_ref.float(), rtol=1e-2, atol=1e-2) + + assert torch.equal( + fused_gdn_gating(A_log, a_view, dt_bias), + fused_gdn_gating(A_log, a_view.contiguous(), dt_bias), + ) + + +@skip_no_cuda +def test_transpose_helpers_strided_views(): + """The prefill transpose/split helpers must read column-slice views in place.""" + from tensorrt_llm._torch.modules.mamba.fuse_elementwise_ops import ( + extract_transpose_prefill_slice, + transpose_and_split_qkv, + ) + + torch.manual_seed(42) + device = torch.device("cuda") + wide = torch.randn(100, 1536, dtype=torch.bfloat16, device=device) + view = wide[:, :1024] # rows contiguous, row stride 1536 + + out = extract_transpose_prefill_slice(view, 100, 0, 1024) + assert torch.equal(out, view.T.contiguous()) + + prefill_t = torch.randn(1024, 40, dtype=torch.bfloat16, device=device) + decode_view = wide[40:80, :1024] + outs_view = transpose_and_split_qkv(prefill_t, decode_view, 256, 256, 512, 16, 16, 16, 32) + outs_packed = transpose_and_split_qkv( + prefill_t, decode_view.contiguous(), 256, 256, 512, 16, 16, 16, 32 + ) + for got, ref in zip(outs_view, outs_packed): + assert torch.equal(got, ref)