diff --git a/requirements.txt b/requirements.txt index 6b9b856..bebc374 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,4 @@ modelscope peft>=0.11,<0.21 safetensors tqdm -transformers>=4.33,<5.15.0 +transformers>=4.33,<5.17.0 diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index 7ce61f8..5075dc2 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -5,10 +5,11 @@ import torch.nn.functional as F from contextlib import contextmanager from copy import deepcopy -from megatron.core import mpu from megatron.core.extensions.transformer_engine import TEColumnParallelLinear, TENorm, TERowParallelLinear from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.ssm.gated_delta_net import GatedDeltaNetSubmodules +from megatron.core.tensor_parallel import gather_from_sequence_parallel_region from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.spec_utils import ModuleSpec @@ -16,10 +17,11 @@ from transformers.utils import is_torch_npu_available from typing import List, Optional -from mcore_bridge.utils import get_local_layer_specs, get_logger +from mcore_bridge.utils import get_env_args, get_local_layer_specs, get_logger +from mcore_bridge.utils.megatron_utils import reconstruct_tensor_cp -from ..modules import (GatedDeltaNet, QSAIndexer, Qwen4ExpTextGatedResidual, Qwen4ExpTextPLELayer, TransformerBlock, - TransformerLayer) +from ..modules import (GatedDeltaNet, QSAIndexer, QSASparseCoreAttention, Qwen4ExpTextGatedResidual, + Qwen4ExpTextPLELayer, TransformerBlock, TransformerLayer, qsa_sparse_supported) from ..register import ModelLoader from .qwen3_next import Qwen3NextBridge, Qwen3NextRMSNorm, Qwen3NextSelfAttention @@ -42,7 +44,7 @@ def _apply_gated_norm(self, x, gate): x = x.reshape(-1, x.shape[-1]) y = self.out_norm(x) gate = gate.reshape(-1, gate.shape[-1]) - output_gate_type = getattr(self.config, 'output_gate_type', None) + output_gate_type = self.config.output_gate_type gate_act = torch.sigmoid if output_gate_type == 'sigmoid' else F.silu y = y * gate_act(gate.float()) return y.to(x_dtype) @@ -57,14 +59,18 @@ def __init__(self, config, submodules, layer_number: int = 1, **kwargs): self.ple = Qwen4ExpTextPLELayer( config, config.ple_layer_ids.index(self.layer_number), pg_collection=self.pg_collection) is_linear_attention = config.linear_attention_freq[self.layer_number - 1] - if not is_linear_attention and getattr(config, 'indexer_n_heads', None) is not None: + if not is_linear_attention and config.indexer_n_heads is not None: self.self_attention.indexer = QSAIndexer(config) + if qsa_sparse_supported(config.kv_channels or 0): + attn = self.self_attention + attn.core_attention = QSASparseCoreAttention( + attn.core_attention, config, softmax_scale=getattr(config, 'softmax_scale', None)) self.attn_hyper_connection = Qwen4ExpTextGatedResidual(config) self.mlp_hyper_connection = Qwen4ExpTextGatedResidual(config) def forward(self, hidden_states: torch.Tensor, **kwargs): attention_mask = kwargs.get('attention_mask') - packed_seq_params = kwargs.get('packed_seq_params') + packed_seq_params: PackedSeqParams = kwargs.get('packed_seq_params') attn_kwargs = dict( attention_mask=attention_mask, inference_context=kwargs.get('inference_context'), @@ -82,11 +88,14 @@ def forward(self, hidden_states: torch.Tensor, **kwargs): # attention sub-block (mirrors transformers Qwen4ExpTextDecoderLayer.forward) hidden_states, hyper_input, injection_weights = self.attn_hyper_connection(hidden_states) - qsa_mask = self._qsa_select_mask(hidden_states, attn_kwargs) - if qsa_mask is not None: - # failed to full atention - attn_kwargs = dict(attn_kwargs, attention_mask=qsa_mask) - with self._patch_apply_rotary_pos_emb(), self._qsa_arbitrary_mask(qsa_mask is not None): + qsa_selection, sparse = self._qsa_select(hidden_states, attn_kwargs, kwargs.get('position_ids')) + if qsa_selection is not None: + # sparse: int64 indices consumed by QSASparseCoreAttention; mask + # fallback: bool mask consumed by TE under attn_mask_type=arbitrary. + attn_kwargs = dict(attn_kwargs, attention_mask=qsa_selection) + # `arbitrary` mask type is only needed for the bool-mask (TE) path; the + # sparse kernel reads the indices and ignores attn_mask_type. + with self._patch_apply_rotary_pos_emb(), self._qsa_arbitrary_mask(qsa_selection is not None and not sparse): hidden_states, _ = self.self_attention(hidden_states=hidden_states, **attn_kwargs) injection = hidden_states.unsqueeze(-2) * injection_weights.unsqueeze(-1) hidden_states = hyper_input + injection.flatten(-2) @@ -126,35 +135,130 @@ def _qsa_arbitrary_mask(self, enabled: bool): for t, old in saved: t.attn_mask_type = old - def _warn_qsa_fallback_once(self, reason: str) -> None: - # warning_once dedupes on the message, so every QSA layer can call this and - # the user still sees it exactly once per distinct reason. - get_logger().warning_once(f'QSA sparse selection disabled: {reason}') + def _qsa_select(self, hidden_states, attn_kwargs, position_ids=None): + """Choose the QSA selection representation for this forward. - def _qsa_select_mask(self, hidden_states, attn_kwargs): - # return None means full attention - # TODO: support padding_free & cp + Returns ``(selection, is_sparse)``. ``is_sparse`` means ``selection`` is the + int64 index tensor consumed by ``QSASparseCoreAttention`` (sbhd and thd, + with or without SP/CP); otherwise it is the bool TE mask from the legacy + path, or ``None`` for full attention. CP needs the allgather comm type + (the selection has to see every key before attention runs; ring/p2p + cannot provide that), mirroring mcore DSA's restriction. + """ indexer = getattr(self.self_attention, 'indexer', None) + sparse_ok = isinstance(getattr(self.self_attention, 'core_attention', None), QSASparseCoreAttention) if indexer is None: - return None - if attn_kwargs.get('packed_seq_params') is not None: - self._warn_qsa_fallback_once( - 'packing/padding_free is enabled (qkv_format=thd), which TE cannot combine with a ' - 'custom attention mask. QSA layers fall back to full attention -- training will ' - 'differ from sparse inference beyond the indexer budget. Pass `--padding_free false` ' - 'to enable QSA sparse selection.') - return None + return None, False + packed_seq_params: PackedSeqParams = attn_kwargs.get('packed_seq_params') + is_thd = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + cp_size = self.config.context_parallel_size + needs_kernel = is_thd or cp_size > 1 + + # sbhd with CP==1 + if not needs_kernel: + return self._qsa_select_mask(hidden_states, attn_kwargs), False + + # From here the mask path is not an option, so every failure raises instead of + # silently degrading to dense attention (which would diverge from the sparse + # rollout without telling anyone). + if not sparse_ok: + raise RuntimeError(f'QSA needs the sparse kernel here ({"packing/thd" if is_thd else f"CP={cp_size}"}), ' + 'but QSASparseCoreAttention was not installed -- triton is missing or ' + f'kv_channels={getattr(self.config, "kv_channels", None)} is not a power of two. ' + 'Use --padding_free false with context_parallel_size 1 to take the bool-mask path.') + if cp_size > 1 and getattr(self.config, 'cp_comm_type', None) != 'allgather': + raise RuntimeError(f"QSA sparse selection with context_parallel_size={cp_size} requires " + f"cp_comm_type='allgather' (got {getattr(self.config, 'cp_comm_type', None)!r}): the " + 'selection has to see every key before attention runs, which ring/p2p cannot provide.') + rotary_pos_emb = attn_kwargs.get('rotary_pos_emb') + if rotary_pos_emb is None: + raise RuntimeError('QSA sparse selection needs rotary_pos_emb (blocks rotate at their first ' + 'token position) but it was not passed to the layer.') + if is_thd: + indices = self._qsa_select_indices_thd(hidden_states, rotary_pos_emb, packed_seq_params, position_ids) + else: + indices = self._qsa_select_indices_sbhd(hidden_states, rotary_pos_emb) + return indices, True + + def _qsa_select_indices_sbhd(self, hidden_states, rotary_pos_emb): + if self.config.sequence_parallel and self.config.tensor_model_parallel_size > 1: + hidden_states = gather_from_sequence_parallel_region(hidden_states, tensor_parallel_output_grad=False) + if self.config.context_parallel_size > 1: + hidden_states = reconstruct_tensor_cp(hidden_states, None, dim=0) + rotary_pos_emb = reconstruct_tensor_cp(rotary_pos_emb, None, dim=0) + return self.self_attention.indexer.selection_as_token_indices(hidden_states, rotary_pos_emb) + + def _qsa_select_indices_thd(self, hidden_states, rotary_pos_emb, packed_seq_params, position_ids=None): + if self.config.sequence_parallel and self.config.tensor_model_parallel_size > 1: + hidden_states = gather_from_sequence_parallel_region(hidden_states, tensor_parallel_output_grad=False) + psp_for_cp = None if self.config.context_parallel_size > 1: - self._warn_qsa_fallback_once( - f'context_parallel_size={self.config.context_parallel_size} > 1 is not supported by ' - 'the QSA indexer yet (block pooling needs keys from other CP ranks). QSA layers fall ' - 'back to full attention -- training will differ from sparse inference beyond the ' - 'indexer budget.') + # TE's packed CP partition (thd_get_partitioned_indices) requires + # int32 cu; the training pipeline produces int32, but normalize callers + # that hand us int64. + if packed_seq_params.cu_seqlens_q is not None \ + and packed_seq_params.cu_seqlens_q.dtype != torch.int32: + psp_for_cp = copy.copy(packed_seq_params) + psp_for_cp.cu_seqlens_q = packed_seq_params.cu_seqlens_q.to(torch.int32) + if packed_seq_params.cu_seqlens_q_padded is not None: + psp_for_cp.cu_seqlens_q_padded = packed_seq_params.cu_seqlens_q_padded.to(torch.int32) + else: + psp_for_cp = packed_seq_params + hidden_states = reconstruct_tensor_cp(hidden_states, psp_for_cp, dim=0) + # Per-token rotary angles. Without rope fusion gpt_model already indexes + # the freq table by position_ids, so what arrives is per-token (zigzag- + # sharded under CP -- undo it like hidden). With fusion the raw table + # arrives and must be indexed by the (CP-reconstructed) per-doc ids. + freqs = rotary_pos_emb + fused_table = freqs.shape[0] != hidden_states.shape[0] + if self.config.context_parallel_size > 1: + if fused_table: + if position_ids is None: + raise RuntimeError('QSA thd selection under CP needs position_ids to index the fused rotary ' + 'table (apply_rope_fusion=true hands over the raw table, not per-token ' + 'freqs). Pass position_ids, or set --apply_rope_fusion false.') + pos = reconstruct_tensor_cp(position_ids, psp_for_cp, dim=1) + freqs = freqs[pos.reshape(-1)] + else: + freqs = reconstruct_tensor_cp(freqs, psp_for_cp, dim=0) + elif fused_table: + # Same problem without CP, and here there is no reconstruct step to hide + # behind: the indexer would slice the raw table's first T rows, treating + # row i as token i's angle. In a packed batch token i sits at in-document + # position i - cu[doc], so those angles belong to the wrong positions -- + # silently degrading the selection instead of failing. + raise RuntimeError(f'QSA thd selection got a fused rotary table ({freqs.shape[0]} rows for ' + f'{hidden_states.shape[0]} tokens): apply_rope_fusion=true hands over the raw ' + 'table rather than per-token freqs. Set --apply_rope_fusion false.') + # the CP reconstruct (like TE's thd kernels) works in the padded pack + # space, so align against the padded cu when present + cu = packed_seq_params.cu_seqlens_q_padded + if cu is None: + cu = packed_seq_params.cu_seqlens_q + if cu is None: + raise RuntimeError('QSA thd selection needs packed_seq_params.cu_seqlens_q to find document ' + 'boundaries, but it is missing.') + cu = Qwen4ExpTextPLELayer._normalize_cu_seqlens(cu, hidden_states.shape[0]) + hidden_tok = hidden_states.reshape(hidden_states.shape[0], -1) + return self.self_attention.indexer.select_token_indices_thd(hidden_tok, freqs, cu) + + def _qsa_select_mask(self, hidden_states, attn_kwargs): + # Bool-mask QSA on TE's `arbitrary` mask. Only reached for sbhd with CP==1 -- + # _qsa_selection() routes thd and CP>1 to the kernel, because TE rejects an + # arbitrary mask under thd and this path never gathers keys across CP ranks. + # Returning None means full attention, which here only happens when the + # sequence is short enough that selection is a no-op anyway (selection_as_mask + # short-circuits at max_blocks <= block_topk). + indexer = getattr(self.self_attention, 'indexer', None) + if indexer is None: return None rotary_pos_emb = attn_kwargs.get('rotary_pos_emb') if rotary_pos_emb is None: - return None - return indexer.select_mask(hidden_states, rotary_pos_emb) + raise RuntimeError('QSA bool-mask selection needs rotary_pos_emb (blocks rotate at their first ' + 'token position) but it was not passed to the layer.') + if self.config.sequence_parallel and self.config.tensor_model_parallel_size > 1: + hidden_states = gather_from_sequence_parallel_region(hidden_states, tensor_parallel_output_grad=False) + return indexer.selection_as_mask(hidden_states, rotary_pos_emb) class Qwen4ExpTransformerBlock(TransformerBlock): @@ -162,8 +266,9 @@ class Qwen4ExpTransformerBlock(TransformerBlock): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) config = self.config - hc_count = getattr(config, 'hc_count', 0) or 0 - if hc_count > 1 and self.has_final_layernorm_in_this_stage(): + if config.hc_count is None: + raise ValueError('Qwen4Exp requires config.hc_count (checkpoint has hc_count=4).') + if config.hc_count > 1 and self.has_final_layernorm_in_this_stage(): # Final contraction (use_combine=False matches the checkpoint: # hyper_connection_mixer has no block_inject_weight). self.hyper_connection_mixer = Qwen4ExpTextGatedResidual(config, use_combine=False) @@ -221,14 +326,7 @@ def _set_layer_hc(self, mg_layer, hf_state_dict, to_mcore: bool): self._set_state_dict(hyper_connection, weight_key, hf_state_dict, f'{key}.{weight_key}', to_mcore) # --- PLE ----------------------------------------------------------------- - # (mcore attribute name, hf checkpoint suffix). Both sides now use the - # transformers buffer names; the table is kept so these buffers stay on - # the dedicated PLE conversion path (pp broadcast, no TP split). - _PLE_NGRAM_BUFFERS = ( - ('layer_multipliers', 'layer_multipliers'), - ('ngram_heads_offsets', 'ngram_heads_offsets'), - ('ngram_heads_vocab_sizes', 'ngram_heads_vocab_sizes'), - ) + _PLE_NGRAM_BUFFERS = ('layer_multipliers', 'ngram_heads_offsets', 'ngram_heads_vocab_sizes') def _get_tp_split_dim(self, mg_key): # PLE weights are replicated across TP; in particular `conv1d.weight` @@ -244,65 +342,6 @@ def _get_pp_src_rank(self, has_module: bool) -> int: dist.all_reduce(holder, op=dist.ReduceOp.MAX, group=self.pp_group) return int(holder.item()) - def _set_ple_ngram_embedding(self, ple, hf_state_dict, to_mcore: bool, pp_src_rank: int): - # The checkpoint shards the (padded) table into `parts` uniform row - # blocks, so shard boundaries must be derived from the padded size. - total = parts = dim = None - if ple is not None: - total = ple.ple_embedding.ngram_embedding.num_embeddings - parts = ple.ple_embedding.split_ngram_parts - dim = ple.ple_embedding.head_dim - if not to_mcore and self.pp_size > 1: - obj = [(total, parts, dim)] - dist.broadcast_object_list(obj, src=pp_src_rank, group=self.pp_group) - total, parts, dim = obj[0] - if to_mcore and ple is None: - return - shard_size = (total + parts - 1) // parts - tp_size = mpu.get_tensor_model_parallel_world_size() - tp_rank = mpu.get_tensor_model_parallel_rank() - tp_ranks = dist.get_process_group_ranks(self.tp_group) - emb = ple.ple_embedding.ngram_embedding if ple is not None else None - dtype = emb.weight.dtype if emb is not None else self.config.params_dtype - device = emb.weight.device if emb is not None else torch.cuda.current_device() - per_partition = (emb.num_embeddings_per_partition if emb is not None else (total + tp_size - 1) // tp_size) - tp_start = tp_rank * per_partition if emb is not None else 0 - tp_end = min((tp_rank + 1) * per_partition, total) if emb is not None else 0 - if to_mcore: - for i in range(parts): - key = f'ple.ple_embedding.ngram_embedding.shard_{i}.weight' - if key not in hf_state_dict: - continue - cs, ce = i * shard_size, min((i + 1) * shard_size, total) - s, e = max(cs, tp_start), min(ce, tp_end) - if s < e: - weight = hf_state_dict[key].load() - emb.weight.data[s - tp_start:e - tp_start] = weight[s - cs:e - cs].to(emb.weight.dtype) - else: - for i in range(parts): - cs, ce = i * shard_size, min((i + 1) * shard_size, total) - pieces = [] - for r in range(tp_size): - r_start = r * per_partition - r_end = min((r + 1) * per_partition, total) - s, e = max(cs, r_start), min(ce, r_end) - if s >= e: - continue - if emb is not None and r == tp_rank: - piece = emb.weight.data[s - tp_start:e - tp_start].clone() - else: - piece = torch.empty(e - s, dim, dtype=dtype, device=device) - dist.broadcast(piece, src=tp_ranks[r], group=self.tp_group) - pieces.append(piece) - shard = torch.cat(pieces, dim=0) - if self.pp_size > 1: - dist.broadcast(shard, src=pp_src_rank, group=self.pp_group) - # Written directly into the state dict (bypasses _get_weight, - # which normally applies _target_device). - if self._target_device is not None: - shard = shard.to(self._target_device) - hf_state_dict[f'ple.ple_embedding.ngram_embedding.shard_{i}.weight'] = shard - def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool): ple = None if mg_layer is None else getattr(mg_layer, 'ple', None) if to_mcore: @@ -318,12 +357,18 @@ def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool): has_ple = self._reduce_tensor_pp_group(ple is not None, to_mcore) if not has_ple: return - for mg_buf, hf_buf in self._PLE_NGRAM_BUFFERS: + # `ple` is only non-None on the pp stage owning the PLE layer, so the offload + # flag has to be reduced across pp before it can gate the loop below -- that + # loop runs pp collectives (broadcast_object_list) and export_table_to_hf runs + # tp ones, and stages disagreeing on whether to enter would deadlock. + ple_offloaded = self._reduce_tensor_pp_group(ple is not None and ple.ple_embedding.cpu_offload, to_mcore) + skip_ngram_state = not to_mcore and not self._is_saving and (self._peft_format or ple_offloaded) + for buf in () if skip_ngram_state else self._PLE_NGRAM_BUFFERS: if to_mcore: - buffer = getattr(ple.ple_embedding, mg_buf) - buffer.copy_(hf_state_dict[f'ple.ple_embedding.{hf_buf}'].load().to(buffer.device)) + buffer = getattr(ple.ple_embedding, buf) + buffer.copy_(hf_state_dict[f'ple.ple_embedding.{buf}'].load().to(buffer.device)) else: - tensor = getattr(ple.ple_embedding, mg_buf).data.clone() if ple is not None else None + tensor = getattr(ple.ple_embedding, buf).data.clone() if ple is not None else None if self.pp_size > 1: obj = [tensor] dist.broadcast_object_list(obj, src=pp_src_rank, group=self.pp_group) @@ -332,8 +377,12 @@ def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool): # which normally applies _target_device). if tensor is not None and self._target_device is not None: tensor = tensor.to(self._target_device) - hf_state_dict[f'ple.ple_embedding.{hf_buf}'] = tensor - self._set_ple_ngram_embedding(ple, hf_state_dict, to_mcore, pp_src_rank) + hf_state_dict[f'ple.ple_embedding.{buf}'] = tensor + if not skip_ngram_state and to_mcore: + # The table's only ingestion path: fill from the HF checkpoint shards. + ple.ple_embedding.fill_table_from_hf(hf_state_dict) + elif not skip_ngram_state and ple is not None: + ple.ple_embedding.export_table_to_hf(hf_state_dict) self._converting_ple = True try: for mg_key, hf_key in [('key_proj.weight', 'ple.key_proj.weight'), @@ -371,8 +420,8 @@ def _set_final_layernorm(self, lm_model, hf_state_dict, to_mcore): def _convert_post_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcore): res = super()._convert_post_process(mg_model, hf_state_dict, hf_prefix, to_mcore) - lm_model = getattr(mg_model, 'language_model') if self.is_multimodal else mg_model - hc_count = getattr(self.config, 'hc_count', 0) or 0 + lm_model = mg_model.language_model if self.is_multimodal else mg_model + hc_count = self.config.hc_count if hc_count > 1: # The mixer only exists on the stage holding the final layernorm. mixer_keys = ['hc_norm.weight', 'input_mix_weight_down.weight', 'input_mix_weight_up.weight'] diff --git a/src/mcore_bridge/model/modules/__init__.py b/src/mcore_bridge/model/modules/__init__.py index 7f5f2a2..fdba46b 100644 --- a/src/mcore_bridge/model/modules/__init__.py +++ b/src/mcore_bridge/model/modules/__init__.py @@ -5,6 +5,7 @@ from .gated_delta_net import GatedDeltaNet from .gated_self_attention import GatedSelfAttention from .hyper_connection_gated import Qwen4ExpTextGatedResidual, Qwen4ExpTextGroupedRMSNorm +from .kernels import QSASparseCoreAttention, qsa_sparse_supported from .mtp_layer import MultiTokenPredictionLayer from .multi_latent_attention import MLASelfAttention from .ple import Qwen4ExpTextNGramEmbedding, Qwen4ExpTextPLELayer diff --git a/src/mcore_bridge/model/modules/hyper_connection_gated.py b/src/mcore_bridge/model/modules/hyper_connection_gated.py index 78d373e..5c94b58 100644 --- a/src/mcore_bridge/model/modules/hyper_connection_gated.py +++ b/src/mcore_bridge/model/modules/hyper_connection_gated.py @@ -1,22 +1,39 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import torch import torch.nn.functional as F +from megatron.core.extensions.transformer_engine import TELinear from torch import nn from typing import Tuple +def _duplicated_linear(config, input_size: int, output_size: int) -> TELinear: + """Replicated projection: TELinear so LoRA can wrap it. + + ``dispatch_megatron`` only wraps TE classes into ``LoraParallelLinear``, which the + bridge's PEFT export keys off. ``parallel_mode='duplicated'`` because + ``block_inject_weight``'s output is hc_count (4) and cannot be split across TP=8. + """ + return TELinear( + input_size=input_size, + output_size=output_size, + parallel_mode='duplicated', + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=True, + skip_weight_param_allocation=False, + ) + + class Qwen4ExpTextGroupedRMSNorm(nn.Module): def __init__(self, dim: int, group_size: int, eps: float = 1e-6, dtype=None, sequence_parallel: bool = False): super().__init__() self.eps = eps - self.weight = nn.Parameter(torch.zeros(dim)) + self.weight = nn.Parameter(torch.zeros(dim, dtype=dtype)) self.group_size = group_size if dim % group_size != 0: raise ValueError(f'hidden_size ({dim}) must be divisible by group_size ({group_size}).') - # mcore-specific: params_dtype weight + SP grad flag (HF builds fp32). - if dtype is not None: - self.weight.data = self.weight.data.to(dtype) setattr(self.weight, 'sequence_parallel', sequence_parallel) def _norm(self, x: torch.Tensor) -> torch.Tensor: @@ -33,6 +50,20 @@ def extra_repr(self): return f'{tuple(self.weight.shape)}, eps={self.eps}' +@torch.compile +def _mix_elementwise(down_out: torch.Tensor, hc_count: int) -> torch.Tensor: + """silu(down/hc) -- the pre-up-projection half of the gate chain.""" + return F.silu(down_out / hc_count) + + +@torch.compile +def _mix_and_reduce(up_out: torch.Tensor, hyper_input_normed: torch.Tensor, hc_count: int, + hidden_size: int) -> torch.Tensor: + """sigmoid -> unflatten -> multiply -> mean, the post-up-projection half.""" + w = torch.sigmoid(up_out).unflatten(-1, (hc_count, hidden_size)) + return (w * hyper_input_normed.unflatten(-1, (hc_count, hidden_size))).mean(dim=-2) + + class Qwen4ExpTextGatedResidual(nn.Module): def __init__(self, config, use_combine: bool = True): @@ -44,10 +75,9 @@ def __init__(self, config, use_combine: bool = True): # eps=config.rms_norm_eps); layernorm_epsilon is mcore's rms_norm_eps.) self.hc_norm = Qwen4ExpTextGroupedRMSNorm( hc_hidden_size, group_size=self.hidden_size, eps=config.layernorm_epsilon, dtype=config.params_dtype) - self.input_mix_weight_down = nn.Linear(hc_hidden_size, config.hc_lowrank, bias=False, dtype=config.params_dtype) - self.input_mix_weight_up = nn.Linear(config.hc_lowrank, hc_hidden_size, bias=False, dtype=config.params_dtype) - self.block_inject_weight = nn.Linear( - hc_hidden_size, self.hc_count, bias=False, dtype=config.params_dtype) if use_combine else None + self.input_mix_weight_down = _duplicated_linear(config, hc_hidden_size, config.hc_lowrank) + self.input_mix_weight_up = _duplicated_linear(config, config.hc_lowrank, hc_hidden_size) + self.block_inject_weight = _duplicated_linear(config, hc_hidden_size, self.hc_count) if use_combine else None # mcore-specific: SP grad flag on the replicated weights. for param in self.parameters(): setattr(param, 'sequence_parallel', config.sequence_parallel) @@ -58,12 +88,15 @@ def forward(self, hyper_input: torch.Tensor) -> Tuple[torch.Tensor, ...]: raise ValueError(f'Expected {self.hc_count * self.hidden_size} hyper-connection features, ' f'got {hyper_input.shape[-1]}.') hyper_input_normed = self.hc_norm(hyper_input) - input_mix_weight = F.silu(self.input_mix_weight_down(hyper_input_normed) / self.hc_count) - input_mix_weight = torch.sigmoid(self.input_mix_weight_up(input_mix_weight)) - input_mix_weight = input_mix_weight.unflatten(-1, (self.hc_count, self.hidden_size)) - mixed_input = (input_mix_weight * hyper_input_normed.unflatten(-1, - (self.hc_count, self.hidden_size))).mean(dim=-2) + # TELinear returns (output, bias); bias is None here (bias=False). The two + # linears stay outside the compiled helpers: TE modules do work inside a + # compiled region (verified), but keeping them out leaves TE's own fused + # kernels and FP8 bookkeeping untouched and limits the graph to the + # elementwise ops that actually benefit. + input_mix_weight = _mix_elementwise(self.input_mix_weight_down(hyper_input_normed)[0], self.hc_count) + input_mix_weight = self.input_mix_weight_up(input_mix_weight)[0] + mixed_input = _mix_and_reduce(input_mix_weight, hyper_input_normed, self.hc_count, self.hidden_size) if self.block_inject_weight is None: return mixed_input - injection_weights = 2 * torch.sigmoid(self.block_inject_weight(hyper_input_normed) / self.hc_count) + injection_weights = 2 * torch.sigmoid(self.block_inject_weight(hyper_input_normed)[0] / self.hc_count) return mixed_input, hyper_input, injection_weights diff --git a/src/mcore_bridge/model/modules/kernels/__init__.py b/src/mcore_bridge/model/modules/kernels/__init__.py new file mode 100644 index 0000000..2139173 --- /dev/null +++ b/src/mcore_bridge/model/modules/kernels/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .ple_kernels import gather_ple_rows, ple_gate_conv_triton +from .qsa_kernels import QSASparseCoreAttention, qsa_sparse_supported + +__all__ = [ + 'QSASparseCoreAttention', + 'gather_ple_rows', + 'ple_gate_conv_triton', + 'qsa_sparse_supported', +] diff --git a/src/mcore_bridge/model/modules/kernels/ple_kernels.py b/src/mcore_bridge/model/modules/kernels/ple_kernels.py new file mode 100644 index 0000000..ef1f0c6 --- /dev/null +++ b/src/mcore_bridge/model/modules/kernels/ple_kernels.py @@ -0,0 +1,532 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Triton kernels for the qwen4_exp PLE host-offload path. + +``gather_ple_rows`` reads rows straight out of the CPU-pinned n-gram table via +a raw host pointer, so the GPU pulls only the rows it needs over the coherent +link instead of staging the (~102 GB) table into HBM. Mirrors sglang's +``_gather_ple_embedding_from_pinned_kernel``. + +Every kernel here is optional: if triton or CUDA is unavailable the callers in +ple.py fall back to the plain-torch path, which is numerically identical. +""" +import torch + +try: + import triton + import triton.language as tl + HAVE_TRITON = True +except Exception: # pragma: no cover - triton absent + HAVE_TRITON = False + +if HAVE_TRITON: + + @triton.jit + def _gather_ple_rows_from_pinned( + weight_ptr, + ids_ptr, + output_ptr, + embedding_dim, + row_start, + row_end, + BLOCK_D: tl.constexpr, + ): + # One program per flattened (token, hash-head) id. Rows outside this TP + # rank's [row_start, row_end) are written as zero; the caller sums the + # per-rank results across TP to reassemble the full embedding. + row_id = tl.program_id(0) + global_idx = tl.load(ids_ptr + row_id) + in_range = (global_idx >= row_start) & (global_idx < row_end) + local_idx = tl.where(in_range, global_idx - row_start, 0) + offsets = tl.arange(0, BLOCK_D) + mask = offsets < embedding_dim + ptr = weight_ptr.to(tl.int64).to(tl.pointer_type(tl.bfloat16)) + values = tl.load(ptr + local_idx * embedding_dim + offsets, mask=mask, other=0.0) + tl.store( + output_ptr + row_id * embedding_dim + offsets, + tl.where(in_range, values.to(tl.bfloat16), 0.0), + mask=mask, + ) + + +def gather_ple_rows(host_table, ids, row_start, row_end, out=None): + """Gather n-gram rows from the CPU-pinned table with a triton kernel. + + Returns ``None`` when the fast path is not usable (no triton, CPU-resident + ids, or the table is not bf16), so the caller can fall back to the torch path. + + Args: + host_table: ``[n_local, embedding_dim]`` bf16 CPU-pinned table partition. + ids: int64 tensor of any shape, values in ``[0, padded_vocab_size)``. + row_start / row_end: this rank's global row range. + out: optional preallocated ``[*ids.shape, embedding_dim]`` bf16 device tensor. + """ + # Gate on "triton exists + ids live on an accelerator", not on is_cuda: + # triton ships per-vendor backends (ROCm in-tree; NPU/XPU via vendor forks + # such as triton-ascend), so a platform that cannot run this kernel fails + # at launch instead of being silently excluded here. The device-side read + # of the pinned host table is only verified on CUDA so far. + if not HAVE_TRITON or ids.device.type == 'cpu': + return None + if host_table.dtype != torch.bfloat16: + return None + embedding_dim = host_table.shape[-1] + + shape = (*ids.shape, embedding_dim) + if out is None: + out = torch.empty(shape, dtype=torch.bfloat16, device=ids.device) + flat = ids.reshape(-1) + if flat.numel(): + _gather_ple_rows_from_pinned[(flat.numel(), )]( + host_table.data_ptr(), + flat.contiguous(), + out.view(-1, embedding_dim), + embedding_dim=embedding_dim, + row_start=row_start, + row_end=row_end, + BLOCK_D=triton.next_power_of_2(embedding_dim), + ) + return out + + +if HAVE_TRITON: + + @triton.jit(do_not_specialize=['T']) + def _ple_gate_fwd_kernel( + key_ptr, # [T, N*C] pre-norm key projection + query_ptr, # [T, N*C] hc state (the PLE query) + value_ptr, # [T, C] + wk_ptr, + wq_ptr, # zero-centered grouped-norm weights [N*C] + gated_ptr, # fp32 out [T, N*C] + gate_ptr, + rstdk_ptr, + rstdq_ptr, # fp32 out [T, N] + T, + N: tl.constexpr, + C: tl.constexpr, + EPS: tl.constexpr, + SQRTC: tl.constexpr, + BLOCK_C: tl.constexpr, + ): + # Fused: grouped RMSNorm(key) * grouped RMSNorm(query) -> per-group score, + # gate = sigmoid(sign(s)*sqrt(max(|s|,1e-6))), out = gate * value. + pid = tl.program_id(0) + t = pid // N + c = pid % N + if t >= T: + return + offs = tl.arange(0, BLOCK_C) + mask = offs < C + base = t * (N * C) + c * C + + k = tl.load(key_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + q = tl.load(query_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + wk = tl.load(wk_ptr + c * C + offs, mask=mask, other=0.0).to(tl.float32) + wq = tl.load(wq_ptr + c * C + offs, mask=mask, other=0.0).to(tl.float32) + + rk = 1.0 / tl.sqrt(tl.sum(k * k, axis=0) / C + EPS) + rq = 1.0 / tl.sqrt(tl.sum(q * q, axis=0) / C + EPS) + kn = k * rk * (1.0 + wk) + qn = q * rq * (1.0 + wq) + + score = tl.sum(kn * qn, axis=0) / SQRTC + mag = tl.maximum(tl.abs(score), 1e-6) + sgn = tl.where(score >= 0, 1.0, -1.0) + gate = tl.sigmoid(sgn * tl.sqrt(mag)) + + v = tl.load(value_ptr + t * C + offs, mask=mask, other=0.0).to(tl.float32) + tl.store(gated_ptr + base + offs, gate * v, mask=mask) + tl.store(gate_ptr + t * N + c, gate) + tl.store(rstdk_ptr + t * N + c, rk) + tl.store(rstdq_ptr + t * N + c, rq) + + @triton.jit(do_not_specialize=['T']) + def _ple_gate_bwd_kernel( + dgated_ptr, # fp32 in [T, N*C] + key_ptr, + query_ptr, + value_ptr, + wk_ptr, + wq_ptr, + gate_ptr, + rstdk_ptr, + rstdq_ptr, + dkey_ptr, + dquery_ptr, # out, input dtype [T, N*C] + dvalue_ptr, # fp32 out [T, N, C] (host sums over N) + dwk_partial_ptr, + dwq_partial_ptr, # fp32 out [T, N*C] (host sums over T) + T, + N: tl.constexpr, + C: tl.constexpr, + SQRTC: tl.constexpr, + BLOCK_C: tl.constexpr, + ): + pid = tl.program_id(0) + t = pid // N + c = pid % N + if t >= T: + return + offs = tl.arange(0, BLOCK_C) + mask = offs < C + base = t * (N * C) + c * C + sqrtC = SQRTC + + k = tl.load(key_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + q = tl.load(query_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + v = tl.load(value_ptr + t * C + offs, mask=mask, other=0.0).to(tl.float32) + wk = tl.load(wk_ptr + c * C + offs, mask=mask, other=0.0).to(tl.float32) + wq = tl.load(wq_ptr + c * C + offs, mask=mask, other=0.0).to(tl.float32) + g = tl.load(gate_ptr + t * N + c) + rk = tl.load(rstdk_ptr + t * N + c) + rq = tl.load(rstdq_ptr + t * N + c) + dg_out = tl.load(dgated_ptr + base + offs, mask=mask, other=0.0) + + kn = k * rk * (1.0 + wk) + qn = q * rq * (1.0 + wq) + + dgate = tl.sum(dg_out * v, axis=0) + tl.store(dvalue_ptr + (t * N + c) * C + offs, dg_out * g, mask=mask) + + score = tl.sum(kn * qn, axis=0) / sqrtC + mag = tl.maximum(tl.abs(score), 1e-6) + du = dgate * g * (1.0 - g) + ds = tl.where(tl.abs(score) > 1e-6, du / (2.0 * tl.sqrt(mag)), 0.0) + + dkn = qn * (ds / sqrtC) + dqn = kn * (ds / sqrtC) + + tl.store(dwk_partial_ptr + base + offs, dkn * (k * rk), mask=mask) + tl.store(dwq_partial_ptr + base + offs, dqn * (q * rq), mask=mask) + + gk = dkn * (1.0 + wk) + dotk = tl.sum(gk * k, axis=0) + dk = rk * gk - k * (rk * rk * rk) * (dotk / C) + gq = dqn * (1.0 + wq) + dotq = tl.sum(gq * q, axis=0) + dq = rq * gq - q * (rq * rq * rq) * (dotq / C) + + tl.store(dkey_ptr + base + offs, dk.to(dkey_ptr.dtype.element_ty), mask=mask) + tl.store(dquery_ptr + base + offs, dq.to(dquery_ptr.dtype.element_ty), mask=mask) + + @triton.jit(do_not_specialize=['T']) + def _ple_norm_fwd_kernel( + x_ptr, + w_ptr, + out_ptr, + rstd_ptr, + T, + N: tl.constexpr, + C: tl.constexpr, + EPS: tl.constexpr, + BLOCK_C: tl.constexpr, + ): + # Grouped zero-centered RMSNorm: out = x * rstd * (1 + w), fp32 out. + pid = tl.program_id(0) + t = pid // N + c = pid % N + if t >= T: + return + offs = tl.arange(0, BLOCK_C) + mask = offs < C + base = t * (N * C) + c * C + x = tl.load(x_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + w = tl.load(w_ptr + c * C + offs, mask=mask, other=0.0).to(tl.float32) + r = 1.0 / tl.sqrt(tl.sum(x * x, axis=0) / C + EPS) + tl.store(out_ptr + base + offs, x * r * (1.0 + w), mask=mask) + tl.store(rstd_ptr + t * N + c, r) + + @triton.jit(do_not_specialize=['T']) + def _ple_norm_bwd_kernel( + x_ptr, + w_ptr, + rstd_ptr, + dout_ptr, + dx_ptr, + T, + N: tl.constexpr, + C: tl.constexpr, + BLOCK_C: tl.constexpr, + ): + pid = tl.program_id(0) + t = pid // N + c = pid % N + if t >= T: + return + offs = tl.arange(0, BLOCK_C) + mask = offs < C + base = t * (N * C) + c * C + x = tl.load(x_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + w = tl.load(w_ptr + c * C + offs, mask=mask, other=0.0).to(tl.float32) + r = tl.load(rstd_ptr + t * N + c) + dn = tl.load(dout_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + xh = x * r + g = dn * (1.0 + w) + dx = r * (g - xh * (tl.sum(g * xh, axis=0) / C)) + tl.store(dx_ptr + base + offs, dx, mask=mask) + + @triton.jit(do_not_specialize=['T', 'W']) + def _ple_conv_fwd_kernel( + normed_ptr, # fp32 [T, W] + gated_ptr, # fp32 [T, W] + convw_ptr, # [W, K] depthwise weights + segstart_ptr, # int32 [T] + out_ptr, # out dtype [T, W] + conv_ptr, # fp32 out [T, W] pre-SiLU conv (saved for bwd) + T, + W, + K: tl.constexpr, + DIL: tl.constexpr, + BLOCK_W: tl.constexpr, + ): + # Causal dilated depthwise conv; rows never read across their segment + # start. out = gated + silu(conv(normed)). + t = tl.program_id(0) + wb = tl.program_id(1) + if t >= T: + return + offs = wb * BLOCK_W + tl.arange(0, BLOCK_W) + mask = offs < W + seg_lo = tl.load(segstart_ptr + t) + + acc = tl.zeros([BLOCK_W], dtype=tl.float32) + for j in tl.static_range(K): + src = t - (K - 1 - j) * DIL + wgt = tl.load(convw_ptr + offs * K + j, mask=mask, other=0.0).to(tl.float32) + if src >= 0: + ok = src >= seg_lo + x = tl.load(normed_ptr + src * W + offs, mask=mask & ok, other=0.0) + acc += wgt * x + tl.store(conv_ptr + t * W + offs, acc, mask=mask) + silu = acc * tl.sigmoid(acc) + gt = tl.load(gated_ptr + t * W + offs, mask=mask, other=0.0) + tl.store(out_ptr + t * W + offs, (gt + silu).to(out_ptr.dtype.element_ty), mask=mask) + + @triton.jit(do_not_specialize=['T', 'W']) + def _ple_conv_bwd_kernel( + dout_ptr, # incoming grad [T, W] + conv_ptr, # fp32 [T, W] pre-SiLU + normed_ptr, # fp32 [T, W] + convw_ptr, # [W, K] + segstart_ptr, + segend_ptr, # int32 [T] (exclusive) + dnormed_ptr, # fp32 out [T, W] + dconvw_ptr, # fp32 out [W, K] via atomics + dgated_add_ptr, # fp32 out [T, W] (dout passthrough for the residual) + T, + W, + K: tl.constexpr, + DIL: tl.constexpr, + BLOCK_W: tl.constexpr, + ): + t = tl.program_id(0) + wb = tl.program_id(1) + if t >= T: + return + offs = wb * BLOCK_W + tl.arange(0, BLOCK_W) + mask = offs < W + seg_lo = tl.load(segstart_ptr + t) + seg_hi = tl.load(segend_ptr + t) + + do = tl.load(dout_ptr + t * W + offs, mask=mask, other=0.0).to(tl.float32) + tl.store(dgated_add_ptr + t * W + offs, do, mask=mask) + + cv = tl.load(conv_ptr + t * W + offs, mask=mask, other=0.0) + sig = tl.sigmoid(cv) + dconv = do * sig * (1.0 + cv * (1.0 - sig)) + + for j in tl.static_range(K): + src = t - (K - 1 - j) * DIL + if src >= 0: + ok = src >= seg_lo + x = tl.load(normed_ptr + src * W + offs, mask=mask & ok, other=0.0) + tl.atomic_add(dconvw_ptr + offs * K + j, dconv * x, mask=mask & ok) + + acc = tl.zeros([BLOCK_W], dtype=tl.float32) + for j in tl.static_range(K): + dst = t + (K - 1 - j) * DIL + wgt = tl.load(convw_ptr + offs * K + j, mask=mask, other=0.0).to(tl.float32) + if dst < T: + ok = dst < seg_hi + do2 = tl.load(dout_ptr + dst * W + offs, mask=mask & ok, other=0.0).to(tl.float32) + cv2 = tl.load(conv_ptr + dst * W + offs, mask=mask & ok, other=0.0) + sig2 = tl.sigmoid(cv2) + acc += wgt * do2 * sig2 * (1.0 + cv2 * (1.0 - sig2)) + tl.store(dnormed_ptr + t * W + offs, acc, mask=mask) + + +if HAVE_TRITON: + import math as _math + + def _uniform_seg_bounds(total, seq_len, device): + # The PLE forward runs on padded [rows, seq_len] batches, so every row + # is its own segment: the dilated conv must not read across rows. + pos = torch.arange(total, device=device) + lo = (pos // seq_len * seq_len).to(torch.int32) + hi = (lo + seq_len).to(torch.int32) + return lo, hi + + class _PLEGateConv(torch.autograd.Function): + + @staticmethod + def forward(ctx, hc_state, key, value, wk, wq, wc, conv_w, n, eps, dilation, seq_len): + T, W = hc_state.shape + C = W // n + Kk = conv_w.shape[-1] + dev = hc_state.device + block_c = triton.next_power_of_2(C) + + gated = torch.empty(T, W, dtype=torch.float32, device=dev) + gate = torch.empty(T, n, dtype=torch.float32, device=dev) + rstdk = torch.empty(T, n, dtype=torch.float32, device=dev) + rstdq = torch.empty(T, n, dtype=torch.float32, device=dev) + if T > 0: + _ple_gate_fwd_kernel[(T * n, )]( + key, + hc_state, + value, + wk, + wq, + gated, + gate, + rstdk, + rstdq, + T, + N=n, + C=C, + EPS=eps, + SQRTC=_math.sqrt(C), + BLOCK_C=block_c) + + normed = torch.empty(T, W, dtype=torch.float32, device=dev) + rstdc = torch.empty(T, n, dtype=torch.float32, device=dev) + if T > 0: + _ple_norm_fwd_kernel[(T * n, )](gated, wc, normed, rstdc, T, N=n, C=C, EPS=eps, BLOCK_C=block_c) + + seg_lo, seg_hi = _uniform_seg_bounds(T, seq_len, dev) + convw2d = conv_w.reshape(W, Kk).contiguous() + out = torch.empty(T, W, dtype=hc_state.dtype, device=dev) + conv_pre = torch.empty(T, W, dtype=torch.float32, device=dev) + BW = 256 + if T > 0: + _ple_conv_fwd_kernel[(T, triton.cdiv(W, BW))]( + normed, gated, convw2d, seg_lo, out, conv_pre, T, W, K=Kk, DIL=dilation, BLOCK_W=BW) + + ctx.save_for_backward(hc_state, key, value, wk, wq, wc, convw2d, gate, rstdk, rstdq, rstdc, seg_lo, seg_hi, + conv_pre) + ctx.dims = (n, eps, dilation, Kk, conv_w.dtype) + return out + + @staticmethod + def backward(ctx, dout): + (hc_state, key, value, wk, wq, wc, convw2d, gate, rstdk, rstdq, rstdc, seg_lo, seg_hi, + conv_pre) = ctx.saved_tensors + n, eps, dilation, Kk, conv_w_dtype = ctx.dims + T, W = hc_state.shape + C = W // n + dev = hc_state.device + dout = dout.contiguous() + block_c = triton.next_power_of_2(C) + + # Recompute the forward intermediates (the fp32 gated/normed are not + # saved to keep the recompute footprint small). + gated = torch.empty(T, W, dtype=torch.float32, device=dev) + _g = torch.empty(T, n, dtype=torch.float32, device=dev) + _rk = torch.empty(T, n, dtype=torch.float32, device=dev) + _rq = torch.empty(T, n, dtype=torch.float32, device=dev) + if T > 0: + _ple_gate_fwd_kernel[(T * n, )]( + key, + hc_state, + value, + wk, + wq, + gated, + _g, + _rk, + _rq, + T, + N=n, + C=C, + EPS=eps, + SQRTC=_math.sqrt(C), + BLOCK_C=block_c) + normed = torch.empty(T, W, dtype=torch.float32, device=dev) + _rc = torch.empty(T, n, dtype=torch.float32, device=dev) + if T > 0: + _ple_norm_fwd_kernel[(T * n, )](gated, wc, normed, _rc, T, N=n, C=C, EPS=eps, BLOCK_C=block_c) + + dnormed = torch.empty(T, W, dtype=torch.float32, device=dev) + dconvw = torch.zeros(W, Kk, dtype=torch.float32, device=dev) + dgated = torch.empty(T, W, dtype=torch.float32, device=dev) + BW = 256 + if T > 0: + _ple_conv_bwd_kernel[(T, triton.cdiv(W, BW))]( + dout, + conv_pre, + normed, + convw2d, + seg_lo, + seg_hi, + dnormed, + dconvw, + dgated, + T, + W, + K=Kk, + DIL=dilation, + BLOCK_W=BW) + + # norm_conv backward: dwc on host, dx via kernel (fp32). + x_hat = (gated.view(T, n, C) * rstdc.unsqueeze(-1)).view(T, W) + dwc = (dnormed * x_hat).sum(dim=0).to(wc.dtype) + dgated_norm = torch.empty(T, W, dtype=torch.float32, device=dev) + if T > 0: + _ple_norm_bwd_kernel[(T * n, )](gated, wc, rstdc, dnormed, dgated_norm, T, N=n, C=C, BLOCK_C=block_c) + dgated += dgated_norm + + dkey = torch.empty_like(key) + dquery = torch.empty_like(hc_state) + dvalue_pern = torch.empty(T, n, C, dtype=torch.float32, device=dev) + dwk_part = torch.empty(T, W, dtype=torch.float32, device=dev) + dwq_part = torch.empty(T, W, dtype=torch.float32, device=dev) + if T > 0: + _ple_gate_bwd_kernel[(T * n, )]( + dgated, + key, + hc_state, + value, + wk, + wq, + gate, + rstdk, + rstdq, + dkey, + dquery, + dvalue_pern, + dwk_part, + dwq_part, + T, + N=n, + C=C, + SQRTC=_math.sqrt(C), + BLOCK_C=block_c) + dvalue = dvalue_pern.sum(dim=1).to(value.dtype) + dwk = dwk_part.sum(dim=0).to(wk.dtype) + dwq = dwq_part.sum(dim=0).to(wq.dtype) + dconv_w = dconvw.view(convw2d.shape).view(W, 1, Kk).to(conv_w_dtype) + + return (dquery, dkey, dvalue, dwk, dwq, dwc, dconv_w, None, None, None, None) + + +def ple_gate_conv_triton(hc_state, key, value, norm_key_w, norm_query_w, norm_conv_w, conv1d_weight, n, eps, dilation, + seq_len): + """Fused PLE increment (gate chain + norm_conv + causal dilated conv + SiLU + + residual). fp32 accumulation, output dtype = ``hc_state.dtype``. + + Returns ``None`` when the fast path is unavailable so callers fall back. + """ + if not HAVE_TRITON or hc_state.device.type == 'cpu': + return None + return _PLEGateConv.apply(hc_state.contiguous(), key.contiguous(), value.contiguous(), norm_key_w, norm_query_w, + norm_conv_w, conv1d_weight, n, eps, dilation, seq_len) diff --git a/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py b/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py new file mode 100644 index 0000000..f0b8586 --- /dev/null +++ b/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py @@ -0,0 +1,611 @@ +# VENDORED FILE -- DO NOT EDIT LOCALLY. +# +# refer : radixark/miles, miles_plugins/models/qwen3_8_next/ops/kernel/qsa_block_sparse_attn.py +# +# CONTRACT THIS KERNEL RELIES ON (verified against our indexer, see qsa_kernels.py): +# The selection rows must be whole, contiguous blocks of `block_size` tokens. The +# kernel tests membership at block granularity and restores exactness with the lo/hi +# range test, so a hand-built row like [0, 5] with block_size=4 is NOT computed as +# "attend {0,5}" -- it attends the whole 0..5 range (measured: 16.62 vs 50.0 on a +# probe where only the selected keys should contribute). Our QSAIndexer satisfies +# this by construction (`top_blocks * R + arange(R)`); there is no runtime check, so +# any new caller must uphold it by construction too. +"""Tensor-core QSA sparse attention for training: forward + backward. + +Same semantics as ``qsa_sparse_attn.py`` -- each query attends exactly the tokens in +its selection row -- but reached a different way, because the gather-per-query form +cannot use ``tl.dot``: with a distinct key set per query it has to materialise a +``[BQ, BK, D]`` tile and reduce it with ALU math. Measured at the production shape +(T=25k, 12 q-heads, D=256, budget 2048) that costs 2.5 s forward and 13.9 s +forward+backward per layer, which is ~330x slower than a DENSE causal flash kernel over +the same tensors (9.4 ms / 39.2 ms) even though dense does 6x more FLOPs. QSA at 25k +tokens only removes ~6x of the work, so paying 300x for the privilege is a large loss: +on the 16-node agentic run those three QSA layers per pipeline stage were ~49 s of a +~52 s micro-batch. + +So this kernel walks key tiles with ``tl.dot`` and masks each (query, key) pair to the +query's own selection. Tiles nobody in the query tile selected are skipped entirely via +a CSR-style per-query-tile list, so it beats a dense sweep by the coverage ratio rather +than merely matching it. The result is exact, not an approximation: + +- selection membership is tested at BLOCK granularity (the indexer picks blocks of + ``BLK`` consecutive tokens), which is a superset of the row's token list, because + expanding a block drops only the tokens past the query's own position; +- the dropped ones come back out via ``lo``/``hi``, the inclusive key range the caller + already computes (sequence start .. query position). + +bf16 inputs feed the dots with fp32 accumulation, the flash-attention convention, which +is also what the sglang kernel this mirrors does -- the old kernel's fp32 ALU path was +the odd one out. +""" + +import torch +import triton +import triton.language as tl +from torch import Tensor + + +@triton.jit +def _qsa_bs_fwd_kernel( + Q, + K, + V, + SEL, + LO, + HI, + BLKBASE, + TOKBASE, + KLIST, + KCNT, + OUT, + LSE, + stride_qt, + stride_qh, + stride_kt, + stride_kh, + stride_vt, + stride_vh, + stride_st, + stride_kl, + stride_ot, + stride_oh, + T, + NB, + scale, + GROUP: tl.constexpr, + D: tl.constexpr, + BQ: tl.constexpr, + BK: tl.constexpr, + BLK: tl.constexpr, +): + pid_t = tl.program_id(0) + pid_h = tl.program_id(1) + kv_head = pid_h // GROUP + + offs_q = pid_t * BQ + tl.arange(0, BQ) + offs_d = tl.arange(0, D) + q_mask = offs_q < T + + q = tl.load( + Q + offs_q[:, None] * stride_qt + pid_h * stride_qh + offs_d[None, :], + mask=q_mask[:, None], + other=0.0, + ) + lo = tl.load(LO + offs_q, mask=q_mask, other=0) + hi = tl.load(HI + offs_q, mask=q_mask, other=-1) + blk_base = tl.load(BLKBASE + offs_q, mask=q_mask, other=0) + tok_base = tl.load(TOKBASE + offs_q, mask=q_mask, other=0) + + m_i = tl.full((BQ, ), float('-inf'), tl.float32) + l_i = tl.zeros((BQ, ), tl.float32) + acc = tl.zeros((BQ, D), tl.float32) + + # Only the key tiles this query tile actually selected, so the kernel beats a dense + # sweep by the coverage ratio instead of merely matching it. The list is the union + # over the tile's queries; per-query exactness comes from the mask below. + n_tiles = tl.load(KCNT + pid_t) + for i in range(0, n_tiles): + kt = tl.load(KLIST + pid_t * stride_kl + i) + offs_k = kt * BK + tl.arange(0, BK) + k_in = offs_k < T + + # per-sequence block grid: after the packed-indexer fix a sequence's blocks start + # at its own first token, which is not a multiple of BLK in a packed batch. + # Looked up per TOKEN, not per 4-token group: a sequence whose start is not a + # multiple of BLK has its per-sequence blocks straddling the global groups, so a + # group-granular lookup would assign some tokens to the wrong block. + blk = blk_base[:, None] + (offs_k[None, :] - tok_base[:, None]) // BLK + sel = tl.load( + SEL + offs_q[:, None] * stride_st + blk, + mask=q_mask[:, None] & k_in[None, :] & (blk >= 0) & (blk < NB), + other=0, + ) + ok = (sel != 0) & (offs_k[None, :] <= hi[:, None]) & (offs_k[None, :] >= lo[:, None]) & k_in[None, :] + + k_tile = tl.load( + K + offs_k[:, None] * stride_kt + kv_head * stride_kh + offs_d[None, :], + mask=k_in[:, None], + other=0.0, + ) + s = tl.dot(q, tl.trans(k_tile)) * scale + s = tl.where(ok, s, float('-inf')) + + m_new = tl.maximum(m_i, tl.max(s, axis=1)) + m_use = tl.where(m_new == float('-inf'), 0.0, m_new) + p = tl.exp(s - m_use[:, None]) + p = tl.where(ok, p, 0.0) + alpha = tl.where(m_i == float('-inf'), 0.0, tl.exp(m_i - m_use)) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + + v_tile = tl.load( + V + offs_k[:, None] * stride_vt + kv_head * stride_vh + offs_d[None, :], + mask=k_in[:, None], + other=0.0, + ) + acc += tl.dot(p.to(v_tile.dtype), v_tile) + m_i = m_new + + l_safe = tl.where(l_i == 0.0, 1.0, l_i) + out = acc / l_safe[:, None] + tl.store( + OUT + offs_q[:, None] * stride_ot + pid_h * stride_oh + offs_d[None, :], + out, + mask=q_mask[:, None], + ) + lse = tl.where(m_i == float('-inf'), float('-inf'), m_i + tl.log(l_safe)) + tl.store(LSE + pid_h * T + offs_q, lse, mask=q_mask) + + +@triton.jit +def _qsa_bs_dq_kernel( + Q, + K, + V, + SEL, + LO, + HI, + BLKBASE, + TOKBASE, + KLIST, + KCNT, + OUT, + LSE, + DO, + DQ, + DELTA, + stride_qt, + stride_qh, + stride_kt, + stride_kh, + stride_vt, + stride_vh, + stride_st, + stride_kl, + stride_ot, + stride_oh, + T, + NB, + scale, + GROUP: tl.constexpr, + D: tl.constexpr, + BQ: tl.constexpr, + BK: tl.constexpr, + BLK: tl.constexpr, +): + pid_t = tl.program_id(0) + pid_h = tl.program_id(1) + kv_head = pid_h // GROUP + + offs_q = pid_t * BQ + tl.arange(0, BQ) + offs_d = tl.arange(0, D) + q_mask = offs_q < T + + q = tl.load(Q + offs_q[:, None] * stride_qt + pid_h * stride_qh + offs_d[None, :], mask=q_mask[:, None], other=0.0) + do = tl.load( + DO + offs_q[:, None] * stride_ot + pid_h * stride_oh + offs_d[None, :], mask=q_mask[:, None], other=0.0) + lse = tl.load(LSE + pid_h * T + offs_q, mask=q_mask, other=0.0) + delta = tl.load(DELTA + pid_h * T + offs_q, mask=q_mask, other=0.0) + lse_safe = tl.where(lse == float('-inf'), 0.0, lse) + alive = lse != float('-inf') + + lo = tl.load(LO + offs_q, mask=q_mask, other=0) + hi = tl.load(HI + offs_q, mask=q_mask, other=-1) + blk_base = tl.load(BLKBASE + offs_q, mask=q_mask, other=0) + tok_base = tl.load(TOKBASE + offs_q, mask=q_mask, other=0) + + dq = tl.zeros((BQ, D), tl.float32) + n_tiles = tl.load(KCNT + pid_t) + for i in range(0, n_tiles): + kt = tl.load(KLIST + pid_t * stride_kl + i) + offs_k = kt * BK + tl.arange(0, BK) + k_in = offs_k < T + + # per-sequence block grid: after the packed-indexer fix a sequence's blocks start + # at its own first token, which is not a multiple of BLK in a packed batch. + # Looked up per TOKEN, not per 4-token group: a sequence whose start is not a + # multiple of BLK has its per-sequence blocks straddling the global groups, so a + # group-granular lookup would assign some tokens to the wrong block. + blk = blk_base[:, None] + (offs_k[None, :] - tok_base[:, None]) // BLK + sel = tl.load( + SEL + offs_q[:, None] * stride_st + blk, + mask=q_mask[:, None] & k_in[None, :] & (blk >= 0) & (blk < NB), + other=0, + ) + ok = ((sel != 0) + & (offs_k[None, :] <= hi[:, None]) + & (offs_k[None, :] >= lo[:, None]) + & k_in[None, :] + & alive[:, None]) + + k_tile = tl.load( + K + offs_k[:, None] * stride_kt + kv_head * stride_kh + offs_d[None, :], mask=k_in[:, None], other=0.0) + v_tile = tl.load( + V + offs_k[:, None] * stride_vt + kv_head * stride_vh + offs_d[None, :], mask=k_in[:, None], other=0.0) + + s = tl.dot(q, tl.trans(k_tile)) * scale + p = tl.exp(s - lse_safe[:, None]) + p = tl.where(ok, p, 0.0) + + dp = tl.dot(do, tl.trans(v_tile)) + ds = (p * (dp - delta[:, None]) * scale).to(k_tile.dtype) + + dq += tl.dot(ds, k_tile) + + tl.store(DQ + offs_q[:, None] * stride_qt + pid_h * stride_qh + offs_d[None, :], dq, mask=q_mask[:, None]) + + +@triton.jit +def _qsa_bs_dkdv_kernel( + Q, + K, + V, + SEL, + LO, + HI, + BLKBASE, + TOKBASE, + QLIST, + QCNT, + LSE, + DO, + DELTA, + DK, + DV, + stride_qt, + stride_qh, + stride_kt, + stride_kh, + stride_vt, + stride_vh, + stride_st, + stride_ql, + stride_ot, + stride_oh, + T, + NB, + scale, + GROUP: tl.constexpr, + D: tl.constexpr, + BQ: tl.constexpr, + BK: tl.constexpr, + BLK: tl.constexpr, +): + """dK/dV keyed on the KEY tile, so nothing needs atomics. + + Launched per KV head, not per query head: with GQA the ``GROUP`` query heads sharing a + KV head all contribute to the same dK/dV, so they have to be summed here. Writing them + from separate programs would have them overwrite each other (the gather kernel got away + with it only because it used atomic_add). + """ + pid_k = tl.program_id(0) + kv_head = tl.program_id(1) + + offs_k = pid_k * BK + tl.arange(0, BK) + offs_d = tl.arange(0, D) + k_in = offs_k < T + + k_tile = tl.load( + K + offs_k[:, None] * stride_kt + kv_head * stride_kh + offs_d[None, :], mask=k_in[:, None], other=0.0) + v_tile = tl.load( + V + offs_k[:, None] * stride_vt + kv_head * stride_vh + offs_d[None, :], mask=k_in[:, None], other=0.0) + dk = tl.zeros((BK, D), tl.float32) + dv = tl.zeros((BK, D), tl.float32) + + n_q = tl.load(QCNT + pid_k) + for i in range(0, n_q): + qt = tl.load(QLIST + pid_k * stride_ql + i) + offs_q = qt * BQ + tl.arange(0, BQ) + q_mask = offs_q < T + lo = tl.load(LO + offs_q, mask=q_mask, other=0) + hi = tl.load(HI + offs_q, mask=q_mask, other=-1) + blk_base = tl.load(BLKBASE + offs_q, mask=q_mask, other=0) + tok_base = tl.load(TOKBASE + offs_q, mask=q_mask, other=0) + + blk = blk_base[:, None] + (offs_k[None, :] - tok_base[:, None]) // BLK + sel = tl.load( + SEL + offs_q[:, None] * stride_st + blk, + mask=q_mask[:, None] & k_in[None, :] & (blk >= 0) & (blk < NB), + other=0, + ) + ok = ((sel != 0) + & (offs_k[None, :] <= hi[:, None]) + & (offs_k[None, :] >= lo[:, None]) + & k_in[None, :] + & q_mask[:, None]) + + for gh in range(0, GROUP): + qh = kv_head * GROUP + gh + q = tl.load( + Q + offs_q[:, None] * stride_qt + qh * stride_qh + offs_d[None, :], + mask=q_mask[:, None], + other=0.0, + ) + do = tl.load( + DO + offs_q[:, None] * stride_ot + qh * stride_oh + offs_d[None, :], + mask=q_mask[:, None], + other=0.0, + ) + lse = tl.load(LSE + qh * T + offs_q, mask=q_mask, other=0.0) + delta = tl.load(DELTA + qh * T + offs_q, mask=q_mask, other=0.0) + okh = ok & (lse[:, None] != float('-inf')) + + lse_safe = tl.where(lse == float('-inf'), 0.0, lse) + sc = tl.dot(q, tl.trans(k_tile)) * scale + p = tl.exp(sc - lse_safe[:, None]) + p = tl.where(okh, p, 0.0) + + dp = tl.dot(do, tl.trans(v_tile)) + ds = (p * (dp - delta[:, None]) * scale).to(k_tile.dtype) + + dk += tl.dot(tl.trans(ds), q) + dv += tl.dot(tl.trans(p.to(do.dtype)), do) + + tl.store(DK + offs_k[:, None] * stride_kt + kv_head * stride_kh + offs_d[None, :], dk, mask=k_in[:, None]) + tl.store(DV + offs_k[:, None] * stride_vt + kv_head * stride_vh + offs_d[None, :], dv, mask=k_in[:, None]) + + +def selection_to_block_bitmap(indices: Tensor, num_tokens: int, block_size: int) -> Tensor: + """``[T, K]`` token indices (``-1`` pad) -> ``[T, ceil(T / block_size)]`` uint8 flags. + + A block is flagged when any of its tokens appears in the row. Tokens that the caller + clamped away inside an otherwise selected block are re-excluded by the ``lo``/``hi`` + range test in the kernel, so this stays exact while being ``block_size``x smaller. + """ + num_blocks = -(-num_tokens // block_size) + flags = torch.zeros(indices.shape[0], num_blocks, dtype=torch.uint8, device=indices.device) + valid = indices >= 0 + rows = torch.arange(indices.shape[0], device=indices.device).unsqueeze(1).expand_as(indices) + blocks = torch.where(valid, indices // block_size, torch.zeros_like(indices)) + flags[rows[valid], blocks[valid].long()] = 1 + return flags + + +def build_tile_index(sel: Tensor, bq: int, bk: int, block_size: int) -> tuple[Tensor, Tensor]: + """``[T, NB]`` block flags -> (``klist`` [NQT, maxc] int32, ``kcnt`` [NQT] int32). + + ``klist[i]`` lists, ascending, the key tiles that at least one query in query-tile + ``i`` selected. Per-query exactness still comes from the in-kernel mask; this only + decides which tiles are worth visiting. + """ + T, nb = sel.shape + bpt = bk // block_size + nqt = -(-T // bq) + nkt = -(-nb // bpt) + pad_b = nkt * bpt - nb + pad_q = nqt * bq - T + if pad_b or pad_q: + sel = torch.nn.functional.pad(sel, (0, pad_b, 0, pad_q)) + tile = sel.view(nqt, bq, nkt, bpt).amax(dim=3).amax(dim=1) > 0 # [nqt, nkt] + kcnt = tile.sum(dim=1).to(torch.int32) + maxc = max(int(kcnt.max().item()), 1) + order = torch.argsort((~tile).to(torch.int8), dim=1, stable=True) + klist = order[:, :maxc].contiguous().to(torch.int32) + return klist, kcnt + + +def build_tile_index_pair(sel: Tensor, bq: int, bk: int, block_size: int): + """Both directions of the tile map: (klist, kcnt) per query tile, (qlist, qcnt) per key tile. + + The transposed half is what lets dK/dV be keyed on the key tile and so avoid atomics. + """ + T, nb = sel.shape + bpt = bk // block_size + nqt = -(-T // bq) + nkt = -(-nb // bpt) + pad_b = nkt * bpt - nb + pad_q = nqt * bq - T + padded = sel + if pad_b or pad_q: + padded = torch.nn.functional.pad(sel, (0, pad_b, 0, pad_q)) + tile = padded.view(nqt, bq, nkt, bpt).amax(dim=3).amax(dim=1) > 0 + + def compact(mat): + cnt = mat.sum(dim=1).to(torch.int32) + maxc = max(int(cnt.max().item()), 1) + order = torch.argsort((~mat).to(torch.int8), dim=1, stable=True) + return order[:, :maxc].contiguous().to(torch.int32), cnt + + klist, kcnt = compact(tile) + qlist, qcnt = compact(tile.t().contiguous()) + return klist, kcnt, qlist, qcnt + + +class _QSABlockSparseAttn(torch.autograd.Function): + + @staticmethod + def forward(ctx, q, k, v, sel, lo, hi, blk_base, tok_base, scale, block_size): + T, Hq, D = q.shape + S, Hkv, _ = k.shape + assert Hq % Hkv == 0 + group = Hq // Hkv + qc, kc, vc = q.contiguous(), k.contiguous(), v.contiguous() + selc = sel.contiguous() + BQ_, BK_ = 64, 64 + klist, kcnt = build_tile_index(selc, BQ_, BK_, block_size) + o = torch.empty(T, Hq, D, device=q.device, dtype=torch.float32) + lse = torch.empty(Hq, T, device=q.device, dtype=torch.float32) + BQ, BK = 64, 64 + grid = (triton.cdiv(T, BQ), Hq) + _qsa_bs_fwd_kernel[grid]( + qc, + kc, + vc, + selc, + lo, + hi, + blk_base, + tok_base, + klist, + kcnt, + o, + lse, + qc.stride(0), + qc.stride(1), + kc.stride(0), + kc.stride(1), + vc.stride(0), + vc.stride(1), + selc.stride(0), + klist.stride(0), + o.stride(0), + o.stride(1), + T, + selc.shape[1], + scale, + GROUP=group, + D=D, + BQ=BQ, + BK=BK, + BLK=block_size, + num_warps=8, + num_stages=2, + ) + ctx.save_for_backward(qc, kc, vc, selc, lo, hi, blk_base, tok_base, klist, kcnt, o, lse) + ctx.scale = scale + ctx.group = group + ctx.block_size = block_size + return o.to(q.dtype) + + @staticmethod + def backward(ctx, grad_out): + qc, kc, vc, selc, lo, hi, blk_base, tok_base, _klist_fwd, _kcnt_fwd, o, lse = ctx.saved_tensors + T, Hq, D = qc.shape + do = grad_out.contiguous().to(qc.dtype) + # delta once, in torch: both backward kernels need it and neither should redo it + delta = (do.float() * o.float()).sum(-1).transpose(0, 1).contiguous() + dq = torch.empty(T, Hq, D, device=qc.device, dtype=torch.float32) + dk = torch.zeros(kc.shape, device=kc.device, dtype=torch.float32) + dv = torch.zeros(vc.shape, device=vc.device, dtype=torch.float32) + + BQ, BK = 64, 32 + klist, kcnt, qlist, qcnt = build_tile_index_pair(selc, BQ, BK, ctx.block_size) + common = ( + qc.stride(0), + qc.stride(1), + kc.stride(0), + kc.stride(1), + vc.stride(0), + vc.stride(1), + selc.stride(0), + ) + _qsa_bs_dq_kernel[(triton.cdiv(T, BQ), Hq)]( + qc, + kc, + vc, + selc, + lo, + hi, + blk_base, + tok_base, + klist, + kcnt, + o, + lse, + do, + dq, + delta, + *common, + klist.stride(0), + do.stride(0), + do.stride(1), + T, + selc.shape[1], + ctx.scale, + GROUP=ctx.group, + D=D, + BQ=BQ, + BK=BK, + BLK=ctx.block_size, + num_warps=8, + num_stages=1, + ) + _qsa_bs_dkdv_kernel[(triton.cdiv(T, BK), kc.shape[1])]( + qc, + kc, + vc, + selc, + lo, + hi, + blk_base, + tok_base, + qlist, + qcnt, + lse, + do, + delta, + dk, + dv, + *common, + qlist.stride(0), + do.stride(0), + do.stride(1), + T, + selc.shape[1], + ctx.scale, + GROUP=ctx.group, + D=D, + BQ=BQ, + BK=BK, + BLK=ctx.block_size, + num_warps=8, + num_stages=1, + ) + return dq.to(qc.dtype), dk.to(kc.dtype), dv.to(vc.dtype), None, None, None, None, None, None, None + + +def qsa_block_sparse_attention_triton( + q: Tensor, + k: Tensor, + v: Tensor, + sel_blocks: Tensor, + lo: Tensor, + hi: Tensor, + blk_base: Tensor, + tok_base: Tensor, + scale: float, + block_size: int = 4, +) -> Tensor: + """``q`` [T, Hq, D], ``k``/``v`` [S, Hkv, D], ``sel_blocks`` [T, NB] uint8. + + ``lo``/``hi`` are the inclusive key range per query; ``blk_base``/``tok_base`` place + the query's sequence in the packed block grid (both zero for a single sequence). + """ + return _QSABlockSparseAttn.apply(q, k, v, sel_blocks, lo, hi, blk_base, tok_base, scale, block_size) + + +def qsa_sparse_attention_from_indices(q: Tensor, + k: Tensor, + v: Tensor, + indices: Tensor, + scale: float, + block_size: int = 4) -> Tensor: + """Drop-in for the gather kernel: derives the bitmap and range from ``indices``.""" + T = q.shape[0] + sel = selection_to_block_bitmap(indices, T, block_size) + valid = indices >= 0 + big = torch.iinfo(torch.int32).max + lo = torch.where(valid, indices, torch.full_like(indices, big)).min(dim=1).values.to(torch.int32) + hi = torch.where(valid, indices, torch.full_like(indices, -1)).max(dim=1).values.to(torch.int32) + zeros = torch.zeros(T, dtype=torch.int32, device=q.device) + return qsa_block_sparse_attention_triton(q, k, v, sel, lo, hi, zeros, zeros, scale, block_size) diff --git a/src/mcore_bridge/model/modules/kernels/qsa_kernels.py b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py new file mode 100644 index 0000000..9595717 --- /dev/null +++ b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py @@ -0,0 +1,289 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""QSA sparse attention wrappers around the vendored tensor-core triton kernel. + +The kernel itself lives in ``qsa_block_sparse_attn.py``, vendored verbatim from +miles PR #2777 (commit 0f5dff4). This file owns only the glue mcore needs: +sbhd<->thd flattening, context parallelism, and the ``core_attention`` shim. + +WHY A KERNEL AT ALL (the decision this file encodes) + QSA needs a per-query key set, which as an attention mask is TE's + ``arbitrary`` type. TE refuses that under packing: + ``dot_product_attention.py:1347`` asserts ``"padding" in attn_mask_type`` + whenever ``qkv_format == "thd"``, so a bool mask cannot express QSA there. + Context parallelism is blocked for a different reason -- block pooling needs + keys from other CP ranks, which the mask path never gathers. + + So the layer picks by data shape, with no user-facing switch: + + thd (padding_free) or CP>1 -> this kernel + sbhd and CP==1 -> the indexer's bool mask on TE + + The mask path is not a degraded fallback in that last case: both paths run + the same ``QSAIndexer._score_and_topk_blocks`` top-block selection, so they are + mathematically equivalent there, and the mask path measures *faster* + (8192: 8.4 s/it vs 25.9 s/it) because TE keeps the whole thing fused. +""" +import torch + +from .qsa_block_sparse_attn import qsa_sparse_attention_from_indices + +try: + import triton # noqa: F401 (import guard for the vendored kernel) + HAVE_TRITON = True +except Exception: # pragma: no cover - triton absent + HAVE_TRITON = False + + +def qsa_sparse_supported(head_dim: int) -> bool: + """Whether the sparse kernel can run for this head dim. + + Triton must be importable and ``head_dim`` must be a power of two (the + kernel tiles the head with ``tl.arange`` blocks). + """ + return HAVE_TRITON and head_dim > 0 and not (head_dim & (head_dim - 1)) + + +def _cp_query_global_positions(seq_len: int, cp_size: int, cp_rank: int, device) -> torch.Tensor: + """This CP rank's logical token positions under mcore zigzag sharding. + + Matches ``split_cp_inputs`` / ``rope_utils.get_pos_emb_on_this_cp_rank``: the + sequence is viewed as ``2 * cp_size`` chunks and the rank owns the + ``cp_rank``-th front chunk plus its mirrored back chunk. + """ + chunk = seq_len // (2 * cp_size) + front = torch.arange(cp_rank * chunk, (cp_rank + 1) * chunk, device=device, dtype=torch.int64) + back_chunk = 2 * cp_size - cp_rank - 1 + back = torch.arange(back_chunk * chunk, (back_chunk + 1) * chunk, device=device, dtype=torch.int64) + return torch.cat((front, back), dim=0) + + +def _cp_query_global_positions_thd(cu_seqlens: torch.Tensor, cp_size: int, cp_rank: int, device) -> torch.Tensor: + """Local packed-token positions per sample under zigzag thd CP sharding. + + Each sample is padded to a multiple of ``2 * cp_size`` by the data pipeline; + the rank owns the ``cp_rank``-th front chunk and the mirrored back chunk of + every sample (the same partition ``split_cp_inputs`` applies per sample). + """ + cu = cu_seqlens.to(device=device, dtype=torch.int64) + starts, ends = cu[:-1], cu[1:] + half = (ends - starts) // (2 * cp_size) # per-sample chunk length + # front chunk cp_rank; back chunk (2*cp-1-cp_rank), whose start is + # ends - (cp_rank+1)*half since it is the (cp_rank+1)-th chunk from the end + seg_starts = torch.stack((starts + cp_rank * half, ends - (cp_rank + 1) * half), dim=1).reshape(-1) + seg_lens = torch.stack((half, half), dim=1).reshape(-1) + nz = seg_lens > 0 + seg_starts, seg_lens = seg_starts[nz], seg_lens[nz] + seg_ids = torch.repeat_interleave(torch.arange(seg_lens.numel(), device=device), seg_lens) + offsets = torch.arange(int(seg_lens.sum().item()), device=device) + offsets = offsets - torch.repeat_interleave(torch.cumsum(seg_lens, 0) - seg_lens, seg_lens) + return seg_starts.index_select(0, seg_ids) + offsets + + +def _cp_gathered_to_logical_order(seq_len: int, cp_size: int, device) -> torch.Tensor: + """Index restoring rank-major gathered chunks to logical order: applying + ``gathered[idx]`` yields the same layout as ``_undo_attention_load_balancing`` + (the inverse of ``split_cp_inputs``).""" + chunk = seq_len // (2 * cp_size) + order = [2 * i for i in range(cp_size)] + [2 * cp_size - 2 * i - 1 for i in range(cp_size)] + return torch.cat([torch.arange(c * chunk, (c + 1) * chunk, device=device) for c in order]) + + +def qsa_sparse_attention_thd(q, k, v, indices, scale, block_size): + """``q`` [T, Hq, D], ``k``/``v`` [S, Hkv, D], ``indices`` [T, K] (-1 pad). + + Token-space indices: works for packed (thd) inputs directly, and for any + pre-flattened token dimension. + + ``block_size`` must be the indexer's ``compress_ratio``. The kernel tests + selection membership per block of that many consecutive tokens, so passing a + different value silently changes which keys are attended -- see the contract + note in qsa_block_sparse_attn.py. + """ + if not HAVE_TRITON or q.device.type == 'cpu': + raise RuntimeError('QSA sparse attention requires triton and accelerator tensors ' + f'(HAVE_TRITON={HAVE_TRITON}, device={q.device.type}). This path is only ' + 'selected for packing (thd) or CP>1, where no dense fallback is correct.') + if q.shape[-1] & (q.shape[-1] - 1): + raise RuntimeError(f'QSA sparse attention needs a power-of-two head dim, got {q.shape[-1]}.') + if q.shape[0] != k.shape[0]: + # The kernel takes its key bound from the query count (T, Hq, D = q.shape, + # then `offs_k < T`), so unequal lengths would silently drop every key past + # len(q). Callers must equalise first -- _forward_cp does this by scattering + # the local query shard into a full-length buffer. + raise ValueError(f'QSA sparse attention needs len(q) == len(k), got {q.shape[0]} vs {k.shape[0]}.') + return qsa_sparse_attention_from_indices(q, k, v, indices.contiguous(), scale, block_size) + + +def qsa_sparse_attention(q, k, v, indices, scale, block_size): + """QSA sparse attention over mcore layouts. + + sbhd: ``q`` [s, b, Hq, D], ``k``/``v`` [s, b, Hkv, D], ``indices`` + [b, s, K] in per-sample sequence space. + thd: ``q`` [T, Hq, D], ``k``/``v`` [T, Hkv, D], ``indices`` [T, K] in + pack space (document boundaries already clamped torch-side). + + Raises rather than returning ``None``: the caller only reaches here for + packing or CP>1, and for those a dense fallback would silently diverge from + sparse inference. + """ + if q.dim() == 3: + return qsa_sparse_attention_thd(q, k, v, indices, scale, block_size) + if q.dim() != 4: + raise ValueError(f'qsa_sparse_attention expected 3D (thd) or 4D (sbhd) q, got {tuple(q.shape)}') + s, b, hq, d = q.shape + sk = k.shape[0] + # Flatten to token space, batch-major: token t = r * sk + p. The per-sample + # sequence-space indices are offset by the batch start; -1 padding stays -1. + # The offset is a whole multiple of sk, so block alignment survives it only + # when sk % block_size == 0; guard rather than corrupt the selection. + if sk % block_size: + raise ValueError(f'sbhd QSA needs the kv sequence length ({sk}) to be a multiple of ' + f'block_size ({block_size}); otherwise flattening to token space shifts ' + 'each sample off the block grid the kernel indexes by.') + q_f = q.permute(1, 0, 2, 3).reshape(b * s, hq, d) + k_f = k.permute(1, 0, 2, 3).reshape(b * sk, *k.shape[2:]) + v_f = v.permute(1, 0, 2, 3).reshape(b * sk, *v.shape[2:]) + off = torch.arange(b, device=q.device).view(b, 1, 1) * sk + idx_f = torch.where(indices >= 0, indices + off, indices.new_full((), -1)).reshape(b * s, -1) + out_f = qsa_sparse_attention_thd(q_f, k_f, v_f, idx_f, scale, block_size) + return out_f.view(b, s, hq, d).permute(1, 0, 2, 3) + + +class QSASparseCoreAttention(torch.nn.Module): + """Drop-in ``core_attention`` that runs the QSA sparse triton kernel. + + The selection indices ride in on the ``attention_mask`` argument (int64) -- + the one positional that mcore threads unchanged through both the plain and + the activation-checkpointed core-attention paths. Carrying them there (rather + than on a module attribute) means recompute under + ``recompute_modules=['core_attn']`` sees the same indices, and a stacked + microbatch cannot overwrite the selection between forward and its backward. + + Any bool/None mask (or non-int tensor) falls through to the wrapped + core_attention, keeping the no-op / short-sequence path on TE's fused kernel. + + Context parallelism (``cp_comm_type=allgather``): q stays on the local CP + shard while k/v are gathered across the CP group and restored to logical + order; indices (full-sequence) are sliced to this rank's query rows. The + gather's backward reduce-scatters dk/dv to the local shards. + """ + + def __init__(self, core_attention, config, softmax_scale=None): + super().__init__() + self.core_attention = core_attention + # None -> resolved to 1/sqrt(head_dim) at call time, matching TE. + self.softmax_scale = softmax_scale + self.config = config + # The kernel tests selection membership per block of compress_ratio + # consecutive tokens, so this must be the same value the indexer expanded + # its top-k blocks with -- see the contract in qsa_block_sparse_attn.py. + self.block_size = config.indexer_compress_ratio + if not self.block_size: + raise ValueError('QSASparseCoreAttention needs config.indexer_compress_ratio to size the ' + f'kernel block grid, got {self.block_size!r}.') + + def forward(self, + query, + key, + value, + attention_mask, + attn_mask_type=None, + attention_bias=None, + packed_seq_params=None, + **kwargs): + if attention_mask is not None and attention_mask.dtype in (torch.int32, torch.int64): + scale = self.softmax_scale if self.softmax_scale is not None else query.shape[-1]**-0.5 + cp_size = self.config.context_parallel_size + if cp_size > 1: + out = self._forward_cp(query, key, value, attention_mask, scale, packed_seq_params) + else: + out = qsa_sparse_attention(query, key, value, attention_mask, scale, self.block_size) + # mcore's sbhd core attention returns heads flattened ([s, b, h*d]); + # the kernel yields [s, b, h, d]. thd stays [t, h, d] -- mcore + # reshapes it itself right after (see attention.py thd branch). + if out.dim() == 4: + out = out.reshape(out.shape[0], out.shape[1], -1) + return out + return self.core_attention( + query, + key, + value, + attention_mask, + attn_mask_type=attn_mask_type, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + **kwargs) + + def _forward_cp(self, query, key, value, indices, scale, packed_seq_params): + from megatron.core import mpu + from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region + cp_size = mpu.get_context_parallel_world_size() + cp_rank = mpu.get_context_parallel_rank() + device = query.device + thd = query.dim() == 3 + if thd: + # the gathered k/v live in the padded pack space, so the query + # positions must come from the padded cu as well + cu_q = ( + packed_seq_params.cu_seqlens_q_padded + if packed_seq_params.cu_seqlens_q_padded is not None else packed_seq_params.cu_seqlens_q) + q_pos = _cp_query_global_positions_thd(cu_q, cp_size, cp_rank, device) + # rank-major gathered positions -> permutation back to global packed + # order (same construction as mcore DSA's packed kv reorder) + gathered_pos = torch.cat([_cp_query_global_positions_thd(cu_q, cp_size, r, device) for r in range(cp_size)]) + kv_reorder = torch.argsort(gathered_pos) + else: + sq, b = query.shape[0], query.shape[1] + q_pos = _cp_query_global_positions(sq * cp_size, cp_size, cp_rank, device) + kv_reorder = _cp_gathered_to_logical_order(sq * cp_size, cp_size, device) + # gather k/v across CP with a DIFFERENTIABLE all-gather (backward is a + # reduce-scatter, so dk/dv reach the local shards) and undo the zigzag + # shard into full logical order. A raw torch.distributed.all_gather is + # not differentiable -- reconstruct_tensor_cp is not usable here. + cp_group = mpu.get_context_parallel_group() + kv_reorder_t = kv_reorder + + def _gather_full(t): + g = gather_from_sequence_parallel_region(t, tensor_parallel_output_grad=True, group=cp_group) + return g.index_select(0, kv_reorder_t) + + key_full = _gather_full(key) + value_full = _gather_full(value) + # The kernel derives the key bound from the query count (T, Hq, D = q.shape, + # then `offs_k < T`), so it structurally requires len(q) == len(k). Under CP + # the queries are a 1/cp_size shard while k/v are now full length, so scatter + # the local queries back into a full-length buffer, run, and take our rows + # out again. The padding rows carry an all-`-1` selection, which the kernel + # skips, so they cost tile launches but produce nothing. + if thd: + local_idx = indices[q_pos] + out_full = qsa_sparse_attention(*self._scatter_q_to_full(query, key_full, value_full, local_idx, q_pos), + scale, self.block_size) + return out_full.index_select(0, q_pos) + # sbhd: token-space kernel on the batch-major flattening (t = r*sk + p) + local_idx = indices[:, q_pos] + sk = key_full.shape[0] + k_f = key_full.permute(1, 0, 2, 3).reshape(b * sk, key_full.shape[2], key_full.shape[3]) + v_f = value_full.permute(1, 0, 2, 3).reshape(b * sk, value_full.shape[2], value_full.shape[3]) + off = torch.arange(b, device=device).view(b, 1, 1) * sk + idx_f = torch.where(local_idx >= 0, local_idx + off, local_idx.new_full((), -1)).reshape(sq * b, -1) + q_f = query.permute(1, 0, 2, 3).reshape(sq * b, query.shape[2], query.shape[3]) + # batch-major token ids of this rank's rows: sample r contributes q_pos + r*sk + rows = (q_pos[None, :] + torch.arange(b, device=device).view(b, 1) * sk).reshape(-1) + out_f = qsa_sparse_attention(*self._scatter_q_to_full(q_f, k_f, v_f, idx_f, rows), scale, self.block_size) + out_f = out_f.index_select(0, rows) + return out_f.view(b, sq, query.shape[2], query.shape[3]).permute(1, 0, 2, 3) + + @staticmethod + def _scatter_q_to_full(q, k, v, indices, rows): + """Place ``q``/``indices`` rows at ``rows`` inside a len(k)-row buffer. + + index_copy keeps this differentiable: backward gathers the same rows, so the + padded positions contribute no gradient. + """ + n = k.shape[0] + q_full = q.new_zeros((n, *q.shape[1:])) + q_full = q_full.index_copy(0, rows, q) + idx_full = indices.new_full((n, indices.shape[1]), -1) + idx_full = idx_full.index_copy(0, rows, indices) + return q_full, k, v, idx_full diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py index 1121bd2..e97a1a0 100644 --- a/src/mcore_bridge/model/modules/ple.py +++ b/src/mcore_bridge/model/modules/ple.py @@ -3,14 +3,18 @@ import math import torch import torch.nn.functional as F +from megatron.core import parallel_state +from megatron.core.extensions.transformer_engine import TELinear from megatron.core.tensor_parallel import VocabParallelEmbedding from megatron.core.tensor_parallel.mappings import (gather_from_sequence_parallel_region, scatter_to_sequence_parallel_region) from torch import nn from typing import List, Optional +from ...utils import get_env_args, get_logger from ...utils.megatron_utils import reconstruct_tensor_cp, split_cp_inputs from .hyper_connection_gated import Qwen4ExpTextGroupedRMSNorm +from .kernels import gather_ple_rows, ple_gate_conv_triton _MASK64 = (1 << 64) - 1 _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15 @@ -19,6 +23,15 @@ _PRIME_1 = 10007 +def use_ple_cpu_offload() -> bool: + return get_env_args('PLE_CPU_OFFLOAD', bool, False) + + +def use_ple_fused_kernel() -> bool: + # On by default; PLE_FUSED_KERNEL=0 forces the eager chain. + return get_env_args('PLE_FUSED_KERNEL', bool, True) + + def _splitmix64(value: int) -> int: value = (value + _SPLITMIX_GAMMA) & _MASK64 value = ((value ^ (value >> 30)) * _SPLITMIX_M1) & _MASK64 @@ -99,17 +112,24 @@ def __init__(self, config, ple_layer_index: int): # validates them for qwen4_exp); a silently substituted default would # corrupt the weight conversion. (transformers reads eos_token_id # directly and has no split_ngram_parts: its table is replicated.) - eos_token_id = getattr(config, 'eos_token_id', None) + eos_token_id = config.eos_token_id + split_ngram_parts = config.split_ngram_parts ple_seed = getattr(config, 'ple_seed', None) - split_ngram_parts = getattr(config, 'split_ngram_parts', None) - if eos_token_id is None or ple_seed is None or split_ngram_parts is None: - raise ValueError(f'eos_token_id/ple_seed/split_ngram_parts must be provided by the model ' - f'config (got {eos_token_id!r}/{ple_seed!r}/{split_ngram_parts!r}).') + if ple_seed is None: + raise ValueError('ple_seed must be provided by the model config (the parser derives it from ' + 'text_config.seed); a silently substituted default would desynchronize the ' + 'n-gram hash multipliers from transformers.') + if eos_token_id is None or split_ngram_parts is None: + raise ValueError(f'eos_token_id/split_ngram_parts must be provided by the model ' + f'config (got {eos_token_id!r}/{split_ngram_parts!r}).') self.eos_token_id = int(eos_token_id) self.split_ngram_parts = int(split_ngram_parts) self.head_dim = head_dim_per_ngram # mcore-specific: the bridge weight conversion reads it off the module. - # Multipliers (splitmix64 derived, checkpoint-persistent). + # Multipliers (splitmix64 derived, checkpoint-persistent). The seed comes + # from the config (parser defaults it to Qwen4ExpTextConfig.seed's 1234); + # the multipliers must stay bit-identical to transformers or the n-gram + # lookups desynchronize. multipliers = _build_layer_multipliers(config.padded_vocab_size, self.ngram_size, self.ple_layer_index, int(ple_seed)) self.register_buffer('layer_multipliers', torch.tensor(multipliers, dtype=torch.long), persistent=True) @@ -130,13 +150,81 @@ def __init__(self, config, ple_layer_index: int): self.register_buffer('ngram_heads_offsets', torch.tensor(self.head_offsets, dtype=torch.long), persistent=True) ngram_vocab_divisor = config.make_ngram_vocab_size_divisible_by padded_vocab_size = math.ceil(self.total_vocab_size / ngram_vocab_divisor) * ngram_vocab_divisor - # mcore-specific: TP-sharded table (a replicated nn.Embedding would be ~80GB). - self.ngram_embedding = VocabParallelEmbedding( - padded_vocab_size, - head_dim_per_ngram, - init_method=torch.nn.init.normal_, - config=config, - ) + self.padded_vocab_size = padded_vocab_size + self.cpu_offload = use_ple_cpu_offload() + if self.cpu_offload: + # Warn rather than raise: whether this is safe depends on the train type, + # which the model config does not carry (see use_ple_cpu_offload). + get_logger().warning_once('PLE_CPU_OFFLOAD=1: the n-gram table is host-resident and receives no ' + 'gradient, so it stays frozen.') + self._init_host_table(config, padded_vocab_size, head_dim_per_ngram) + else: + # mcore-specific: TP-sharded table (a replicated nn.Embedding would be ~80GB). + self.ngram_embedding = VocabParallelEmbedding( + padded_vocab_size, + head_dim_per_ngram, + init_method=torch.nn.init.normal_, + config=config, + ) + + def fill_table_from_hf(self, hf_state_dict): + """Populate the (local TP partition of the) n-gram table from the HF + checkpoint shards at load time. + """ + total = self.padded_vocab_size + parts = self.split_ngram_parts + shard_size = (total + parts - 1) // parts + tp_rank = parallel_state.get_tensor_model_parallel_rank() + if self.cpu_offload: + tp_start, tp_end = self.vocab_start, self.vocab_end + dtype = self.host_table.dtype + else: + emb = self.ngram_embedding + per_partition = emb.num_embeddings_per_partition + tp_start = tp_rank * per_partition + tp_end = min((tp_rank + 1) * per_partition, total) + dtype = emb.weight.dtype + for i in range(parts): + key = f'ple.ple_embedding.ngram_embedding.shard_{i}.weight' + if key not in hf_state_dict: + continue + cs, ce = i * shard_size, min((i + 1) * shard_size, total) + s, e = max(cs, tp_start), min(ce, tp_end) + if s >= e: + continue + weight = hf_state_dict[key].load() + if self.cpu_offload: + self.host_table[s - tp_start:e - tp_start] = weight[s - cs:e - cs].to( + dtype=dtype, device=self.host_table.device) + else: + emb.weight.data[s - tp_start:e - tp_start] = weight[s - cs:e - cs].to(dtype) + + @torch.no_grad() + def export_table_to_hf(self, hf_state_dict, prefix=''): + """Reverse of ``fill_table_from_hf``: write the offloaded table back as HF + shards so a full-parameter checkpoint is self-contained. + """ + if not self.cpu_offload: + return + total = self.padded_vocab_size + parts = self.split_ngram_parts + shard_size = (total + parts - 1) // parts + tp_rank = parallel_state.get_tensor_model_parallel_rank() + tp_group = parallel_state.get_tensor_model_parallel_group() + for i in range(parts): + cs, ce = i * shard_size, min((i + 1) * shard_size, total) + # Reduce on GPU: NCCL has no CPU backend, and the host table is pinned + # CPU. Each rank scatters its owned rows into a full shard, sums across + # TP (rows are disjoint, so sum == gather), then rank 0 keeps the CPU copy. + device = torch.cuda.current_device() + local = torch.zeros(ce - cs, self.host_table.shape[-1], dtype=self.host_table.dtype, device=device) + s, e = max(cs, self.vocab_start), min(ce, self.vocab_end) + if s < e: + local[s - cs:e - cs] = self.host_table[s - self.vocab_start:e - self.vocab_start].to(device) + if self._tp_size > 1: + torch.distributed.all_reduce(local, group=tp_group) + if tp_rank == 0: + hf_state_dict[f'{prefix}ple.ple_embedding.ngram_embedding.shard_{i}.weight'] = local.cpu() def _shift_right_ignore_eos(self, token_ids: torch.Tensor, shift: int) -> torch.Tensor: # Mirrors transformers `_shift_right_ignore_eos`: segment-aware shift, @@ -189,8 +277,42 @@ def forward(self, input_ids: torch.Tensor) -> torch.Tensor: blocks.append(ngram_ids + head_offsets.view(1, 1, -1)) ngram_ids = torch.cat(blocks, dim=-1)[:, -input_ids.shape[1]:] + if self.cpu_offload: + return self._host_lookup(ngram_ids) return self.ngram_embedding(ngram_ids).flatten(-2) + def _init_host_table(self, config, padded_vocab_size, head_dim): + tp_size = parallel_state.get_tensor_model_parallel_world_size() + tp_rank = parallel_state.get_tensor_model_parallel_rank() + per_partition = (padded_vocab_size + tp_size - 1) // tp_size + self.vocab_start = tp_rank * per_partition + self.vocab_end = min(padded_vocab_size, self.vocab_start + per_partition) + n_local = max(0, self.vocab_end - self.vocab_start) + self.host_table = torch.empty(n_local, head_dim, dtype=config.params_dtype, pin_memory=True) + self._tp_size = tp_size + self._tp_group = parallel_state.get_tensor_model_parallel_group() + + def _host_lookup(self, ngram_ids): + # ngram_ids: [rows, L, nH] global ids in [0, padded_vocab_size). + # Triton fast path reads rows straight out of the pinned host table; + # otherwise fall back to a torch gather. Either way, all-reduce across + # TP (partitions are disjoint, so the sum is the full embedding). + rows = gather_ple_rows(self.host_table, ngram_ids.reshape(-1), self.vocab_start, self.vocab_end) + if rows is not None: + rows = rows.view(*ngram_ids.shape, self.host_table.shape[-1]) + if self._tp_size > 1: + torch.distributed.all_reduce(rows, group=self._tp_group) + return rows.flatten(-2) + ids_cpu = ngram_ids.detach().to('cpu') + local_mask = (ids_cpu >= self.vocab_start) & (ids_cpu < self.vocab_end) + local_ids = (ids_cpu - self.vocab_start).clamp(0, max(self.host_table.shape[0] - 1, 0)) + rows = self.host_table[local_ids] # [rows, L, nH, head_dim] + rows = rows * local_mask.unsqueeze(-1).to(rows.dtype) + rows = rows.to(ngram_ids.device, non_blocking=True) + if self._tp_size > 1: + torch.distributed.all_reduce(rows, group=self._tp_group) + return rows.flatten(-2) + class Qwen4ExpTextPLELayer(nn.Module): """Inject hashed n-gram features into every hyper-connection stream; @@ -223,8 +345,26 @@ def __init__(self, config, ple_layer_index: int, pg_collection=None): conv_dilation = int(config.ngram_size) self.short_conv_state_len = (conv_kernel_size - 1) * conv_dilation # Replicated projections in params_dtype. - self.key_proj = nn.Linear(ple_embed_dim, hc_hidden_size, bias=False, dtype=config.params_dtype) - self.value_proj = nn.Linear(ple_embed_dim, self.hidden_size, bias=False, dtype=config.params_dtype) + self.key_proj = TELinear( + input_size=ple_embed_dim, + output_size=hc_hidden_size, + parallel_mode='duplicated', + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=True, + skip_weight_param_allocation=False, + ) + self.value_proj = TELinear( + input_size=ple_embed_dim, + output_size=self.hidden_size, + parallel_mode='duplicated', + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=True, + skip_weight_param_allocation=False, + ) # mcore's config field layernorm_epsilon corresponds to HF's rms_norm_eps; # the grouped norm is the mcore subclass adding dtype/SP-flag construction. self.norm_key = Qwen4ExpTextGroupedRMSNorm( @@ -276,8 +416,12 @@ def compute(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch """Mirrors transformers ``Qwen4ExpTextPLELayer.forward`` (training variant); hidden_states/input_ids: [rows, L, nH]/[rows, L].""" embeddings = self.ple_embedding(input_ids) # mcore-specific: no past_key_values cache arg - key_normed = self.norm_key(self.key_proj(embeddings)).unflatten(-1, (self.hc_count, self.hidden_size)) - value = self.value_proj(embeddings) + if use_ple_fused_kernel() and embeddings.device.type != 'cpu': + fused = self._compute_fused(hidden_states, embeddings) + if fused is not None: + return fused + key_normed = self.norm_key(self.key_proj(embeddings)[0]).unflatten(-1, (self.hc_count, self.hidden_size)) + value = self.value_proj(embeddings)[0] query_normed = self.norm_query(hidden_states).unflatten(-1, (self.hc_count, self.hidden_size)) gate = (key_normed * query_normed).sum(dim=-1, keepdim=True) / math.sqrt(self.hidden_size) gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign() @@ -289,6 +433,31 @@ def compute(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch output = gated_value + self._short_conv(gated_value_normed) return output + def _compute_fused(self, hidden_states: torch.Tensor, embeddings: torch.Tensor): + """Fused triton chain (gate + norm_conv + causal conv + SiLU + residual); + fp32 accumulation, output dtype = params_dtype. Returns None when the + kernel path is unavailable so the caller falls back to torch.""" + rows, seq_len = embeddings.shape[:2] + key = self.key_proj(embeddings)[0].reshape(rows * seq_len, -1) + value = self.value_proj(embeddings)[0].reshape(rows * seq_len, -1) + query = hidden_states.reshape(rows * seq_len, -1) + out = ple_gate_conv_triton( + query, + key, + value, + self.norm_key.weight, + self.norm_query.weight, + self.norm_conv.weight, + self.conv1d.weight, + self.hc_count, + self.config.layernorm_epsilon, + int(self.config.ngram_size), + seq_len, + ) + if out is None: + return None + return out.view(rows, seq_len, -1) + @staticmethod def _normalize_cu_seqlens(cu: Optional[torch.Tensor], total: int) -> Optional[torch.Tensor]: """Align a (possibly padded/offset) cu_seqlens against the gathered full length. @@ -327,9 +496,9 @@ def forward( sequence, and the additive output is scattered back. """ sp_on = ( - self.pg_collection is not None and getattr(self.config, 'sequence_parallel', False) - and getattr(self.config, 'tensor_model_parallel_size', 1) > 1) - cp_on = getattr(self.config, 'context_parallel_size', 1) > 1 + self.pg_collection is not None and self.config.sequence_parallel + and self.config.tensor_model_parallel_size > 1) + cp_on = self.config.context_parallel_size > 1 if not (sp_on or cp_on): return self._forward_impl(hidden_states, input_ids, packed_seq_params) @@ -383,11 +552,7 @@ def _forward_impl( thd = packed_seq_params is not None and getattr(packed_seq_params, 'qkv_format', 'bshd') == 'thd' if thd: num_samples = packed_seq_params.num_samples - # PackedSeqParams.max_seqlen_q is declared `int` in mcore and swift - # normalizes it to int, so `.item()` would raise AttributeError; - # tolerate a 0-d tensor from other callers. - max_seqlen_q = packed_seq_params.max_seqlen_q - max_len = int(max_seqlen_q.item() if torch.is_tensor(max_seqlen_q) else max_seqlen_q) + max_len = int(packed_seq_params.max_seqlen_q) cu = packed_seq_params.cu_seqlens_q total = hidden_states.shape[0] hid = hidden_states.new_zeros((num_samples, max_len, hidden_states.shape[-1])) diff --git a/src/mcore_bridge/model/modules/qsa_indexer.py b/src/mcore_bridge/model/modules/qsa_indexer.py index 193bcbb..080e28a 100644 --- a/src/mcore_bridge/model/modules/qsa_indexer.py +++ b/src/mcore_bridge/model/modules/qsa_indexer.py @@ -1,6 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import math import torch +from megatron.core.extensions.transformer_engine import TELinear from torch import nn @@ -26,8 +27,54 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor: return torch.cat((-x2, x1), dim=-1) +def _materialize_rope(freqs: torch.Tensor, seq_len: int, dtype: torch.dtype, mscale: float): + """Materialize batch-aware RoPE cos/sin from mcore rotary angles. + + ``[s, freq_b, 1, rot] -> [freq_b, s, rot]``. mcore stores angles rather than + cos/sin, so they are materialized the way ``_patch_apply_rotary_pos_emb`` + does, keeping the indexer's RoPE identical to the attention's. + + Keeping ``freq_b`` as its own dim is what makes MRoPE correct: its positions + differ per sample, so flattening here would fold batch into the rotary + feature dim and make ``rot`` come out as ``b * rot``. + """ + f = freqs[:seq_len].squeeze(2).permute(1, 0, 2) + cos = (torch.cos(f) * mscale).to(dtype) + sin = (torch.sin(f) * mscale).to(dtype) + return cos, sin + + class QSAIndexer(nn.Module): - # refer: transformers Qwen4ExpTextQSAIndexer + """QSA block selection: score compressed key blocks, keep the top-k per query. + + refer: transformers ``Qwen4ExpTextQSAIndexer``. That reference packs scoring and + mask construction into one ``forward`` with a per-query Python loop; here the + scoring is factored out so the two output encodings provably agree. + + Two layers, not four peers:: + + _score_and_topk_blocks(hidden, freqs) the actual selection + | -> (top_blocks, keep, n_blocks) + +-- selection_as_mask(...) encode as [b, 1, s, s] bool + +-- selection_as_token_indices(...) encode as [b, s, K] int64 + + select_token_indices_thd(...) SEPARATE implementation for thd; + does NOT reuse the scorer + + Both encoders expand the same blocks to tokens and append the query's own + partial-block tail, so they select an identical key set -- only the wire format + differs. Which one the layer asks for (``qwen4_exp.py:_qsa_select``):: + + sbhd, CP==1 -> selection_as_mask -> TE `arbitrary` mask + thd or CP>1 -> selection_as_token_indices -> triton sparse kernel + thd -> select_token_indices_thd -> triton sparse kernel + + thd needs its own path because the causal prefix is per-document (``cu_seqlens``) + rather than ``(arange(s) + 1) // R``, and because TE rejects an ``arbitrary`` + mask once ``qkv_format == 'thd'``. Indices also scale as O(s*K) instead of the + mask's O(s^2), which is what makes very long sequences feasible. + """ + def __init__(self, config): super().__init__() self.config = config @@ -37,11 +84,16 @@ def __init__(self, config): self.compress_ratio = config.indexer_compress_ratio self.token_budget = config.indexer_budget self.block_topk = self.token_budget // self.compress_ratio - # Replicated projection (reference uses ReplicatedLinear). - self.index_qk_proj = nn.Linear( - config.hidden_size, (self.index_n_heads + self.index_kv_heads) * self.index_head_dim, + self.index_qk_proj = TELinear( + input_size=config.hidden_size, + output_size=(self.index_n_heads + self.index_kv_heads) * self.index_head_dim, + parallel_mode='duplicated', + config=config, + init_method=config.init_method, bias=False, - dtype=config.params_dtype) + skip_bias_add=True, + skip_weight_param_allocation=False, + ) self.q_layernorm = Qwen4ExpTextRMSNorm( self.index_head_dim, eps=config.layernorm_epsilon, @@ -55,11 +107,15 @@ def __init__(self, config): setattr(self.index_qk_proj.weight, 'sequence_parallel', config.sequence_parallel) def forward(self, *args, **kwargs): - raise RuntimeError('QSAIndexer performs selection via `select_mask`, not `forward`.') + raise RuntimeError('QSAIndexer selects via selection_as_mask / selection_as_token_indices / ' + 'select_token_indices_thd, not forward.') @torch.no_grad() - def select_mask(self, hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: - """Return the QSA selection as a bool mask, or ``None`` when it is a no-op. + def _score_and_topk_blocks(self, hidden_states: torch.Tensor, freqs: torch.Tensor): + """The selection itself: score every (query, block) pair and keep the top-k. + + Shared by ``selection_as_mask`` / ``selection_as_token_indices`` so both + encodings describe the same choice. Args: hidden_states: ``[s, b, h]`` (mcore layout), pre-attention input -- @@ -70,28 +126,21 @@ def select_mask(self, hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch keeping the indexer's RoPE identical to the attention's. Returns: - ``[b, 1, s, s]`` bool mask where True marks a *masked-out* key (TE's - ``arbitrary`` convention), or ``None`` if every visible key is - selected -- in which case the caller keeps the plain causal path and - pays nothing. - - Only the causal, unpacked layout is handled; callers must not invoke this - for packed/THD or context-parallel inputs (see the guard in the layer). + ``None`` when selection is a no-op (the causal prefix never exceeds + the budget), else ``(top_blocks, keep, n_blocks)`` where + ``top_blocks``/``keep`` are ``[b, s, k]`` (``k = + min(block_topk, max_blocks)``) and ``n_blocks`` is ``[s]``. """ s, b, _ = hidden_states.shape R = self.compress_ratio max_blocks = s // R - # Selection is a no-op while the causal prefix never exceeds the budget: - # `topk(min(block_topk, num_blocks))` then keeps every block and the tail - # re-adds the remainder, so the mask would be exactly causal. Skipping it - # keeps short sequences on TE's fused causal kernel. if max_blocks <= self.block_topk: return None device = hidden_states.device # ---- project to indexer q/k ---- # [s, b, h] -> [s, b, (nh + nkv) * d] - qk = self.index_qk_proj(hidden_states) + qk = self.index_qk_proj(hidden_states)[0] q, token_k = torch.split( qk, [self.index_n_heads * self.index_head_dim, self.index_kv_heads * self.index_head_dim], dim=-1) # -> [b, s, nh, d] / [b, s, d] @@ -100,11 +149,10 @@ def select_mask(self, hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch q = self.q_layernorm(q) # ---- materialize cos/sin from mcore freqs ---- - # freqs: [s, 1, 1, rot_dim] -> [s, rot_dim]; mscale mirrors the attention path. - mscale = getattr(self.config, 'attention_scaling', 1.0) or 1.0 - f = freqs.reshape(freqs.shape[0], -1)[:s] - cos = (torch.cos(f) * mscale).to(q.dtype) - sin = (torch.sin(f) * mscale).to(q.dtype) + # freqs is [s, freq_b, 1, rot]; mrope makes dim 1 the real batch + # (rope_utils.py:344), and _materialize_rope keeps it so it broadcasts. + mscale = self.config.attention_scaling + cos, sin = _materialize_rope(freqs, s, q.dtype, mscale) rot = cos.shape[-1] def apply_rope(t, cos_, sin_): @@ -112,8 +160,8 @@ def apply_rope(t, cos_, sin_): t_rope = (t_rope * cos_) + (_rotate_half(t_rope) * sin_) return torch.cat((t_rope, t_pass), dim=-1) - # queries rotate at their own position: cos [s, rot] -> [1, s, 1, rot] - q = apply_rope(q, cos[None, :, None, :], sin[None, :, None, :]) + # queries rotate at their own position: cos [bf, s, rot] -> [bf, s, 1, rot] + q = apply_rope(q, cos.unsqueeze(2), sin.unsqueeze(2)) # ---- pool every block once (shared across queries) ---- usable = max_blocks * R @@ -122,7 +170,7 @@ def apply_rope(t, cos_, sin_): pooled = self.k_layernorm(pooled) # blocks rotate at their first token's position starts = torch.arange(max_blocks, device=device) * R - block_keys = apply_rope(pooled, cos[starts][None], sin[starts][None]) # [b, nb, d] + block_keys = apply_rope(pooled, cos[:, starts], sin[:, starts]) # [b, nb, d] # ---- score all (query, block) pairs ---- scores = torch.einsum('bqhd,bkd->bqhk', q.float(), block_keys.float()) @@ -133,10 +181,34 @@ def apply_rope(t, cos_, sin_): block_ids = torch.arange(max_blocks, device=device) scores = scores.masked_fill((block_ids[None, :] >= n_blocks[:, None])[None], float('-inf')) - # ---- top-k blocks -> token mask ---- k = min(self.block_topk, max_blocks) top_blocks = scores.topk(k, dim=-1).indices # [b, s, k] keep = top_blocks < n_blocks[None, :, None] # drop the -inf padding slots + return top_blocks, keep, n_blocks + + @torch.no_grad() + def selection_as_mask(self, hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Encode the selection as a bool mask, or ``None`` when it is a no-op. + + Returns: + ``[b, 1, s, s]`` bool mask where True marks a *masked-out* key (TE's + ``arbitrary`` convention), or ``None`` if every visible key is + selected -- in which case the caller keeps the plain causal path and + pays nothing. + + Only the causal, unpacked layout is handled; callers must not invoke this + for packed/THD or context-parallel inputs (see the guard in the layer). + """ + core = self._score_and_topk_blocks(hidden_states, freqs) + # when seq lengths less than budget + if core is None: + return None + top_blocks, keep, n_blocks = core + s, b, _ = hidden_states.shape + R = self.compress_ratio + device = hidden_states.device + + # ---- top-k blocks -> token mask ---- tok = (top_blocks.unsqueeze(-1) * R + torch.arange(R, device=device)).flatten(-2) keep_tok = keep.unsqueeze(-1).expand(-1, -1, -1, R).flatten(-2) @@ -152,3 +224,151 @@ def apply_rope(t, cos_, sin_): # silently attending to the future) allowed &= (pos[None, :] <= pos[:, None])[None] return ~allowed.unsqueeze(1) # True == masked out + + @torch.no_grad() + def selection_as_token_indices(self, hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Encode the selection as token indices for the sparse kernel. + + Returns: + ``[b, s, K]`` int64 (``K = block_topk*R + R``) with ``-1`` marking + unused slots, or ``None`` when selection is a no-op. Indices live in + per-sample sequence space, matching ``qsa_sparse_attention``'s sbhd + branch. Top-block tokens and the query's own partial-block tail are + disjoint, so the kernel attends each selected key exactly once. + + Same layout constraints as ``selection_as_mask`` (causal, unpacked). + """ + core = self._score_and_topk_blocks(hidden_states, freqs) + if core is None: + return None + top_blocks, keep, n_blocks = core + s, b, _ = hidden_states.shape + R = self.compress_ratio + device = hidden_states.device + k = top_blocks.shape[-1] + arange_r = torch.arange(R, device=device) + + # top-k blocks expanded to tokens: [b, s, k, R] -> [b, s, k*R]; dropped + # (-inf) slots carry -1 so the kernel skips them. + tok = top_blocks.unsqueeze(-1) * R + arange_r + top_idx = tok.flatten(-2) + top_keep = keep.unsqueeze(-1).expand(-1, -1, -1, R).reshape(b, s, k * R) + top_idx = torch.where(top_keep, top_idx, top_idx.new_full((), -1)) + + # tail: the query's own partial block, causally truncated: [b, s, R] + pos = torch.arange(s, device=device) + tail_idx = (n_blocks * R)[:, None] + arange_r[None, :] # [s, R] + tail_idx = torch.where(tail_idx <= pos[:, None], tail_idx, tail_idx.new_full((), -1)) + tail_idx = tail_idx[None].expand(b, s, R) + + return torch.cat([top_idx, tail_idx], dim=-1).to(torch.int64) + + @torch.no_grad() + def select_token_indices_thd(self, hidden_tok: torch.Tensor, freqs: torch.Tensor, + cu_seqlens: torch.Tensor) -> torch.Tensor: + """QSA selection for packed (thd) inputs, indices in pack space. + + A standalone implementation: it does *not* call + ``_score_and_topk_blocks``, because under packing the causal prefix is + per-document (derived from ``cu_seqlens``) rather than a single + ``(arange(s) + 1) // R``, and the batch dim is already flattened into ``T``. + + Args: + hidden_tok: ``[T, h]`` packed tokens (the dummy batch dim squeezed out). + freqs: rotary angles ``[T, ...]``; each row already encodes that token's + *in-document* position (mcore builds them from per-doc position ids + under padding_free), so they are reused as-is here. + cu_seqlens: ``[D+1]`` int document boundaries (``cu[0] == 0``, + ``cu[-1] == T``). + + Returns: + ``[T, K]`` int64 pack-space indices (``K = block_topk*R + R``), ``-1`` + unused; every index stays inside its query's document and causal + prefix. ``None`` when selection is a no-op for every document, in which + case TE's packed causal kernel reproduces the selection exactly. + """ + T, _ = hidden_tok.shape + R = self.compress_ratio + device = hidden_tok.device + doc_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).long() # [D] + D = doc_lens.numel() + full_blocks = doc_lens // R # complete R-blocks per document + # No-op when no document's causal prefix can exceed the budget. + if int(full_blocks.max().item()) <= self.block_topk: + return None + + # ---- per-token document id / in-doc position ---- + token_doc = torch.repeat_interleave(torch.arange(D, device=device), doc_lens) # [T] + pos_in_doc = torch.arange(T, device=device) - cu_seqlens[token_doc].long() # [T] + + # ---- project to indexer q/k ---- + qk = self.index_qk_proj(hidden_tok)[0] + q, token_k = torch.split( + qk, [self.index_n_heads * self.index_head_dim, self.index_kv_heads * self.index_head_dim], dim=-1) + q = q.view(T, self.index_n_heads, self.index_head_dim) + raw_keys = token_k.view(T, self.index_kv_heads, self.index_head_dim).squeeze(1) # [T, d] + q = self.q_layernorm(q) + + # ---- rope on q at its own (in-doc) position ---- + # Deliberately not _materialize_rope: thd has no batch dim to preserve + # (samples are already flattened into T), so the trailing dims are folded + # into rot here instead of being kept separate. + mscale = self.config.attention_scaling + f = freqs.reshape(freqs.shape[0], -1)[:T] + cos = (torch.cos(f) * mscale).to(q.dtype) + sin = (torch.sin(f) * mscale).to(q.dtype) + rot = cos.shape[-1] + + def apply_rope(t, cos_, sin_): + t_rope, t_pass = t[..., :rot], t[..., rot:] + t_rope = (t_rope * cos_) + (_rotate_half(t_rope) * sin_) + return torch.cat((t_rope, t_pass), dim=-1) + + q = apply_rope(q, cos.unsqueeze(1), sin.unsqueeze(1)) # [T, nh, d] + + # ---- pool complete blocks in pack space (shared across queries) ---- + NB = int(full_blocks.sum().item()) + token_in_full = pos_in_doc < (full_blocks * R)[token_doc] # [T] + block_in_doc = pos_in_doc // R # [T] + block_offset = torch.cumsum(full_blocks, 0) - full_blocks # exclusive [D] + global_block = block_offset[token_doc] + block_in_doc # [T] + + pooled_sum = torch.zeros(NB, self.index_head_dim, device=device, dtype=torch.float32) + gb = global_block[token_in_full] + pooled_sum.index_add_(0, gb, raw_keys[token_in_full].float()) + pooled = (pooled_sum / R).to(raw_keys.dtype) # exactly R tokens per full block + pooled = self.k_layernorm(pooled) # [NB, d] + + # blocks rotate at their first token's in-doc position + block_doc = torch.repeat_interleave(torch.arange(D, device=device), full_blocks) # [NB] + block_in_doc_idx = torch.arange(NB, device=device) - block_offset[block_doc] # [NB] + first_pack = cu_seqlens[block_doc].long() + block_in_doc_idx * R # [NB] + block_keys = apply_rope(pooled, cos[first_pack], sin[first_pack]) # [NB, d] + + # ---- score every (token, block) pair ---- + scores = torch.einsum('thd,kd->thk', q.float(), block_keys.float()) + scores = torch.relu(scores).sum(dim=1) / math.sqrt(self.index_head_dim) # [T, NB] + + # ---- restrict to same-document, causally-before blocks ---- + q_nblocks = (pos_in_doc + 1) // R # [T] + valid = (block_doc[None, :] == token_doc[:, None]) & \ + (block_in_doc_idx[None, :] < q_nblocks[:, None]) # [T, NB] + scores = scores.masked_fill(~valid, float('-inf')) + + # ---- top-k blocks -> token indices ---- + k = min(self.block_topk, NB) + top_blocks = scores.topk(k, dim=-1).indices # [T, k] into [0, NB) + keep = valid.gather(1, top_blocks) # [T, k] + arange_r = torch.arange(R, device=device) + base = cu_seqlens[block_doc[top_blocks]].long() + block_in_doc_idx[top_blocks] * R # [T, k] + top_idx = (base.unsqueeze(-1) + arange_r).flatten(-2) # [T, k*R] + top_keep = keep.unsqueeze(-1).expand(-1, -1, R).reshape(T, k * R) + top_idx = torch.where(top_keep, top_idx, top_idx.new_full((), -1)) + + # ---- tail: the query's own partial block, causally truncated ---- + tail_base = cu_seqlens[token_doc].long() + q_nblocks * R # [T] + tail_idx = tail_base.unsqueeze(-1) + arange_r[None, :] # [T, R] + token_pos = torch.arange(T, device=device) + tail_idx = torch.where(tail_idx <= token_pos[:, None], tail_idx, tail_idx.new_full((), -1)) + + return torch.cat([top_idx, tail_idx], dim=-1).to(torch.int64) diff --git a/tests/test_qsa_indexer.py b/tests/test_qsa_indexer.py new file mode 100644 index 0000000..22f9b3c --- /dev/null +++ b/tests/test_qsa_indexer.py @@ -0,0 +1,67 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch + +from mcore_bridge.model.modules.qsa_indexer import _materialize_rope, _rotate_half + + +def test_materialize_rope_preserves_mrope_batch_dimension(): + """MRoPE positions differ per sample, so ``freq_b`` must stay its own dim. + + The bug this pins: flattening ``[s, b, 1, rot]`` to ``[s, b * rot]`` makes + ``rot`` come out as ``b * rot``, and every downstream ``[..., :rot]`` slice + then reads the wrong half -- silently, and only when ``b > 1``. + """ + seq_len, batch_size, rope_dim = 16, 2, 64 + freqs = torch.randn(seq_len, batch_size, 1, rope_dim) + + cos, sin = _materialize_rope(freqs, seq_len, torch.float32, 1.0) + expected = freqs.squeeze(2).permute(1, 0, 2) + + # rot must stay rope_dim, not batch_size * rope_dim + assert cos.shape == (batch_size, seq_len, rope_dim), \ + f'expected [b, s, rot] = {(batch_size, seq_len, rope_dim)}, got {tuple(cos.shape)}' + torch.testing.assert_close(cos, expected.cos()) + torch.testing.assert_close(sin, expected.sin()) + + # Perturbing sample 1 must leave sample 0 untouched: no cross-sample bleed. + changed_freqs = freqs.clone() + changed_freqs[:, 1].add_(0.5) + changed_cos, changed_sin = _materialize_rope(changed_freqs, seq_len, torch.float32, 1.0) + + torch.testing.assert_close(changed_cos[0], cos[0]) + torch.testing.assert_close(changed_sin[0], sin[0]) + assert not torch.equal(changed_cos[1], cos[1]), 'sample 1 should have changed' + assert not torch.equal(changed_sin[1], sin[1]), 'sample 1 should have changed' + + +def test_materialize_rope_applies_mscale_and_dtype(): + """``mscale`` mirrors the attention path's attention_scaling; dtype is honoured.""" + freqs = torch.randn(8, 1, 1, 32) + mscale = 1.7 + + cos, sin = _materialize_rope(freqs, 8, torch.bfloat16, mscale) + ref = freqs.squeeze(2).permute(1, 0, 2) + + assert cos.dtype is torch.bfloat16 and sin.dtype is torch.bfloat16 + torch.testing.assert_close(cos.float(), (ref.cos() * mscale).bfloat16().float()) + torch.testing.assert_close(sin.float(), (ref.sin() * mscale).bfloat16().float()) + + +def test_materialize_rope_truncates_to_seq_len(): + """A longer freq table is sliced to ``seq_len`` (CP hands over full-length tables).""" + freqs = torch.randn(64, 3, 1, 16) + + cos, _ = _materialize_rope(freqs, 20, torch.float32, 1.0) + + assert cos.shape == (3, 20, 16) + torch.testing.assert_close(cos, freqs[:20].squeeze(2).permute(1, 0, 2).cos()) + + +def test_rotate_half_matches_reference(): + """``_rotate_half`` is the standard (-x2, x1) split the attention path uses.""" + x = torch.randn(2, 5, 8) + + got = _rotate_half(x) + + x1, x2 = x[..., :4], x[..., 4:] + torch.testing.assert_close(got, torch.cat((-x2, x1), dim=-1)) diff --git a/tests/test_qwen4_exp_units.py b/tests/test_qwen4_exp_units.py new file mode 100644 index 0000000..c7166b0 --- /dev/null +++ b/tests/test_qwen4_exp_units.py @@ -0,0 +1,508 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import math +import pytest +import torch +import torch.nn.functional as F +from torch import nn + +# Every test here builds a real TELinear/triton kernel, so a GPU is required. +# Skip (rather than fail) on CPU-only machines so the file is safe to collect in CI. +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason='qwen4_exp unit tests require a GPU') + +# bf16 rounding budget. The gated-residual chain is compiled, which reassociates +# the mean-reduction in fp32 and drifts by ~1 ULP on the output and ~2 ULP on the +# input grad (relative L1 ~2e-3), and the exact value is not reproducible between +# runs. 4 ULP leaves headroom for that without hiding a real regression. +_BF16_ULP = 2**-7 +_TOL = 4 * _BF16_ULP + + +def _make_config(hidden_size=256, + hc_count=4, + hc_lowrank=64, + index_n_heads=4, + index_kv_heads=1, + index_head_dim=64, + compress_ratio=4, + budget=64, + dtype=torch.float32): + from megatron.core.transformer.transformer_config import TransformerConfig + cfg = TransformerConfig(num_layers=1, hidden_size=hidden_size, num_attention_heads=8, params_dtype=dtype) + cfg.hc_count = hc_count + cfg.hc_lowrank = hc_lowrank + cfg.layernorm_epsilon = 1e-6 + cfg.indexer_n_heads = index_n_heads + cfg.indexer_kv_heads = index_kv_heads + cfg.indexer_head_dim = index_head_dim + cfg.indexer_compress_ratio = compress_ratio + cfg.indexer_budget = budget + cfg.attention_scaling = 1.0 + return cfg + + +# -------------------------------------------------------------------------- +# 1. QSA indexer selection vs the transformers reference loop +# -------------------------------------------------------------------------- +def _rotate_half(x): + x1 = x[..., :x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2:] + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rope(q, cos, sin, unsqueeze_dim=1): + """Verbatim from transformers modeling_qwen4_exp.apply_rotary_pos_emb (k=None branch).""" + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + rotary_dim = cos.shape[-1] + q_rope, q_nope = q[..., :rotary_dim], q[..., rotary_dim:] + q_rope = (q_rope * cos) + (_rotate_half(q_rope) * sin) + return torch.cat([q_rope, q_nope], dim=-1) + + +def _reference_select(indexer, hidden_states, cos, sin): + """The reference double loop: per query, re-pool blocks and take top-k. + + Mirrors transformers Qwen4ExpTextQSAIndexer's selection (modeling_qwen4_exp.py + ~L667-702): only blocks fully inside the causal prefix are scored, the top + `block_topk` are expanded back to token ids, and the trailing tokens that do + not fill a whole block ("tail") are always visible. + """ + s = hidden_states.shape[0] + R = indexer.compress_ratio + topk = indexer.block_topk + d = indexer.index_head_dim + nh, nkv = indexer.index_n_heads, indexer.index_kv_heads + + qk = indexer.index_qk_proj(hidden_states) + if isinstance(qk, tuple): + qk = qk[0] + q, token_k = torch.split(qk, [nh * d, nkv * d], dim=-1) + # Mirrors the reference exactly (modeling_qwen4_exp.py ~L645-651): + # * q goes through q_layernorm AND RoPE at its own position + # * raw_keys stays RAW here -- k_layernorm is applied once, later, to the + # pooled block keys only. Norming token_k up front (as an earlier version + # of this test did) applies the norm twice and silently changes the scores. + q = indexer.q_layernorm(q.view(s, nh, d)) + q = _apply_rope(q.unsqueeze(0), cos.unsqueeze(0), sin.unsqueeze(0), unsqueeze_dim=2)[0] + raw_keys = token_k.view(s, d) + + allowed = torch.zeros(s, s, dtype=torch.bool, device=hidden_states.device) + for i in range(s): + n_blocks = (i + 1) // R + if n_blocks > 0: + pooled = raw_keys[:n_blocks * R].view(n_blocks, R, d).float().mean(dim=1).to(raw_keys.dtype) + pooled = indexer.k_layernorm(pooled) + starts = torch.arange(n_blocks, device=hidden_states.device) * R + bk = _apply_rope(pooled.unsqueeze(1), cos[starts], sin[starts]).squeeze(1) + score = torch.relu(q[i].float() @ bk.float().T).sum(dim=0) / math.sqrt(d) + keep = score.topk(min(topk, n_blocks), dim=-1).indices + for b in keep.tolist(): + allowed[i, b * R:(b + 1) * R] = True + # tail: visible tokens that do not fill a complete block + allowed[i, n_blocks * R:i + 1] = True + return ~allowed[None, None] + + +def test_indexer_matches_reference(): + from mcore_bridge.model.modules.qsa_indexer import QSAIndexer + print('\n[1] QSAIndexer.selection_as_mask vs transformers reference loop') + ok = True + # Shapes chosen so the selection is genuinely sparse: with budget/ratio = + # block_topk, sparsity needs num_blocks > block_topk, i.e. s > budget+ratio-1. + # A case at/below the budget is included too, where the model's own math makes + # every block selected (selection_as_mask must then report a no-op). + cases = [ + # (seq_len, compress_ratio, budget, seed) -- includes non-multiples of the + # ratio to exercise the tail, and two ratios to catch hardcoded 4s. + (200, 4, 64, 0), + (201, 4, 64, 1), + (203, 4, 64, 2), + (128, 2, 32, 3), + (130, 8, 64, 4), + ] + for s, ratio, budget, seed in cases: + torch.manual_seed(seed) + cfg = _make_config(compress_ratio=ratio, budget=budget) + idx = QSAIndexer(cfg).cuda() + with torch.no_grad(): + idx.index_qk_proj.weight.normal_(0, 0.02) + idx.q_layernorm.weight.normal_(0, 0.02) + idx.k_layernorm.weight.normal_(0, 0.02) + + hs = torch.randn(s, 1, cfg.hidden_size, device='cuda') + d = cfg.indexer_head_dim + # mcore stores rope angles; selection_as_mask materializes cos/sin from them, so + # feed angles here and derive the same cos/sin for the reference. + freqs = torch.randn(s, 1, 1, d, device='cuda') + cos, sin = freqs[:, 0, 0].cos(), freqs[:, 0, 0].sin() + + got = idx.selection_as_mask(hs, freqs) + n_blocks = s // ratio + expect_noop = n_blocks <= budget // ratio + if expect_noop: + good = got is None + print(f' s={s:4d} ratio={ratio} budget={budget}: below-budget no-op -> ' + f"{'None as expected' if good else 'UNEXPECTED mask'}") + else: + ref = _reference_select(idx, hs, cos, sin) + good = got is not None and torch.equal(got, ref) + if got is None: + print(f' s={s:4d} ratio={ratio} budget={budget}: MISMATCH (got None, ' + 'expected a sparse mask)') + else: + nsel = int((~ref).sum()) + print(f' s={s:4d} ratio={ratio} budget={budget}: ' + f"blocks={n_blocks} selected_keys={nsel} -> {'MATCH' if good else 'MISMATCH'}") + ok &= good + assert ok, 'numerical check failed; see printed relL1/diff above' + + +def test_indexer_is_forward_only(): + """The indexer participates in the forward but must take no gradient. + + The selection is a discrete top-k turned into a bool mask, so the backbone + still differentiates normally while index_qk_proj / q_layernorm / + k_layernorm stay frozen. If a refactor ever makes the mask carry grad, the + "forward-only" design claim silently stops holding. + """ + from mcore_bridge.model.modules.qsa_indexer import QSAIndexer + print('\n[2] indexer is forward-only (mask is a constant, backbone still differentiates)') + torch.manual_seed(0) + cfg = _make_config(compress_ratio=4, budget=64) + idx = QSAIndexer(cfg).cuda() + with torch.no_grad(): + idx.index_qk_proj.weight.normal_(0, 0.02) + + s = 200 + hs = torch.randn(s, 1, cfg.hidden_size, device='cuda', requires_grad=True) + freqs = torch.randn(s, 1, 1, cfg.indexer_head_dim, device='cuda') + mask = idx.selection_as_mask(hs, freqs) + + mask_is_const = mask is not None and mask.dtype == torch.bool and not mask.requires_grad + # a toy attention consuming the mask -- the backbone must still get grads + q = hs.squeeze(1) + logits = (q @ q.T).masked_fill(mask[0, 0], float('-inf')) + logits.softmax(-1).sum().backward() + backbone_ok = hs.grad is not None and torch.isfinite(hs.grad).all() and hs.grad.abs().sum() > 0 + indexer_ok = all(p.grad is None for p in idx.parameters()) + + print(f' mask is non-differentiable constant: {mask_is_const} (dtype={mask.dtype})') + print(f' backbone received gradient: {backbone_ok}') + print(f' indexer params still have no grad: {indexer_ok}') + assert mask_is_const, f'QSA mask must be a non-differentiable bool constant (got {mask.dtype})' + assert backbone_ok, 'backbone did not receive a finite non-zero gradient through the mask' + assert indexer_ok, 'indexer parameters received gradient; the forward-only design no longer holds' + + +# -------------------------------------------------------------------------- +# 2. gated hyper-connection numerical equivalence +# -------------------------------------------------------------------------- +def _gated_residual_eager(m, hyper_input): + """The gated residual written out plainly, with no compiled helpers.""" + normed = m.hc_norm(hyper_input) + w = F.silu(m.input_mix_weight_down(normed)[0] / m.hc_count) + w = torch.sigmoid(m.input_mix_weight_up(w)[0]).unflatten(-1, (m.hc_count, m.hidden_size)) + mixed = (w * normed.unflatten(-1, (m.hc_count, m.hidden_size))).mean(dim=-2) + if m.block_inject_weight is None: + return mixed, None + inj = 2 * torch.sigmoid(m.block_inject_weight(normed)[0] / m.hc_count) + return mixed, inj + + +def test_gated_residual_equivalence(): + from mcore_bridge.model.modules.hyper_connection_gated import Qwen4ExpTextGatedResidual + print(f'\n[3] Qwen4ExpTextGatedResidual vs plain eager (tol {_TOL:.4g} = 4 bf16 ULP)') + ok = True + for use_combine in (True, False): + for s in (64, 256): + torch.manual_seed(7) + cfg = _make_config(dtype=torch.bfloat16) + m = Qwen4ExpTextGatedResidual(cfg, use_combine=use_combine).cuda() + x = torch.randn(s, 1, cfg.hc_count * cfg.hidden_size, device='cuda', dtype=torch.bfloat16) + xa = x.clone().requires_grad_(True) + xb = x.clone().requires_grad_(True) + + out_a = m(xa) + got = out_a[0] if isinstance(out_a, tuple) else out_a + ref, ref_inj = _gated_residual_eager(m, xb) + + grad_seed = torch.randn_like(got) + got.backward(grad_seed) + ref.backward(grad_seed) + + fwd_d = (got.float() - ref.float()).abs().max().item() + grad_d = (xa.grad.float() - xb.grad.float()).abs().max().item() + inj_d = 0.0 + if isinstance(out_a, tuple): + inj_d = (out_a[2].float() - ref_inj.float()).abs().max().item() + finite = bool(torch.isfinite(got).all() and torch.isfinite(xa.grad).all()) + good = fwd_d <= _TOL and grad_d <= _TOL and inj_d <= _TOL and finite + ok &= good + print(f' combine={use_combine!s:5s} s={s:4d}: fwd={fwd_d:.6g} grad={grad_d:.6g} ' + f"inj={inj_d:.6g} finite={finite} -> {'OK' if good else 'FAIL'}") + assert ok, 'numerical check failed; see printed relL1/diff above' + + +# -------------------------------------------------------------------------- +# 3. host-offload table lookup +# -------------------------------------------------------------------------- +def test_host_lookup_reference(): + """Single-rank host lookup must equal a plain gather of the same rows.""" + from types import SimpleNamespace + + from mcore_bridge.model.modules.ple import Qwen4ExpTextNGramEmbedding + print('\n[4] host-offload lookup vs plain gather (single rank)') + torch.manual_seed(11) + total, head_dim = 1024, 16 + host_table = torch.randn(total, head_dim) + dummy = SimpleNamespace(vocab_start=0, vocab_end=total, host_table=host_table, _tp_size=1, _tp_group=None) + ngram_ids = torch.randint(0, total, (2, 5, 3)) + got = Qwen4ExpTextNGramEmbedding._host_lookup(dummy, ngram_ids) + ref = host_table[ngram_ids].flatten(-2) + d = (got.float() - ref.float()).abs().max().item() + ok = d <= 1e-5 + print(f' max diff = {d:.3e} -> {"OK" if ok else "FAIL"}') + assert ok, 'numerical check failed; see printed relL1/diff above' + + +def test_host_lookup_tp_partition(): + """Each TP rank gathers its own slice; the pieces reassemble the full gather.""" + from types import SimpleNamespace + + from mcore_bridge.model.modules.ple import Qwen4ExpTextNGramEmbedding + print('\n[5] host-offload TP=2 partition reassembly') + total, head_dim = 1024, 16 + torch.manual_seed(13) + full = torch.randn(total, head_dim) + ngram_ids = torch.randint(0, total, (2, 5, 3)) + acc = None + for (vs, ve) in [(0, 512), (512, 1024)]: + dummy = SimpleNamespace( + vocab_start=vs, vocab_end=ve, host_table=full[vs:ve].contiguous(), _tp_size=1, _tp_group=None) + part = Qwen4ExpTextNGramEmbedding._host_lookup(dummy, ngram_ids) + acc = part if acc is None else acc + part + ref = full[ngram_ids].flatten(-2) + d = (acc.float() - ref.float()).abs().max().item() + ok = d <= 1e-4 + print(f' max diff = {d:.3e} -> {"OK" if ok else "FAIL"}') + assert ok, 'numerical check failed; see printed relL1/diff above' + + +# -------------------------------------------------------------------------- +# 4. PLE fused gate+conv triton kernel +# -------------------------------------------------------------------------- +def _fp32_chain(key, query, value, wk, wq, wc, conv_w, n, C, eps, dilation, seq_len): + """The fused chain in pure fp32 (what the kernel numerically implements): + grouped zero-centered RMSNorms, gate transform, sigmoid*value, norm_conv, + causal dilated depthwise conv + SiLU + residual; one cast to bf16 at the end.""" + import torch.nn.functional as F + T = key.shape[0] + k = key.view(T, n, C).float() + q = query.view(T, n, C).float() + rk = torch.rsqrt(k.pow(2).mean(-1, keepdim=True) + eps) + rq = torch.rsqrt(q.pow(2).mean(-1, keepdim=True) + eps) + kn = k * rk * (1 + wk.view(n, C).float()) + qn = q * rq * (1 + wq.view(n, C).float()) + score = (kn * qn).sum(-1, keepdim=True) / math.sqrt(C) + u = score.abs().clamp_min(1e-6).sqrt() * score.sign() + gated = (torch.sigmoid(u) * value.view(T, 1, C).float()).view(T, n * C) + g = gated.view(T, n, C) + rc = torch.rsqrt(g.pow(2).mean(-1, keepdim=True) + eps) + normed = (g * rc * (1 + wc.view(n, C).float())).view(T, n * C) + rows = T // seq_len + x = normed.view(rows, seq_len, n * C).transpose(1, 2) + K = conv_w.shape[-1] + pad = (K - 1) * dilation + x = F.pad(x, (pad, 0))[..., -(pad + seq_len):] + conv = F.silu(F.conv1d(x, conv_w.float(), dilation=dilation, groups=n * C)).transpose(1, 2) + out = gated.view(rows, seq_len, n * C) + conv + return out.view(T, n * C) + + +def test_ple_fused_kernel(): + """The triton gate+conv chain vs its fp32 formulation, fwd+bwd. + + bf16 inputs (as in training), fp32 accumulation inside the kernels, one + dtype cast at the output -- so the fp32 reference is the strict contract; + against the module's bf16-intermediate torch path the kernel drifts by + design (see ple_fused_parity.py for the measured distribution).""" + from mcore_bridge.model.modules.kernels.ple_kernels import HAVE_TRITON, ple_gate_conv_triton + print('\n[6] PLE fused gate+conv kernel vs fp32 reference (fwd+bwd)') + if not HAVE_TRITON: + print(' triton unavailable -> SKIP (treated as pass)') + return True + torch.manual_seed(17) + rows, L, n, C, K, dilation = 2, 33, 4, 128, 4, 3 + T = rows * L + dev = 'cuda' + key = torch.randn(T, n * C, device=dev, dtype=torch.bfloat16).requires_grad_(True) + query = torch.randn(T, n * C, device=dev, dtype=torch.bfloat16).requires_grad_(True) + value = torch.randn(T, C, device=dev, dtype=torch.bfloat16).requires_grad_(True) + wk = (torch.randn(n * C, device=dev, dtype=torch.bfloat16) * 0.1).requires_grad_(True) + wq = (torch.randn(n * C, device=dev, dtype=torch.bfloat16) * 0.1).requires_grad_(True) + wc = (torch.randn(n * C, device=dev, dtype=torch.bfloat16) * 0.1).requires_grad_(True) + conv_w = (torch.randn(n * C, 1, K, device=dev, dtype=torch.bfloat16) * 0.2).requires_grad_(True) + eps = 1e-6 + + out_k = ple_gate_conv_triton(query, key, value, wk, wq, wc, conv_w, n, eps, dilation, L) + assert out_k is not None, 'kernel path unavailable despite HAVE_TRITON' + g = torch.randn_like(out_k) + out_k.backward(g) + grads_k = {p: p.grad.float().clone() for p in (key, query, value, wk, wq, wc, conv_w)} + out_k = out_k.float().clone() + + for p in (key, query, value, wk, wq, wc, conv_w): + p.grad = None + out_r = _fp32_chain(key, query, value, wk, wq, wc, conv_w, n, C, eps, dilation, L).to(torch.bfloat16) + out_r.backward(g) + + ok = True + o_rel = ((out_k - out_r.float()).abs().sum() / out_r.float().abs().sum()).item() + good = o_rel <= 1e-2 + ok &= good + print(f' fwd relL1={o_rel:.3e} -> {"OK" if good else "FAIL"}') + names = { + id(key): 'key', + id(query): 'query', + id(value): 'value', + id(wk): 'wk', + id(wq): 'wq', + id(wc): 'wc', + id(conv_w): 'conv_w' + } + for p in (key, query, value, wk, wq, wc, conv_w): + a, b = grads_k[p], p.grad.float() + rel = ((a - b).abs().sum() / a.abs().sum().clamp_min(1e-12)).item() + good = rel <= 1e-2 + ok &= good + print(f' grad[{names[id(p)]:6s}] relL1={rel:.3e} -> {"OK" if good else "FAIL"}') + assert ok, 'numerical check failed; see printed relL1/diff above' + + +def test_indexer_indices_match_reference(): + """selection_as_token_indices (the sparse-kernel input) must select exactly the tokens the + transformers reference loop selects. A drift here means the kernel attends a + different set than the validated mask path -- invisible without this test.""" + from mcore_bridge.model.modules.qsa_indexer import QSAIndexer + print('\n[7] QSAIndexer.selection_as_token_indices == reference allowed set (sbhd)') + ok = True + cases = [ + (200, 4, 64, 0), + (201, 4, 64, 1), + (128, 2, 32, 3), + (130, 8, 64, 4), + (60, 4, 64, 5), # below budget -> no-op + ] + for s, ratio, budget, seed in cases: + torch.manual_seed(seed) + cfg = _make_config(compress_ratio=ratio, budget=budget) + idx = QSAIndexer(cfg).cuda() + with torch.no_grad(): + idx.index_qk_proj.weight.normal_(0, 0.02) + idx.q_layernorm.weight.normal_(0, 0.02) + idx.k_layernorm.weight.normal_(0, 0.02) + hs = torch.randn(s, 1, cfg.hidden_size, device='cuda') + d = cfg.indexer_head_dim + freqs = torch.randn(s, 1, 1, d, device='cuda') + cos, sin = freqs[:, 0, 0].cos(), freqs[:, 0, 0].sin() + indices = idx.selection_as_token_indices(hs, freqs) + n_blocks = s // ratio + if n_blocks <= budget // ratio: + good = indices is None + print(f' s={s:4d} ratio={ratio}: below-budget no-op -> ' + f"{'None as expected' if good else 'UNEXPECTED indices'}") + else: + ref_allowed = ~_reference_select(idx, hs, cos, sin)[0, 0] # [s, s] + allowed = torch.zeros(s, s, dtype=torch.bool, device='cuda') + row = indices[0] + qq, kk = torch.nonzero(row >= 0, as_tuple=True) + allowed[qq, row[qq, kk]] = True + good = bool(torch.equal(allowed, ref_allowed)) + print(f' s={s:4d} ratio={ratio}: allowed set == reference -> ' + f"{'MATCH' if good else 'MISMATCH'}") + ok &= good + assert ok, 'numerical check failed; see printed relL1/diff above' + + +def test_indexer_indices_packed(): + """select_token_indices_thd (thd) must reproduce running selection_as_token_indices on each + document independently, and never emit a cross-document or future index.""" + from mcore_bridge.model.modules.qsa_indexer import QSAIndexer + print('\n[8] QSAIndexer.select_token_indices_thd == per-doc selection_as_token_indices (thd)') + torch.manual_seed(5) + cfg = _make_config(compress_ratio=4, budget=32) # block_topk = 8 + idx = QSAIndexer(cfg).cuda() + with torch.no_grad(): + idx.index_qk_proj.weight.normal_(0, 0.02) + idx.q_layernorm.weight.normal_(0, 0.02) + idx.k_layernorm.weight.normal_(0, 0.02) + doc_lens = [48, 40, 12] # blocks 12/10/3 -> first two sparse, third causal + T = sum(doc_lens) + cu = torch.tensor([0] + list(torch.cumsum(torch.tensor(doc_lens), 0).tolist()), dtype=torch.long, device='cuda') + h, d = cfg.hidden_size, cfg.indexer_head_dim + hidden_tok = torch.randn(T, h, device='cuda') + pos_in_doc = torch.cat([torch.arange(x) for x in doc_lens]).cuda() + base = torch.randn(64, d, device='cuda') # freqs as a function of in-doc position + freqs = base[pos_in_doc].reshape(T, 1, 1, d) + packed = idx.select_token_indices_thd(hidden_tok, freqs, cu) + assert packed is not None, 'expected sparse to engage for packed docs' + ok = True + off = 0 + for L in doc_lens: + hid_d = hidden_tok[off:off + L].unsqueeze(1) # [L, 1, h] + freqs_d = base[:L].reshape(L, 1, 1, d) + ref = idx.selection_as_token_indices(hid_d, freqs_d) + if ref is None: # doc below budget -> full causal + pos = torch.arange(L, device='cuda') + ref_allowed = pos[None, :] <= pos[:, None] + else: + ref_allowed = torch.zeros(L, L, dtype=torch.bool, device='cuda') + rr = ref[0] + qq, kk = torch.nonzero(rr >= 0, as_tuple=True) + ref_allowed[qq, rr[qq, kk]] = True + seg = packed[off:off + L] + valid = seg >= 0 + in_doc = bool(((seg[valid] - off) >= 0).all() and ((seg[valid] - off) < L).all()) + got_allowed = torch.zeros(L, L, dtype=torch.bool, device='cuda') + qq, kk = torch.nonzero(valid, as_tuple=True) + got_allowed[qq, seg[qq, kk] - off] = True + good = in_doc and bool(torch.equal(ref_allowed, got_allowed)) + ok &= good + print(f' doc L={L:3d}: in-doc={in_doc} allowed==per-doc -> ' + f"{'MATCH' if good else 'MISMATCH'}") + off += L + assert ok, 'numerical check failed; see printed relL1/diff above' + + +def test_indexer_mrope_batch(): + """Selection must stay per-sample when freqs carry a real batch dim. + + mrope hands mcore a ``[s, b, 1, rot]`` freq tensor (rope_utils.py:344 keys off + ``freqs.shape[1] > 1``). Flattening that to ``[s, b*rot]`` folds batch into the + rotary feature dim, so ``rot`` becomes ``b*rot`` and every sample gets the wrong + angles. Every other indexer test runs b=1, where ``b*rot == rot`` hides it.""" + from mcore_bridge.model.modules.qsa_indexer import QSAIndexer + print('\n[9] QSAIndexer per-sample selection under mrope freqs (b=2)') + torch.manual_seed(11) + s, ratio, budget = 200, 4, 64 + cfg = _make_config(compress_ratio=ratio, budget=budget) + idx = QSAIndexer(cfg).cuda() + with torch.no_grad(): + idx.index_qk_proj.weight.normal_(0, 0.02) + idx.q_layernorm.weight.normal_(0, 0.02) + idx.k_layernorm.weight.normal_(0, 0.02) + d = cfg.indexer_head_dim + hs = torch.randn(s, 2, cfg.hidden_size, device='cuda') + # distinct per-sample angles: the whole point of mrope + freqs = torch.randn(s, 2, 1, d, device='cuda') + + both = idx.selection_as_token_indices(hs, freqs) + assert both is not None, 'expected an active selection for this length' + ok = True + for i in range(2): + # sample i alone must reproduce row i of the batched call + alone = idx.selection_as_token_indices(hs[:, i:i + 1], freqs[:, i:i + 1]) + same = torch.equal(alone[0], both[i]) + ok &= same + print(f' sample {i}: batched == standalone -> {"OK" if same else "FAIL"}') + assert ok, 'per-sample selection changed when batched (mrope batch dim leaked into rot)'