From d253e7d927e1a052843b1bba41f7490e8ebda372 Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Wed, 26 Aug 2026 17:23:59 +0800
Subject: [PATCH 01/10] wip

---
 src/mcore_bridge/config/model_config.py       | 20 ++++++++
 src/mcore_bridge/config/parser.py             | 51 +++++++++++++++++++
 src/mcore_bridge/model/constant.py            |  3 ++
 src/mcore_bridge/model/gpt_model.py           |  2 +-
 src/mcore_bridge/model/gpts/__init__.py       |  2 +-
 src/mcore_bridge/model/gpts/qwen3_next.py     |  3 +-
 src/mcore_bridge/model/mm_gpts/__init__.py    |  2 +-
 src/mcore_bridge/model/mm_gpts/qwen3_vl.py    | 14 +++--
 src/mcore_bridge/model/modules/__init__.py    |  3 ++
 .../model/modules/gated_delta_net.py          | 24 ++++++---
 .../model/modules/transformer_block.py        | 19 +++++++
 11 files changed, 126 insertions(+), 17 deletions(-)

diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py
index a4720038..c5231404 100644
--- a/src/mcore_bridge/config/model_config.py
+++ b/src/mcore_bridge/config/model_config.py
@@ -196,6 +196,26 @@ class ModelConfig(TransformerConfig):
     attention_output_gate: bool = False
     linear_decoupled_in_proj: bool = False
 
+    # qwen3.8-flash-next (HC + PLE + QSA)
+    hc_count: Optional[int] = None
+    hc_lowrank: Optional[int] = None
+    ple_layer_ids: Optional[List[int]] = None
+    ple_embed_dim: Optional[int] = None
+    ple_conv_kernel_size: Optional[int] = None
+    ngram_size: Optional[int] = None
+    heads_per_ngram: Optional[int] = None
+    ngram_vocab_size_base: Optional[int] = None
+    make_ngram_vocab_size_divisible_by: Optional[int] = None
+    split_ngram_parts: Optional[int] = None
+    ple_seed: Optional[int] = None
+    eos_token_id: Optional[int] = None
+    indexer_n_heads: Optional[int] = None
+    indexer_kv_heads: Optional[int] = None
+    indexer_head_dim: Optional[int] = None
+    indexer_budget: Optional[int] = None
+    indexer_compress_ratio: Optional[int] = None
+    output_gate_type: Optional[str] = None
+
     # nemotron_h (hybrid mamba2 + attention + moe)
     hybrid_layer_pattern: Optional[str] = None
 
diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py
index ebd30864..a2a8e818 100644
--- a/src/mcore_bridge/config/parser.py
+++ b/src/mcore_bridge/config/parser.py
@@ -52,6 +52,23 @@
     'linear_key_head_dim': ['linear_key_head_dim'],
     'linear_value_head_dim': ['linear_value_head_dim'],
     'linear_conv_kernel_dim': ['linear_conv_kernel_dim'],
+    # qwen3_8_flash_next
+    'hc_count': ['hc_count'],
+    'hc_lowrank': ['hc_lowrank'],
+    'ple_layer_ids': ['ple_layer_ids'],
+    'ple_embed_dim': ['ple_embed_dim'],
+    'ple_conv_kernel_size': ['ple_conv_kernel_size'],
+    'ngram_size': ['ngram_size'],
+    'heads_per_ngram': ['heads_per_ngram'],
+    'ngram_vocab_size_base': ['ngram_vocab_size_base'],
+    'make_ngram_vocab_size_divisible_by': ['make_ngram_vocab_size_divisible_by'],
+    'split_ngram_parts': ['split_ngram_parts'],
+    'indexer_n_heads': ['indexer_n_heads'],
+    'indexer_kv_heads': ['indexer_kv_heads'],
+    'indexer_head_dim': ['indexer_head_dim'],
+    'indexer_budget': ['indexer_budget'],
+    'indexer_compress_ratio': ['indexer_compress_ratio'],
+    'output_gate_type': ['output_gate_type'],
     # dsa
     'dsa_indexer_n_heads': ['index_n_heads'],
     'dsa_indexer_head_dim': ['index_head_dim'],
@@ -244,6 +261,40 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
         if use_mcore_gdn:
             res['experimental_attention_variant'] = 'gated_delta_net'
         res.setdefault('linear_attention_freq', 4)
+    # TODO: confirm model_type, remove one
+    elif hf_model_type in {'qwen3_8_flash_next', 'qwen4_exp'}:
+        use_mcore_gdn = get_env_args('USE_MCORE_GDN', bool, True)
+        res['layernorm_zero_centered_gamma'] = True
+        res['attention_output_gate'] = True
+        res['qk_layernorm'] = True
+        res['linear_decoupled_in_proj'] = True
+        res['moe_shared_expert_gate'] = True
+        if use_mcore_gdn:
+            res['experimental_attention_variant'] = 'gated_delta_net'
+        text_config = getattr(hf_config, 'text_config', hf_config)
+        num_layers = res['num_layers']
+        linear_pattern = ['1' if t == 'linear_attention' else '0' for t in layer_types]
+        res['linear_attention_freq'] = f"[{','.join(linear_pattern)}]"
+        if res.get('num_moe_experts'):
+            res['moe_layer_freq'] = f"[{','.join(['1'] * num_layers)}]"
+        # seed is hardcoded in transformers, not in config
+        res['ple_seed'] = int(getattr(text_config, 'seed', 1234))
+        eos_token_id = getattr(text_config, 'eos_token_id', None)
+        if eos_token_id is not None:
+            res['eos_token_id'] = eos_token_id
+        # These fields must come from the model config: ModelConfig carries no
+        # defaults for them, and a silently substituted value would corrupt the
+        # n-gram hash-table sharding/math at checkpoint conversion.
+        _required = [
+            'hc_count', 'hc_lowrank', 'ple_layer_ids', 'ple_embed_dim', 'ple_conv_kernel_size', 'ngram_size',
+            'heads_per_ngram', 'ngram_vocab_size_base', 'make_ngram_vocab_size_divisible_by', 'split_ngram_parts',
+            'eos_token_id', 'indexer_n_heads', 'indexer_kv_heads', 'indexer_head_dim', 'indexer_budget',
+            'indexer_compress_ratio'
+        ]
+        _missing = [k for k in _required if res.get(k) is None]
+        if _missing:
+            raise ValueError(f'qwen4_exp config is missing required fields: {_missing}. '
+                             'They must be provided by the model config.json.')
     elif llm_model_type == 'minimax_m2':
         res['add_qkv_bias'] = False
     elif llm_model_type == 'olmoe':
diff --git a/src/mcore_bridge/model/constant.py b/src/mcore_bridge/model/constant.py
index d05ad188..008f7da9 100644
--- a/src/mcore_bridge/model/constant.py
+++ b/src/mcore_bridge/model/constant.py
@@ -25,6 +25,9 @@ class MLLMModelType:
     qwen3_omni = 'qwen3_omni'
     qwen3_asr = 'qwen3_asr'
     qwen3_5 = 'qwen3_5'
+    # TODO: confirm and remove one 
+    qwen3_8_flash_next = 'qwen3_8_flash_next'
+    qwen4_exp = 'qwen4_exp'
     ovis2_5 = 'ovis2_5'
 
     internvl_chat = 'internvl_chat'
diff --git a/src/mcore_bridge/model/gpt_model.py b/src/mcore_bridge/model/gpt_model.py
index 8ed99943..5245718c 100644
--- a/src/mcore_bridge/model/gpt_model.py
+++ b/src/mcore_bridge/model/gpt_model.py
@@ -309,7 +309,7 @@ def forward(
                 padding_mask = torch.chunk(padding_mask, tp_size, dim=1)[mpu.get_tensor_model_parallel_rank()]
             extra_block_kwargs['padding_mask'] = padding_mask.contiguous()
 
-        if self.config.moe_n_hash_layers > 0:
+        if self.config.moe_n_hash_layers > 0 or getattr(self.config, 'ple_layer_ids', None):
             extra_block_kwargs['input_ids'] = input_ids
 
         # Run decoder.
diff --git a/src/mcore_bridge/model/gpts/__init__.py b/src/mcore_bridge/model/gpts/__init__.py
index 6b4c0343..0bc7b780 100644
--- a/src/mcore_bridge/model/gpts/__init__.py
+++ b/src/mcore_bridge/model/gpts/__init__.py
@@ -1,3 +1,3 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 from . import (bailing_hybrid, bailing_moe, deepseek_v4, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, nemotron_h, olmoe,
-               qwen3_emb, qwen3_next)
+               qwen3_8_flash_next, qwen3_emb, qwen3_next)
diff --git a/src/mcore_bridge/model/gpts/qwen3_next.py b/src/mcore_bridge/model/gpts/qwen3_next.py
index 0d4e2900..69df8f96 100644
--- a/src/mcore_bridge/model/gpts/qwen3_next.py
+++ b/src/mcore_bridge/model/gpts/qwen3_next.py
@@ -79,8 +79,7 @@ def __init__(self, config: ModelConfig, hidden_size: int, eps: float = 1e-5):
         super().__init__()
         self.config = config
         self.eps = eps
-        # Initialize weight to zeros (Zero-Centered), matching HuggingFace Qwen3NextRMSNorm
-        self.weight = torch.nn.Parameter(torch.zeros(hidden_size))
+        self.weight = torch.nn.Parameter(torch.zeros(hidden_size, dtype=config.params_dtype))
         # Mark weight for SP gradient AllReduce across TP domain (consistent with TENorm/MCoreRMSNorm)
         setattr(self.weight, 'sequence_parallel', config.sequence_parallel)
 
diff --git a/src/mcore_bridge/model/mm_gpts/__init__.py b/src/mcore_bridge/model/mm_gpts/__init__.py
index db3e0479..fdf102e8 100644
--- a/src/mcore_bridge/model/mm_gpts/__init__.py
+++ b/src/mcore_bridge/model/mm_gpts/__init__.py
@@ -1,3 +1,3 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 from . import (gemma4, glm, internvl, kimi_vl, llama4, llava, minicpmv4_6, muse_glimmer, qwen, qwen3_5, qwen3_5_gdn,
-               qwen3_asr, qwen3_omni, qwen3_vl)
+               qwen3_8_flash_next, qwen3_asr, qwen3_omni, qwen3_vl)
diff --git a/src/mcore_bridge/model/mm_gpts/qwen3_vl.py b/src/mcore_bridge/model/mm_gpts/qwen3_vl.py
index 8a5a6362..bb594714 100644
--- a/src/mcore_bridge/model/mm_gpts/qwen3_vl.py
+++ b/src/mcore_bridge/model/mm_gpts/qwen3_vl.py
@@ -80,7 +80,8 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config):
                 deepstack_visual_embeds = visual_res.deepstack_features
             else:
                 image_embeds, deepstack_visual_embeds = visual_res
-            deepstack_visual_embeds = torch.stack(deepstack_visual_embeds, dim=0)
+            deepstack_visual_embeds = (torch.stack(deepstack_visual_embeds, dim=0)
+                                       if deepstack_visual_embeds else None)
             inputs_embeds = inputs_embeds + image_embeds.mean().to(device=inputs_embeds.device) * 0.
             visual_pos_masks = None
         else:
@@ -125,7 +126,7 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config):
                 inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
             image_mask, video_mask = image_mask[..., 0], video_mask[..., 0]
             visual_pos_masks = image_mask | video_mask
-            if image_embeds is not None and video_embeds is not None:
+            if image_embeds is not None and video_embeds is not None and deepstack_visual_embeds:
                 deepstack_image_embeds = [tensor[:image_tokens] for tensor in deepstack_visual_embeds]
                 deepstack_video_embeds = [tensor[image_tokens:] for tensor in deepstack_visual_embeds]
                 deepstack_visual_embeds = []
@@ -137,7 +138,8 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config):
                     embed_joint[video_mask_joint, :] = vid_embed
                     deepstack_visual_embeds.append(embed_joint)
 
-            deepstack_visual_embeds = torch.stack(deepstack_visual_embeds, dim=0)
+            deepstack_visual_embeds = (torch.stack(deepstack_visual_embeds, dim=0)
+                                       if deepstack_visual_embeds else None)
             visual_pos_masks = visual_pos_masks.transpose(0, 1)
             # compat cp
             if self.config.context_parallel_size > 1:
@@ -147,7 +149,8 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config):
                 cu_seqlens = getattr(packed_seq_params, 'cu_seqlens_q', None)
                 cp_mask = split_cp_inputs(cp_mask, cu_seqlens, 0)
                 visual_pos_masks = split_cp_inputs(visual_pos_masks, cu_seqlens, 0)
-                deepstack_visual_embeds = deepstack_visual_embeds[:, cp_mask[(cp_mask != -1)]]
+                if deepstack_visual_embeds is not None:
+                    deepstack_visual_embeds = deepstack_visual_embeds[:, cp_mask[(cp_mask != -1)]]
             # compat sp
             tp_world_size = parallel_state.get_tensor_model_parallel_world_size()
             tp_rank = parallel_state.get_tensor_model_parallel_rank()
@@ -157,7 +160,8 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config):
                 visual_start = 0 if tp_rank == 0 else sum(mask_tokens[:tp_rank])
                 visual_end = visual_start + mask_tokens[tp_rank]
                 visual_pos_masks = visual_pos_masks[tp_rank]
-                deepstack_visual_embeds = deepstack_visual_embeds[:, visual_start:visual_end]
+                if deepstack_visual_embeds is not None:
+                    deepstack_visual_embeds = deepstack_visual_embeds[:, visual_start:visual_end]
         return {
             'inputs_embeds': inputs_embeds,
             'visual_pos_masks': visual_pos_masks,
diff --git a/src/mcore_bridge/model/modules/__init__.py b/src/mcore_bridge/model/modules/__init__.py
index 7269abdc..7f5f2a2e 100644
--- a/src/mcore_bridge/model/modules/__init__.py
+++ b/src/mcore_bridge/model/modules/__init__.py
@@ -4,8 +4,11 @@
 from .dsa_indexer import DSAIndexer
 from .gated_delta_net import GatedDeltaNet
 from .gated_self_attention import GatedSelfAttention
+from .hyper_connection_gated import Qwen4ExpTextGatedResidual, Qwen4ExpTextGroupedRMSNorm
 from .mtp_layer import MultiTokenPredictionLayer
 from .multi_latent_attention import MLASelfAttention
+from .ple import Qwen4ExpTextNGramEmbedding, Qwen4ExpTextPLELayer
+from .qsa_indexer import QSAIndexer
 from .topk_router import TopKRouter
 from .transformer_block import TransformerBlock
 from .transformer_layer import TransformerLayer
diff --git a/src/mcore_bridge/model/modules/gated_delta_net.py b/src/mcore_bridge/model/modules/gated_delta_net.py
index 206150eb..b2e80a75 100644
--- a/src/mcore_bridge/model/modules/gated_delta_net.py
+++ b/src/mcore_bridge/model/modules/gated_delta_net.py
@@ -161,18 +161,28 @@ def _resolve_cu_seqlens(cu_seqlens_padded, cu_seqlens_actual, total_seq_len, cp_
             return None
         return cu_seqlens
 
-    def _set_linear_sequence_parallel(self, enabled: bool) -> dict[str, bool]:
+    def _set_linear_sequence_parallel(self, enabled: bool) -> dict[str, list]:
+        # Walk submodules, not just the top-level module: with LoRA the linear is
+        # wrapped (LoraParallelLinear copies `sequence_parallel` onto itself), so
+        # flipping only the wrapper leaves the inner base_layer still in SP mode
+        # and it gathers a second time over an already-gathered sequence
+        # (upstream megatron-bridge PR #172, fixing #162/#169).
         saved = {}
         for name in ('in_proj', 'in_proj_qkvz', 'in_proj_ba', 'out_proj'):
             module = getattr(self, name, None)
-            if module is not None and hasattr(module, 'sequence_parallel'):
-                saved[name] = module.sequence_parallel
-                module.sequence_parallel = enabled
+            if module is None:
+                continue
+            states = [(sub, sub.sequence_parallel) for sub in module.modules() if hasattr(sub, 'sequence_parallel')]
+            if states:
+                saved[name] = states
+                for sub, _ in states:
+                    sub.sequence_parallel = enabled
         return saved
 
-    def _restore_linear_sequence_parallel(self, saved: dict[str, bool]) -> None:
-        for name, enabled in saved.items():
-            getattr(self, name).sequence_parallel = enabled
+    def _restore_linear_sequence_parallel(self, saved: dict[str, list]) -> None:
+        for states in saved.values():
+            for sub, enabled in states:
+                sub.sequence_parallel = enabled
 
     def forward(
         self,
diff --git a/src/mcore_bridge/model/modules/transformer_block.py b/src/mcore_bridge/model/modules/transformer_block.py
index 0beb20d1..05a333f0 100644
--- a/src/mcore_bridge/model/modules/transformer_block.py
+++ b/src/mcore_bridge/model/modules/transformer_block.py
@@ -358,11 +358,15 @@ def forward(
         #   is called here to be future-proof and corner-case-proof.
         hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)
         enable_hyper_connections = getattr(self.config, 'enable_hyper_connections', False)
+        hc_count = getattr(self.config, 'hc_count', 0) or 0
+        enable_gated_hc = hc_count > 1 and not enable_hyper_connections
         # Expand hidden states for hyper connections at the start of the block
         # Only expand at the first PP stage; subsequent stages receive n-stream from previous stage
         if enable_hyper_connections and self.pre_process:
             hidden_states = HyperConnectionModule.input_expand(hidden_states,
                                                                self.num_residual_streams)  # [s, b, C] -> [s, b, n*C]
+        elif enable_gated_hc and self.pre_process:
+            hidden_states = hidden_states.repeat(1, 1, hc_count)  # [s, b, C] -> [s, b, n*C]
 
         if self.config.sequence_parallel:
             rng_context = tensor_parallel.get_cuda_rng_tracker().fork()
@@ -490,6 +494,11 @@ def forward(
                 self.config.num_residual_streams,
                 self.config.layernorm_epsilon,
             )
+        elif enable_gated_hc and self.has_final_layernorm_in_this_stage():
+            # Gated low-rank contraction (hyper_connection_mixer, use_combine=False
+            # so forward returns only the mixed stream).
+            # [s, b, n*C] -> [s, b, C]
+            hidden_states = self.hyper_connection_mixer(hidden_states)
 
         # Final layer norm.
         if self.final_layernorm is not None:
@@ -515,4 +524,14 @@ def forward(
         if mhc_multistream is not None:
             return hidden_states, mhc_multistream
 
+        # A non-last PP stage hands its hidden_states (still multi-stream under
+        # mHC/gated-HC, since contraction happens on the last stage) straight to
+        # the next stage, and pipeline schedules' deallocate_output_tensor()
+        # asserts the tensor is not a view. The last-stage path guards this
+        # after the final layernorm (make_viewless_tensor above); the same
+        # guard is needed here because the last layer's output can itself be a
+        # view (e.g. TENorm).
+        if (enable_hyper_connections or enable_gated_hc) and not self.has_final_layernorm_in_this_stage():
+            hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)
+
         return hidden_states

From 89a9c2a4d72423502a7a6dd65f499a4afeceef6a Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Wed, 26 Aug 2026 21:18:04 +0800
Subject: [PATCH 02/10] update

---
 src/mcore_bridge/config/parser.py             |   5 +-
 src/mcore_bridge/model/gpts/__init__.py       |   2 +-
 src/mcore_bridge/model/gpts/qwen4_exp.py      | 484 ++++++++++++++++++
 src/mcore_bridge/model/mm_gpts/__init__.py    |   2 +-
 src/mcore_bridge/model/mm_gpts/qwen4_exp.py   |  30 ++
 .../model/modules/hyper_connection_gated.py   |  76 +++
 src/mcore_bridge/model/modules/ple.py         | 403 +++++++++++++++
 7 files changed, 997 insertions(+), 5 deletions(-)
 create mode 100644 src/mcore_bridge/model/gpts/qwen4_exp.py
 create mode 100644 src/mcore_bridge/model/mm_gpts/qwen4_exp.py
 create mode 100644 src/mcore_bridge/model/modules/hyper_connection_gated.py
 create mode 100644 src/mcore_bridge/model/modules/ple.py

diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py
index a2a8e818..1641447f 100644
--- a/src/mcore_bridge/config/parser.py
+++ b/src/mcore_bridge/config/parser.py
@@ -52,7 +52,7 @@
     'linear_key_head_dim': ['linear_key_head_dim'],
     'linear_value_head_dim': ['linear_value_head_dim'],
     'linear_conv_kernel_dim': ['linear_conv_kernel_dim'],
-    # qwen3_8_flash_next
+    # qwen4_exp
     'hc_count': ['hc_count'],
     'hc_lowrank': ['hc_lowrank'],
     'ple_layer_ids': ['ple_layer_ids'],
@@ -261,8 +261,7 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
         if use_mcore_gdn:
             res['experimental_attention_variant'] = 'gated_delta_net'
         res.setdefault('linear_attention_freq', 4)
-    # TODO: confirm model_type, remove one
-    elif hf_model_type in {'qwen3_8_flash_next', 'qwen4_exp'}:
+    elif hf_model_type in {'qwen4_exp'}:
         use_mcore_gdn = get_env_args('USE_MCORE_GDN', bool, True)
         res['layernorm_zero_centered_gamma'] = True
         res['attention_output_gate'] = True
diff --git a/src/mcore_bridge/model/gpts/__init__.py b/src/mcore_bridge/model/gpts/__init__.py
index 0bc7b780..655415bc 100644
--- a/src/mcore_bridge/model/gpts/__init__.py
+++ b/src/mcore_bridge/model/gpts/__init__.py
@@ -1,3 +1,3 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 from . import (bailing_hybrid, bailing_moe, deepseek_v4, glm4, glm_moe_dsa, hunyuan, llm, minimax_m2, nemotron_h, olmoe,
-               qwen3_8_flash_next, qwen3_emb, qwen3_next)
+               qwen3_emb, qwen3_next, qwen4_exp)
diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py
new file mode 100644
index 00000000..a3052903
--- /dev/null
+++ b/src/mcore_bridge/model/gpts/qwen4_exp.py
@@ -0,0 +1,484 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import copy
+import torch
+import torch.distributed as dist
+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.ssm.gated_delta_net import GatedDeltaNetSubmodules
+from megatron.core.transformer.enums import AttnMaskType
+from megatron.core.transformer.identity_op import IdentityOp
+from megatron.core.transformer.spec_utils import ModuleSpec
+from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
+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 ..modules import (GatedDeltaNet, Qwen4ExpTextGatedResidual, Qwen4ExpTextPLELayer, QSAIndexer, TransformerBlock,
+                       TransformerLayer)
+from ..register import ModelLoader
+from .qwen3_next import Qwen3NextBridge, Qwen3NextRMSNorm, Qwen3NextSelfAttention
+
+logger = get_logger()
+
+_HC_WEIGHT_KEYS = (
+    'hc_norm.weight',
+    'input_mix_weight_down.weight',
+    'input_mix_weight_up.weight',
+    'block_inject_weight.weight',
+)
+
+
+class Qwen4ExpGDN(GatedDeltaNet):
+    # upstream uses config.activation_func as the act_fn for both the gated output
+    # norm and the conv1d; but the conv1d path asserts act_fn in ['silu', 'swish'],
+    # so setting it to sigmoid would be rejected. Override only the output gate here.
+    def _apply_gated_norm(self, x, gate):
+        x_dtype = x.dtype
+        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)
+        gate_act = torch.sigmoid if output_gate_type == 'sigmoid' else F.silu
+        y = y * gate_act(gate.float())
+        return y.to(x_dtype)
+
+
+class Qwen4ExpLayer(TransformerLayer):
+    # refer: transformers Qwen4ExpTextDecoderLayer
+    def __init__(self, config, submodules, layer_number: int = 1, **kwargs):
+        super().__init__(config, submodules, layer_number, **kwargs)
+        self.ple = None
+        if self.layer_number in config.ple_layer_ids:
+            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:
+            self.self_attention.indexer = QSAIndexer(config)
+        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')
+        attn_kwargs = dict(
+            attention_mask=attention_mask,
+            inference_context=kwargs.get('inference_context'),
+            rotary_pos_emb=kwargs.get('rotary_pos_emb'),
+            rotary_pos_cos=kwargs.get('rotary_pos_cos'),
+            rotary_pos_sin=kwargs.get('rotary_pos_sin'),
+            attention_bias=kwargs.get('attention_bias'),
+            packed_seq_params=packed_seq_params,
+            sequence_len_offset=kwargs.get('sequence_len_offset'),
+        )
+        if self.ple is not None:
+            input_ids = kwargs.get('input_ids')
+            assert input_ids is not None, 'PLE layers require input_ids in extra_block_kwargs'
+            hidden_states = hidden_states + self.ple(hidden_states, input_ids, packed_seq_params)
+
+        # 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):
+            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)
+
+        # mlp sub-block
+        hidden_states, hyper_input, injection_weights = self.mlp_hyper_connection(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        if isinstance(hidden_states, tuple):
+            hidden_states = hidden_states[0]
+        injection = hidden_states.unsqueeze(-2) * injection_weights.unsqueeze(-1)
+        hidden_states = hyper_input + injection.flatten(-2)
+        return hidden_states, None
+
+    @contextmanager
+    def _qsa_arbitrary_mask(self, enabled: bool):
+        """Temporarily switch the attention to `arbitrary` mask type.
+
+        `Attention.forward` takes no `attn_mask_type` argument -- it reads
+        `self.attn_mask_type` (and core_attention's) -- so a custom mask is
+        silently ignored unless the type is flipped. Restored in `finally` so a
+        raising forward cannot leave the layer stuck in the slower unfused mode.
+        """
+        if not enabled:
+            yield
+            return
+        attn = self.self_attention
+        targets = [attn]
+        core = getattr(attn, 'core_attention', None)
+        if core is not None:
+            targets.append(core)
+        saved = [(t, t.attn_mask_type) for t in targets if hasattr(t, 'attn_mask_type')]
+        for t, _ in saved:
+            t.attn_mask_type = AttnMaskType.arbitrary
+        try:
+            yield
+        finally:
+            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_mask(self, hidden_states, attn_kwargs):
+        # return None means full attention
+        indexer = getattr(self.self_attention, 'indexer', None)
+        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
+        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.')
+            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)
+
+
+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():
+            # 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)
+
+
+class Qwen4ExpBridge(Qwen3NextBridge):
+    hf_mixer_prefix = 'model.'
+
+    def _get_hf_experts_attr(self, is_mtp: bool = False):
+        # The checkpoint stores experts as packed per-layer tensors
+        # (`mlp.experts.gate_up_proj` / `mlp.experts.down_proj`).
+        return True, True
+
+    def _set_layer_attn(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool):
+        mg_attn = None if mg_layer is None else mg_layer.self_attention
+        is_linear_attention = self.config.linear_attention_freq[layer_idx]
+        if is_linear_attention:
+            # GDN weights; this model has no input_layernorm for GDN layers.
+            hf_state_dict.update(
+                self._set_linear_attn_state(mg_attn, hf_state_dict, 'linear_attn.', layer_idx, to_mcore))
+        else:
+            # Dense QSA-equivalent attention (qkv + output gate + q/k norms).
+            hf_state_dict.update(self._set_attn_state(mg_attn, hf_state_dict, 'self_attn.', layer_idx, to_mcore))
+            has_indexer = mg_attn is not None and getattr(mg_attn, 'indexer', None) is not None
+            has_indexer = self._reduce_tensor_pp_group(has_indexer, to_mcore)
+            if has_indexer:
+                indexer = None if mg_attn is None else mg_attn.indexer
+                for mg_key, hf_key in [('index_qk_proj.weight', 'self_attn.indexer.index_qk_proj.weight'),
+                                       ('q_layernorm.weight', 'self_attn.indexer.q_layernorm.weight'),
+                                       ('k_layernorm.weight', 'self_attn.indexer.k_layernorm.weight')]:
+                    self._set_state_dict(indexer, mg_key, hf_state_dict, hf_key, to_mcore)
+        return hf_state_dict
+
+    def _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool, is_mtp: bool = False):
+        mg_mlp = None if mg_layer is None else mg_layer.mlp
+        is_moe = mg_mlp is not None and hasattr(mg_mlp, 'experts')
+        if not to_mcore:
+            is_moe = torch.tensor([is_moe], dtype=torch.bool, device='cuda')
+            if self.pp_size > 1:
+                dist.all_reduce(is_moe, group=self.pp_group)
+        if is_moe:
+            hf_state_dict.update(
+                self._set_moe_state(mg_mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', layer_idx, to_mcore,
+                                    is_mtp=is_mtp))
+        else:
+            hf_state_dict.update(
+                self._set_mlp_state(mg_mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', layer_idx, to_mcore))
+        # No post_attention_layernorm in this model (HC norms replace it).
+        return hf_state_dict
+
+    def _set_layer_hc(self, mg_layer, hf_state_dict, to_mcore: bool):
+        for key in ['attn_hyper_connection', 'mlp_hyper_connection']:
+            hyper_connection = None if mg_layer is None else getattr(mg_layer, key)
+            for weight_key in _HC_WEIGHT_KEYS:
+                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'),
+    )
+
+    def _get_tp_split_dim(self, mg_key):
+        # PLE weights are replicated across TP; in particular `conv1d.weight`
+        # must not use the dim-0 split that applies to the GDN conv1d.
+        if getattr(self, '_converting_ple', False):
+            return None
+        return super()._get_tp_split_dim(mg_key)
+
+    def _get_pp_src_rank(self, has_module: bool) -> int:
+        """Global rank of the PP stage holding the module (all-reduce MAX)."""
+        holder = torch.tensor([dist.get_rank() if has_module else -1], dtype=torch.long, device='cuda')
+        if self.pp_size > 1:
+            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:
+            # Only the stage owning the PLE layer reaches this path, so it
+            # must not run pp collectives (other pp ranks never enter here).
+            if ple is None:
+                return
+            pp_src_rank = None
+        else:
+            # to_hf: every pp rank calls this for every layer, so the pp
+            # collectives below stay in sync across stages.
+            pp_src_rank = self._get_pp_src_rank(ple is not None)
+            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:
+            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))
+            else:
+                tensor = getattr(ple.ple_embedding, mg_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)
+                    tensor = obj[0]
+                # Written directly into the state dict (bypasses _get_weight,
+                # 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)
+        self._converting_ple = True
+        try:
+            for mg_key, hf_key in [('key_proj.weight', 'ple.key_proj.weight'),
+                                   ('value_proj.weight', 'ple.value_proj.weight'),
+                                   ('norm_key.weight', 'ple.norm_key.weight'),
+                                   ('norm_query.weight', 'ple.norm_query.weight'),
+                                   ('norm_conv.weight', 'ple.norm_conv.weight'),
+                                   ('conv1d.weight', 'ple.conv1d.weight')]:
+                self._set_state_dict(ple, mg_key, hf_state_dict, hf_key, to_mcore)
+        finally:
+            self._converting_ple = False
+
+    def _set_layer_state(self, mg_layer, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool):
+        hf_prefix = f'{hf_prefix}{layer_idx}.'
+        if to_mcore:
+            hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix)
+        else:
+            hf_state_dict = {}
+        hf_state_dict.update(self._set_layer_attn(mg_layer, hf_state_dict, layer_idx, to_mcore))
+        hf_state_dict.update(self._set_layer_mlp(mg_layer, hf_state_dict, layer_idx, to_mcore))
+        self._set_layer_hc(mg_layer, hf_state_dict, to_mcore)
+        if (layer_idx + 1) in (self.config.ple_layer_ids or []):
+            self._set_layer_ple(mg_layer, hf_state_dict, to_mcore)
+        if to_mcore:
+            hf_state_dict = {}
+        else:
+            hf_state_dict = self._add_prefix(hf_state_dict, hf_prefix)
+        return hf_state_dict
+
+    def _set_final_layernorm(self, lm_model, hf_state_dict, to_mcore):
+        # This architecture has no final layernorm: the HC norms and the
+        # hyper_connection_mixer contraction replace it, and the checkpoint
+        # carries no `norm` weight.
+        pass
+
+    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
+        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']
+            # super() returns {} in to_mcore mode: read from the incoming full
+            # state dict; in to_hf mode write into the dict super() returns.
+            mixer_sd = hf_state_dict if to_mcore else res
+            for key in mixer_keys:
+                self._set_state_dict(lm_model, f'decoder.hyper_connection_mixer.{key}', mixer_sd,
+                                     f'{self.hf_mixer_prefix}hyper_connection_mixer.{key}', to_mcore)
+        return res
+
+
+class Qwen4ExpLoader(ModelLoader):
+    transformer_block = Qwen4ExpTransformerBlock
+
+    def _get_moe_layer_pattern(self) -> List[bool]:
+        config = self.config
+        freq = config.moe_layer_freq
+        if isinstance(freq, list):
+            return [bool(x) for x in freq]
+        # int N: one MoE every N layers (mcore convention: i % N == N - 1).
+        return [i % freq == freq - 1 for i in range(config.num_layers)]
+
+    def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
+        config = self.config
+        config.hetereogenous_dist_checkpoint = True
+        # Context parallelism: PLE gathers the full sequence internally (undoing the
+        # CP zigzag) and GDN carries its own CP handling (a2a CP<->HP plus CP-aware
+        # cu_seqlens), so CP is no longer blanket-rejected here. Left unasserted so
+        # it can be exercised; QSA layers run dense attention, which mcore's
+        # attention already supports under CP.
+        if getattr(config, 'mtp_num_layers', None):
+            raise NotImplementedError(
+                'Qwen4-Exp MTP is not supported yet')
+        moe_spec = get_gpt_layer_with_transformer_engine_spec(
+            num_experts=config.num_moe_experts,
+            moe_grouped_gemm=config.moe_grouped_gemm,
+            qk_layernorm=config.qk_layernorm,
+            multi_latent_attention=config.multi_latent_attention,
+            use_kitchen=config.use_kitchen,
+        )
+        if config.num_moe_experts is not None:
+            dense_spec = get_gpt_layer_with_transformer_engine_spec(
+                num_experts=None,
+                moe_grouped_gemm=config.moe_grouped_gemm,
+                qk_layernorm=config.qk_layernorm,
+                multi_latent_attention=config.multi_latent_attention,
+                use_kitchen=config.use_kitchen,
+            )
+        else:
+            dense_spec = moe_spec
+        gdn_spec = ModuleSpec(
+            module=Qwen4ExpGDN,
+            submodules=GatedDeltaNetSubmodules(
+                in_proj=TEColumnParallelLinear,
+                out_norm=TENorm,
+                out_proj=TERowParallelLinear,
+            ),
+        )
+        moe_pattern = self._get_moe_layer_pattern()
+        layer_specs = []
+        for layer_idx, is_linear_attention in enumerate(config.linear_attention_freq):
+            from copy import deepcopy
+            layer_spec = deepcopy(moe_spec if moe_pattern[layer_idx] else dense_spec)
+            if is_linear_attention:
+                layer_spec.submodules.self_attention = deepcopy(gdn_spec)
+            else:
+                layer_spec.submodules.self_attention.submodules.linear_qkv = TEColumnParallelLinear
+                layer_spec.submodules.self_attention.module = Qwen3NextSelfAttention
+                if hasattr(layer_spec.submodules.self_attention.submodules, 'q_layernorm'):
+                    layer_spec.submodules.self_attention.submodules.q_layernorm = Qwen3NextRMSNorm
+                if hasattr(layer_spec.submodules.self_attention.submodules, 'k_layernorm'):
+                    layer_spec.submodules.self_attention.submodules.k_layernorm = Qwen3NextRMSNorm
+            # This model has no per-layer layernorms (HC norms replace them).
+            layer_spec.submodules.input_layernorm = IdentityOp
+            if hasattr(layer_spec.submodules, 'pre_mlp_layernorm'):
+                layer_spec.submodules.pre_mlp_layernorm = IdentityOp
+            layer_specs.append(layer_spec)
+
+        local_layer_specs = get_local_layer_specs(config, layer_specs, vp_stage=vp_stage)
+        # No final layernorm in this model; keep the slot so the HC mixer
+        # stage logic (has_final_layernorm_in_this_stage) still triggers.
+        block_spec = TransformerBlockSubmodules(layer_specs=local_layer_specs, layer_norm=IdentityOp)
+        return block_spec
+
+    def _set_transformer_layer(self, transformer_layer_spec):
+        for layer_spec in transformer_layer_spec.layer_specs:
+            layer_spec.module = Qwen4ExpLayer
+
+    def build_model(
+        self,
+        pre_process=True,
+        post_process=True,
+        vp_stage: Optional[int] = None,
+    ):
+        model = super().build_model(pre_process, post_process, vp_stage)
+        lm_model = model.language_model if hasattr(model, 'language_model') else model
+        # The GDN out_norm uses ones-style weights, unlike the zero-centered
+        # HC norms, so opt it out of layernorm_zero_centered_gamma.
+        for layer in lm_model.decoder.layers:
+            if hasattr(layer.self_attention, 'out_norm'):
+                out_norm = layer.self_attention.out_norm
+                out_norm.zero_centered_gamma = False
+                if not is_torch_npu_available():
+                    assert hasattr(out_norm, 'zero_centered_gamma')
+                if hasattr(out_norm, 'config'):
+                    out_norm.config = copy.copy(out_norm.config)
+                    out_norm.config.layernorm_zero_centered_gamma = False
+        return model
diff --git a/src/mcore_bridge/model/mm_gpts/__init__.py b/src/mcore_bridge/model/mm_gpts/__init__.py
index fdf102e8..a931bb97 100644
--- a/src/mcore_bridge/model/mm_gpts/__init__.py
+++ b/src/mcore_bridge/model/mm_gpts/__init__.py
@@ -1,3 +1,3 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 from . import (gemma4, glm, internvl, kimi_vl, llama4, llava, minicpmv4_6, muse_glimmer, qwen, qwen3_5, qwen3_5_gdn,
-               qwen3_8_flash_next, qwen3_asr, qwen3_omni, qwen3_vl)
+               qwen3_asr, qwen3_omni, qwen3_vl, qwen4_exp)
diff --git a/src/mcore_bridge/model/mm_gpts/qwen4_exp.py b/src/mcore_bridge/model/mm_gpts/qwen4_exp.py
new file mode 100644
index 00000000..1bd6903b
--- /dev/null
+++ b/src/mcore_bridge/model/mm_gpts/qwen4_exp.py
@@ -0,0 +1,30 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from mcore_bridge.utils import get_env_args
+
+from ..constant import ModelType
+from ..gpts.qwen4_exp import Qwen4ExpBridge, Qwen4ExpLoader, Qwen4ExpTransformerBlock
+from ..register import ModelMeta, register_model
+from .qwen3_5 import Qwen3_5Vit
+
+
+class Qwen4ExpMMBridge(Qwen4ExpBridge):
+    hf_layers_prefix = 'model.language_model.layers'
+    hf_embed_key = 'model.language_model.embed_tokens.weight'
+    hf_mixer_prefix = 'model.language_model.'
+
+
+class Qwen4ExpMMLoader(Qwen4ExpLoader):
+    transformer_block = Qwen4ExpTransformerBlock
+
+
+use_mcore_gdn = get_env_args('USE_MCORE_GDN', bool, True)
+
+if use_mcore_gdn:
+    register_model(
+        ModelMeta(
+            ModelType.qwen4_exp,
+            ['qwen4_exp'],
+            bridge_cls=Qwen4ExpMMBridge,
+            visual_cls=Qwen3_5Vit,
+            loader=Qwen4ExpMMLoader,
+        ))
diff --git a/src/mcore_bridge/model/modules/hyper_connection_gated.py b/src/mcore_bridge/model/modules/hyper_connection_gated.py
new file mode 100644
index 00000000..faeb1066
--- /dev/null
+++ b/src/mcore_bridge/model/modules/hyper_connection_gated.py
@@ -0,0 +1,76 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import torch
+import torch.nn.functional as F
+from torch import nn
+from typing import Tuple
+
+
+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.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:
+        x = x.reshape(*x.shape[:-1], -1, self.group_size)
+        out = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return out.flatten(-2)
+
+    def forward(self, x):
+        output = self._norm(x.float())
+        output = output * (1.0 + self.weight.float())
+        return output.type_as(x)
+
+    def extra_repr(self):
+        return f'{tuple(self.weight.shape)}, eps={self.eps}'
+
+
+class Qwen4ExpTextGatedResidual(nn.Module):
+    def __init__(self, config, use_combine: bool = True):
+        # Mirrors transformers `Qwen4ExpTextGatedResidual.__init__` (config-driven).
+        super().__init__()
+        self.hc_count = config.hc_count
+        self.hidden_size = config.hidden_size
+        hc_hidden_size = self.hc_count * self.hidden_size
+        # (transformers: Qwen4ExpTextRMSNorm(hc_hidden_size, group_size=self.hidden_size,
+        #  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
+        # mcore-specific: SP grad flag on the replicated weights.
+        for param in self.parameters():
+            setattr(param, 'sequence_parallel', config.sequence_parallel)
+
+    def forward(self, hyper_input: torch.Tensor) -> Tuple[torch.Tensor, ...]:
+        # Mirrors transformers `Qwen4ExpTextGatedResidual.forward` line by line.
+        if hyper_input.shape[-1] != self.hc_count * self.hidden_size:
+            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)
+        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)
+        return mixed_input, hyper_input, injection_weights
+
diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py
new file mode 100644
index 00000000..23bca0e7
--- /dev/null
+++ b/src/mcore_bridge/model/modules/ple.py
@@ -0,0 +1,403 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import copy
+import math
+import torch
+import torch.nn.functional as F
+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.megatron_utils import reconstruct_tensor_cp, split_cp_inputs
+from .hyper_connection_gated import Qwen4ExpTextGroupedRMSNorm
+
+_MASK64 = (1 << 64) - 1
+_SPLITMIX_GAMMA = 0x9E3779B97F4A7C15
+_SPLITMIX_M1 = 0xBF58476D1CE4E5B9
+_SPLITMIX_M2 = 0x94D049BB133111EB
+_PRIME_1 = 10007
+
+
+def _splitmix64(value: int) -> int:
+    value = (value + _SPLITMIX_GAMMA) & _MASK64
+    value = ((value ^ (value >> 30)) * _SPLITMIX_M1) & _MASK64
+    value = ((value ^ (value >> 27)) * _SPLITMIX_M2) & _MASK64
+    return (value ^ (value >> 31)) & _MASK64
+
+
+def _build_layer_multipliers(unigram_vocab_size: int, ngram_size: int, ple_layer_index: int, seed: int) -> List[int]:
+    max_long = (1 << 63) - 1
+    multiplier_max = max_long // max(unigram_vocab_size, 1)
+    half_bound = max(1, multiplier_max // 2)
+    base_seed = seed + _PRIME_1 * ple_layer_index
+    multipliers = []
+    for index in range(ngram_size):
+        value = (base_seed + _SPLITMIX_GAMMA * (index + 1)) & _MASK64
+        multipliers.append(2 * (_splitmix64(value) % half_bound) + 1)
+    return multipliers
+
+
+def _is_prime_64(value: int) -> bool:
+    if value < 2:
+        return False
+    for prime in (2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37):
+        if value % prime == 0:
+            return value == prime
+    exponent = value - 1
+    shifts = 0
+    while exponent % 2 == 0:
+        exponent //= 2
+        shifts += 1
+    for base in (2, 325, 9375, 28178, 450775, 9780504, 1795265022):
+        if base % value == 0:
+            continue
+        witness = pow(base, exponent, value)
+        if witness in (1, value - 1):
+            continue
+        for _ in range(shifts - 1):
+            witness = pow(witness, 2, value)
+            if witness == value - 1:
+                break
+        else:
+            return False
+    return True
+
+
+def _nth_prime_after(start: int, count: int) -> int:
+    # Mirrors transformers `_find_nth_prime_after`.
+    prime = int(start)
+    for _ in range(count):
+        candidate = prime + 1
+        if candidate <= 2:
+            prime = 2
+            continue
+        if candidate % 2 == 0:
+            candidate += 1
+        while not _is_prime_64(candidate):
+            candidate += 2
+        prime = candidate
+    return prime
+
+
+class Qwen4ExpTextNGramEmbedding(nn.Module):
+    def __init__(self, config, ple_layer_index: int):
+        super().__init__()
+        self.ngram_size = config.ngram_size
+        self.context_len = self.ngram_size - 1
+        self.heads_per_ngram = config.heads_per_ngram
+        self.ngram_heads = (self.ngram_size - 1) * self.heads_per_ngram
+        self.ple_layer_index = ple_layer_index
+        # (transformers also stores unigram_vocab_size/ngram_vocab_size_base/
+        # seed as attributes, consumed only by its deferred _init_weights
+        # buffer population; here the buffers are registered at construction,
+        # so no copies are kept.)
+        head_dim_per_ngram = config.ple_embed_dim // self.ngram_heads
+        # mcore-specific fail-loud: these drive the hash math and the checkpoint
+        # shard layout, so they must come from the model config (the parser
+        # 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)
+        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}).')
+        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 = _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)
+
+        # Per-head prime table sizes/offsets (checkpoint-persistent), named as
+        # in transformers.
+        self.head_vocab_sizes = []
+        self.head_offsets = []
+        self.total_vocab_size = 0
+        for head_idx in range(self.ngram_heads):
+            global_head_idx = self.ple_layer_index * self.ngram_heads + head_idx
+            size = _nth_prime_after(config.ngram_vocab_size_base - 1, global_head_idx + 1)
+            self.head_vocab_sizes.append(size)
+            self.head_offsets.append(self.total_vocab_size)
+            self.total_vocab_size += size
+        self.register_buffer('ngram_heads_vocab_sizes', torch.tensor(self.head_vocab_sizes, dtype=torch.long),
+                             persistent=True)
+        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,
+        )
+
+    def _shift_right_ignore_eos(self, token_ids: torch.Tensor, shift: int) -> torch.Tensor:
+        # Mirrors transformers `_shift_right_ignore_eos`: segment-aware shift,
+        # segments are reset after every eos token (request boundary logic).
+        if shift == 0:
+            return token_ids
+        batch_size, seq_len = token_ids.shape
+        positions = torch.arange(seq_len, device=token_ids.device, dtype=torch.long)
+        eos_positions = torch.where(token_ids == self.eos_token_id, positions, -1)
+        previous_eos_inclusive = torch.cummax(eos_positions, dim=1).values
+        previous_eos = torch.cat([eos_positions.new_full((batch_size, 1), -1), previous_eos_inclusive[:, :-1]], dim=1)
+        segment_start = previous_eos + 1
+        position_in_segment = positions.unsqueeze(0) - segment_start
+        source_positions = positions - shift
+        gather_positions = source_positions.clamp_min(0).unsqueeze(0).expand(batch_size, -1)
+        shifted = token_ids.gather(dim=1, index=gather_positions)
+        valid = (position_in_segment >= shift) & (source_positions.unsqueeze(0) >= 0)
+        return torch.where(valid, shifted, token_ids.new_full((), self.eos_token_id))
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        """input_ids: [rows, L] -> embeddings [rows, L, embedding_dim].
+
+        Training variant of transformers ``forward(input_ids, past_key_values,
+        layer_idx)``: the previous context is eos padding (fresh request),
+        prepended here and dropped from the ids afterwards.
+        """
+        input_ids = input_ids.long()
+        # mcore-specific: training has no conv-state cache; a fresh request
+        # starts from eos context, so prepend it (the no-previous-state branch
+        # of transformers `forward` returns exactly this). The latest HF
+        # constructor also takes a cache-only `layer_idx`, which has no
+        # training counterpart.
+        previous_context = input_ids.new_full((input_ids.shape[0], self.context_len), self.eos_token_id)
+        token_history = torch.cat([previous_context, input_ids], dim=-1)
+        shifted_tokens = [self._shift_right_ignore_eos(token_history, shift) for shift in range(self.ngram_size)]
+
+        blocks = []
+        for ngram in range(2, self.ngram_size + 1):
+            start_idx = (ngram - 2) * self.heads_per_ngram
+            end_idx = start_idx + self.heads_per_ngram
+            mixed_ids = shifted_tokens[0] * self.layer_multipliers[0]
+            for position in range(1, ngram):
+                mixed_ids = torch.bitwise_xor(
+                    mixed_ids,
+                    shifted_tokens[position] * self.layer_multipliers[position],
+                )
+            head_vocab_sizes = self.ngram_heads_vocab_sizes[start_idx:end_idx]
+            head_offsets = self.ngram_heads_offsets[start_idx:end_idx]
+            ngram_ids = torch.remainder(mixed_ids.unsqueeze(-1), head_vocab_sizes.view(1, 1, -1))
+            blocks.append(ngram_ids + head_offsets.view(1, 1, -1))
+
+        ngram_ids = torch.cat(blocks, dim=-1)[:, -input_ids.shape[1]:]
+        return self.ngram_embedding(ngram_ids).flatten(-2)
+
+
+class Qwen4ExpTextPLELayer(nn.Module):
+    """Inject hashed n-gram features into every hyper-connection stream;
+    mirrors transformers ``Qwen4ExpTextPLELayer`` (training variant without
+    the inference cache and ``conv_mask``).
+
+    PLE projects each token's concatenated n-gram embedding to a shared value
+    and one key per residual stream. The normalized stream activations gate
+    those values, then a dilated depthwise convolution adds local lexical
+    context. The returned tensor has width ``hc_count * hidden_size``.
+
+    Checkpoint names under the layer prefix ``ple.``:
+        ple_embedding.{layer_multipliers,ngram_heads_offsets,ngram_heads_vocab_sizes}
+        ple_embedding.ngram_embedding.shard_{i}.weight
+        key_proj.weight / value_proj.weight
+        norm_key.weight / norm_query.weight / norm_conv.weight
+        conv1d.weight
+    """
+
+    def __init__(self, config, ple_layer_index: int, pg_collection=None):
+        super().__init__()
+        self.config = config
+        self.pg_collection = pg_collection
+        self.hidden_size = int(config.hidden_size)
+        self.hc_count = int(config.hc_count)
+        ple_embed_dim = int(config.ple_embed_dim)
+        hc_hidden_size = self.hidden_size * self.hc_count
+        self.ple_embedding = Qwen4ExpTextNGramEmbedding(config, ple_layer_index)
+        conv_kernel_size = int(config.ple_conv_kernel_size)
+        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)
+        # 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(hc_hidden_size, group_size=self.hidden_size,
+                                                   eps=config.layernorm_epsilon, dtype=config.params_dtype,
+                                                   sequence_parallel=config.sequence_parallel)
+        self.norm_query = Qwen4ExpTextGroupedRMSNorm(hc_hidden_size, group_size=self.hidden_size,
+                                                     eps=config.layernorm_epsilon, dtype=config.params_dtype,
+                                                     sequence_parallel=config.sequence_parallel)
+        self.norm_conv = Qwen4ExpTextGroupedRMSNorm(hc_hidden_size, group_size=self.hidden_size,
+                                                    eps=config.layernorm_epsilon, dtype=config.params_dtype,
+                                                    sequence_parallel=config.sequence_parallel)
+        self.conv1d = nn.Conv1d(
+            hc_hidden_size,
+            hc_hidden_size,
+            kernel_size=conv_kernel_size,
+            groups=hc_hidden_size,
+            dilation=conv_dilation,
+            bias=False,
+            dtype=config.params_dtype,
+        )
+        nn.init.zeros_(self.conv1d.weight)
+        for name, param in self.named_parameters():
+            if name.startswith(('key_proj', 'value_proj')) or name.endswith('norm.weight') or name.startswith(
+                ('norm_key', 'norm_query', 'norm_conv')) or name.startswith('conv1d'):
+                # Replicated across TP; reduce grads across TP when SP is on.
+                setattr(param, 'sequence_parallel', config.sequence_parallel)
+
+    def _short_conv(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # Training variant of transformers `_short_conv` (no conv state cache):
+        # causal alignment is achieved by left-padding the conv input with
+        # zeros (equivalent to trimming the padded conv output), which is a
+        # no-op on fresh-request context.
+        hidden_states = hidden_states.transpose(1, 2)
+        sequence_length = hidden_states.shape[-1]
+        conv_input = F.pad(hidden_states, (self.short_conv_state_len, 0))
+        conv_input = conv_input[..., -(self.short_conv_state_len + sequence_length):]
+        return F.silu(self.conv1d(conv_input)).transpose(1, 2)
+
+    def compute(self, hidden_states: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
+        """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)
+        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()
+        gated_value = torch.sigmoid(gate) * value.unsqueeze(-2)
+        gated_value_normed = self.norm_conv(gated_value.flatten(-2))
+        gated_value = gated_value.flatten(-2)
+        # (transformers also applies conv_mask here; training relies on the
+        # loss mask instead.)
+        output = gated_value + self._short_conv(gated_value_normed)
+        return output
+
+    @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.
+
+        Mirrors GatedDeltaNet._resolve_cu_seqlens for the SP-only case: accept a
+        global cumulative layout (possibly missing its leading 0 or carrying a
+        per-batch offset), and return a cu whose last entry equals ``total`` so it
+        indexes the gathered full-sequence tensor directly.
+        """
+        if cu is None:
+            return None
+        cu = cu.reshape(-1)
+        if cu.numel() > 0 and int(cu[0]) != 0:
+            if int(cu[-1]) == total:
+                cu = torch.cat([cu.new_zeros(1), cu])
+            elif int(cu[-1]) - int(cu[0]) == total:
+                cu = cu - cu[0]
+        if not (cu.numel() > 0 and int(cu[-1]) == total):
+            raise ValueError(
+                f'PLE cannot align cu_seqlens (last={int(cu[-1]) if cu.numel() else None}, '
+                f'first={int(cu[0]) if cu.numel() else None}) with the gathered sequence '
+                f'length {total} under sequence parallelism.')
+        return cu
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        input_ids: torch.Tensor,
+        packed_seq_params=None,
+    ) -> torch.Tensor:
+        """hidden_states: [s, b, nH] (bsh) or thd [T, 1, nH]; input_ids: [b, s] or [1, T].
+
+        mcore-specific wrapper (no counterpart in transformers): TP needs no
+        special handling (replicated weights, full-width hidden between
+        blocks); under SP/CP the shard is gathered to the full sequence
+        (undoing SP first, then the CP zigzag), ``compute`` runs on the full
+        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
+        if not (sp_on or cp_on):
+            return self._forward_impl(hidden_states, input_ids, packed_seq_params)
+
+        thd = packed_seq_params is not None and getattr(packed_seq_params, 'qkv_format', 'bshd') == 'thd'
+
+        # ---- gather the shard into the full sequence (SP shard is innermost) ----
+        if sp_on:
+            hidden_states = gather_from_sequence_parallel_region(
+                hidden_states, tensor_parallel_output_grad=False, group=self.pg_collection.tp)
+        if cp_on:
+            # For packed (thd) inputs the data pipeline (swift get_batch_on_this_cp_rank)
+            # zigzag-splits hidden/input_ids per sample via cu_seqlens_q while keeping
+            # packed_seq_params itself global, so the undo uses the same cu.
+            psp_for_cp = packed_seq_params if thd else None
+            hidden_states = reconstruct_tensor_cp(hidden_states, psp_for_cp, dim=0)
+            # The data pipeline may hand us either a CP-sharded or a full copy of
+            # input_ids; re-align only when the lengths disagree.
+            if input_ids.shape[-1] != hidden_states.shape[0]:
+                input_ids = reconstruct_tensor_cp(input_ids, psp_for_cp, dim=1)
+        elif input_ids.shape[-1] != hidden_states.shape[0]:
+            raise ValueError(
+                f'PLE input_ids length {input_ids.shape[-1]} does not match the gathered '
+                f'sequence length {hidden_states.shape[0]} under sequence parallelism; '
+                'gpt_model is expected to pass the full, unsharded input_ids.')
+
+        if thd:
+            # SP keeps one global cu_seqlens copy per rank; normalize padded/offset
+            # forms against the gathered total so cu indexes the full sequence.
+            # copy.copy (not dataclasses.replace) preserves dynamically attached
+            # fields such as `num_samples` that the data pipeline relies on.
+            cu = self._normalize_cu_seqlens(
+                getattr(packed_seq_params, 'cu_seqlens_q', None), hidden_states.shape[0])
+            psp = copy.copy(packed_seq_params)
+            psp.cu_seqlens_q = cu
+            packed_seq_params = psp
+
+        out = self._forward_impl(hidden_states, input_ids, packed_seq_params)
+
+        # ---- scatter the additive output back to the caller's shard ----
+        if cp_on:
+            out = split_cp_inputs(out, getattr(packed_seq_params, 'cu_seqlens_q', None) if thd else None, dim=0)
+        if sp_on:
+            out = scatter_to_sequence_parallel_region(out, group=self.pg_collection.tp)
+        return out
+
+    def _forward_impl(
+        self,
+        hidden_states: torch.Tensor,
+        input_ids: torch.Tensor,
+        packed_seq_params=None,
+    ) -> torch.Tensor:
+        """hidden_states: [s, b, nH] (bsh) or thd [T, 1, nH]; input_ids: [b, s] or [1, T]."""
+        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)
+            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]))
+            toks = input_ids.new_full((num_samples, max_len), self.ple_embedding.eos_token_id)
+            for i in range(num_samples):
+                start, end = int(cu[i]), int(cu[i + 1])
+                hid[i, :end - start] = hidden_states[start:end, 0]
+                toks[i, :end - start] = input_ids[0, start:end]
+            res = self.compute(hid, toks)
+            out = res.new_zeros((total, 1, res.shape[-1]))
+            for i in range(num_samples):
+                start, end = int(cu[i]), int(cu[i + 1])
+                out[start:end, 0] = res[i, :end - start]
+            return out
+        else:
+            # [s, b, nH] -> [b, s, nH]; input_ids [b, s]
+            hid = hidden_states.transpose(0, 1)
+            res = self.compute(hid, input_ids)
+            return res.transpose(0, 1).contiguous()

From ad958ec5dc30c4a203916499bac30a767cb6494a Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Wed, 26 Aug 2026 21:20:18 +0800
Subject: [PATCH 03/10] update

---
 src/mcore_bridge/config/parser.py             |   2 +-
 src/mcore_bridge/model/constant.py            |   2 -
 src/mcore_bridge/model/modules/qsa_indexer.py | 148 ++++++++++++++++++
 3 files changed, 149 insertions(+), 3 deletions(-)
 create mode 100644 src/mcore_bridge/model/modules/qsa_indexer.py

diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py
index 1641447f..ef554700 100644
--- a/src/mcore_bridge/config/parser.py
+++ b/src/mcore_bridge/config/parser.py
@@ -261,7 +261,7 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
         if use_mcore_gdn:
             res['experimental_attention_variant'] = 'gated_delta_net'
         res.setdefault('linear_attention_freq', 4)
-    elif hf_model_type in {'qwen4_exp'}:
+    elif hf_model_type == 'qwen4_exp':
         use_mcore_gdn = get_env_args('USE_MCORE_GDN', bool, True)
         res['layernorm_zero_centered_gamma'] = True
         res['attention_output_gate'] = True
diff --git a/src/mcore_bridge/model/constant.py b/src/mcore_bridge/model/constant.py
index 008f7da9..1296f342 100644
--- a/src/mcore_bridge/model/constant.py
+++ b/src/mcore_bridge/model/constant.py
@@ -25,8 +25,6 @@ class MLLMModelType:
     qwen3_omni = 'qwen3_omni'
     qwen3_asr = 'qwen3_asr'
     qwen3_5 = 'qwen3_5'
-    # TODO: confirm and remove one 
-    qwen3_8_flash_next = 'qwen3_8_flash_next'
     qwen4_exp = 'qwen4_exp'
     ovis2_5 = 'ovis2_5'
 
diff --git a/src/mcore_bridge/model/modules/qsa_indexer.py b/src/mcore_bridge/model/modules/qsa_indexer.py
new file mode 100644
index 00000000..967fa506
--- /dev/null
+++ b/src/mcore_bridge/model/modules/qsa_indexer.py
@@ -0,0 +1,148 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import math
+import torch
+from torch import nn
+
+class Qwen4ExpTextRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6, dtype=None, sequence_parallel: bool = False):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(dim, dtype=dtype))
+        # Replicated across TP; reduce grads across TP when SP is on.
+        setattr(self.weight, 'sequence_parallel', sequence_parallel)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = x.float()
+        out = out * torch.rsqrt(out.pow(2).mean(-1, keepdim=True) + self.eps)
+        # zero-centered: (1 + w), and Qwen4ExpText casts after scaling
+        return (out * (1.0 + self.weight.float())).type_as(x)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., :x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+class QSAIndexer(nn.Module):
+    # refer: transformers Qwen4ExpTextQSAIndexer
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.index_n_heads = config.indexer_n_heads
+        self.index_kv_heads = config.indexer_kv_heads
+        self.index_head_dim = config.indexer_head_dim
+        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,
+            bias=False,
+            dtype=config.params_dtype)
+        self.q_layernorm = Qwen4ExpTextRMSNorm(
+            self.index_head_dim, eps=config.layernorm_epsilon, dtype=config.params_dtype,
+            sequence_parallel=config.sequence_parallel)
+        self.k_layernorm = Qwen4ExpTextRMSNorm(
+            self.index_head_dim, eps=config.layernorm_epsilon, dtype=config.params_dtype,
+            sequence_parallel=config.sequence_parallel)
+        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`.')
+
+    @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.
+
+        Args:
+            hidden_states: ``[s, b, h]`` (mcore layout), pre-attention input --
+                the same tensor the reference indexer consumes.
+            freqs: mcore rotary frequencies ``[s, 1, 1, rot_dim]``. mcore stores
+                angles rather than cos/sin, so they are materialized here the way
+                ``_patch_apply_rotary_pos_emb`` does (``cos(freqs) * mscale``),
+                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).
+        """
+        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)
+        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]
+        q = q.view(s, b, self.index_n_heads, self.index_head_dim).permute(1, 0, 2, 3)
+        raw_keys = token_k.view(s, b, self.index_kv_heads, self.index_head_dim).permute(1, 0, 2, 3).squeeze(2)
+        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)
+        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)
+
+        # queries rotate at their own position: cos [s, rot] -> [1, s, 1, rot]
+        q = apply_rope(q, cos[None, :, None, :], sin[None, :, None, :])
+
+        # ---- pool every block once (shared across queries) ----
+        usable = max_blocks * R
+        key_groups = raw_keys[:, :usable].view(b, max_blocks, R, self.index_head_dim)
+        pooled = key_groups.float().mean(dim=2).to(raw_keys.dtype)
+        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]
+
+        # ---- score all (query, block) pairs ----
+        scores = torch.einsum('bqhd,bkd->bqhk', q.float(), block_keys.float())
+        scores = torch.relu(scores).sum(dim=2) / math.sqrt(self.index_head_dim)  # [b, s, nb]
+
+        # ---- restrict to blocks fully inside the causal prefix ----
+        n_blocks = (torch.arange(s, device=device) + 1) // R                      # [s]
+        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
+        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)
+
+        allowed = torch.zeros((b, s, s + 1), dtype=torch.bool, device=device)
+        allowed.scatter_(-1, torch.where(keep_tok, tok, torch.full_like(tok, s)).long(), True)
+        allowed = allowed[..., :s]
+        # tail: visible tokens after the last complete block are always attended
+        pos = torch.arange(s, device=device)
+        tail = (pos[None, :] >= (n_blocks * R)[:, None]) & (pos[None, :] <= pos[:, None])
+        allowed |= tail[None]
+        # causal safety net (the block expansion never crosses it, but keep the
+        # invariant explicit so a future layout change fails loudly instead of
+        # silently attending to the future)
+        allowed &= (pos[None, :] <= pos[:, None])[None]
+        return ~allowed.unsqueeze(1)                       # True == masked out

From 7390ae830ab7f58e4e7420bc75020bb43fdc2d48 Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Wed, 26 Aug 2026 21:28:35 +0800
Subject: [PATCH 04/10] update

---
 src/mcore_bridge/model/gpts/qwen4_exp.py                 | 1 +
 src/mcore_bridge/model/modules/hyper_connection_gated.py | 1 -
 2 files changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py
index a3052903..139950c4 100644
--- a/src/mcore_bridge/model/gpts/qwen4_exp.py
+++ b/src/mcore_bridge/model/gpts/qwen4_exp.py
@@ -133,6 +133,7 @@ def _warn_qsa_fallback_once(self, reason: str) -> None:
 
     def _qsa_select_mask(self, hidden_states, attn_kwargs):
         # return None means full attention
+        # TODO: support padding_free & cp
         indexer = getattr(self.self_attention, 'indexer', None)
         if indexer is None:
             return None
diff --git a/src/mcore_bridge/model/modules/hyper_connection_gated.py b/src/mcore_bridge/model/modules/hyper_connection_gated.py
index faeb1066..c8b9a723 100644
--- a/src/mcore_bridge/model/modules/hyper_connection_gated.py
+++ b/src/mcore_bridge/model/modules/hyper_connection_gated.py
@@ -39,7 +39,6 @@ def extra_repr(self):
 
 class Qwen4ExpTextGatedResidual(nn.Module):
     def __init__(self, config, use_combine: bool = True):
-        # Mirrors transformers `Qwen4ExpTextGatedResidual.__init__` (config-driven).
         super().__init__()
         self.hc_count = config.hc_count
         self.hidden_size = config.hidden_size

From 38d0e774d314221aa832a10832624b6a6a48e1ff Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Wed, 26 Aug 2026 21:42:07 +0800
Subject: [PATCH 05/10] lint

---
 src/mcore_bridge/model/gpts/qwen4_exp.py      | 16 +++--
 src/mcore_bridge/model/mm_gpts/qwen3_vl.py    |  6 +-
 .../model/modules/hyper_connection_gated.py   | 28 ++++-----
 src/mcore_bridge/model/modules/ple.py         | 58 ++++++++++---------
 src/mcore_bridge/model/modules/qsa_indexer.py | 18 ++++--
 5 files changed, 64 insertions(+), 62 deletions(-)

diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py
index 139950c4..47ffe0a1 100644
--- a/src/mcore_bridge/model/gpts/qwen4_exp.py
+++ b/src/mcore_bridge/model/gpts/qwen4_exp.py
@@ -18,7 +18,7 @@
 
 from mcore_bridge.utils import get_local_layer_specs, get_logger
 
-from ..modules import (GatedDeltaNet, Qwen4ExpTextGatedResidual, Qwen4ExpTextPLELayer, QSAIndexer, TransformerBlock,
+from ..modules import (GatedDeltaNet, QSAIndexer, Qwen4ExpTextGatedResidual, Qwen4ExpTextPLELayer, TransformerBlock,
                        TransformerLayer)
 from ..register import ModelLoader
 from .qwen3_next import Qwen3NextBridge, Qwen3NextRMSNorm, Qwen3NextSelfAttention
@@ -54,8 +54,8 @@ def __init__(self, config, submodules, layer_number: int = 1, **kwargs):
         super().__init__(config, submodules, layer_number, **kwargs)
         self.ple = None
         if self.layer_number in config.ple_layer_ids:
-            self.ple = Qwen4ExpTextPLELayer(config, config.ple_layer_ids.index(self.layer_number),
-                                            pg_collection=self.pg_collection)
+            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:
             self.self_attention.indexer = QSAIndexer(config)
@@ -206,8 +206,8 @@ def _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool
                 dist.all_reduce(is_moe, group=self.pp_group)
         if is_moe:
             hf_state_dict.update(
-                self._set_moe_state(mg_mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', layer_idx, to_mcore,
-                                    is_mtp=is_mtp))
+                self._set_moe_state(
+                    mg_mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', layer_idx, to_mcore, is_mtp=is_mtp))
         else:
             hf_state_dict.update(
                 self._set_mlp_state(mg_mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', layer_idx, to_mcore))
@@ -265,8 +265,7 @@ def _set_ple_ngram_embedding(self, ple, hf_state_dict, to_mcore: bool, pp_src_ra
         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)
+        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:
@@ -406,8 +405,7 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
         # it can be exercised; QSA layers run dense attention, which mcore's
         # attention already supports under CP.
         if getattr(config, 'mtp_num_layers', None):
-            raise NotImplementedError(
-                'Qwen4-Exp MTP is not supported yet')
+            raise NotImplementedError('Qwen4-Exp MTP is not supported yet')
         moe_spec = get_gpt_layer_with_transformer_engine_spec(
             num_experts=config.num_moe_experts,
             moe_grouped_gemm=config.moe_grouped_gemm,
diff --git a/src/mcore_bridge/model/mm_gpts/qwen3_vl.py b/src/mcore_bridge/model/mm_gpts/qwen3_vl.py
index bb594714..a6bb82d2 100644
--- a/src/mcore_bridge/model/mm_gpts/qwen3_vl.py
+++ b/src/mcore_bridge/model/mm_gpts/qwen3_vl.py
@@ -80,8 +80,7 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config):
                 deepstack_visual_embeds = visual_res.deepstack_features
             else:
                 image_embeds, deepstack_visual_embeds = visual_res
-            deepstack_visual_embeds = (torch.stack(deepstack_visual_embeds, dim=0)
-                                       if deepstack_visual_embeds else None)
+            deepstack_visual_embeds = (torch.stack(deepstack_visual_embeds, dim=0) if deepstack_visual_embeds else None)
             inputs_embeds = inputs_embeds + image_embeds.mean().to(device=inputs_embeds.device) * 0.
             visual_pos_masks = None
         else:
@@ -138,8 +137,7 @@ def _get_inputs_embeds(self, inputs_embeds, inputs, visual, hf_config):
                     embed_joint[video_mask_joint, :] = vid_embed
                     deepstack_visual_embeds.append(embed_joint)
 
-            deepstack_visual_embeds = (torch.stack(deepstack_visual_embeds, dim=0)
-                                       if deepstack_visual_embeds else None)
+            deepstack_visual_embeds = (torch.stack(deepstack_visual_embeds, dim=0) if deepstack_visual_embeds else None)
             visual_pos_masks = visual_pos_masks.transpose(0, 1)
             # compat cp
             if self.config.context_parallel_size > 1:
diff --git a/src/mcore_bridge/model/modules/hyper_connection_gated.py b/src/mcore_bridge/model/modules/hyper_connection_gated.py
index c8b9a723..78d373e7 100644
--- a/src/mcore_bridge/model/modules/hyper_connection_gated.py
+++ b/src/mcore_bridge/model/modules/hyper_connection_gated.py
@@ -6,12 +6,8 @@
 
 
 class Qwen4ExpTextGroupedRMSNorm(nn.Module):
-    def __init__(self,
-                 dim: int,
-                 group_size: int,
-                 eps: float = 1e-6,
-                 dtype=None,
-                 sequence_parallel: bool = False):
+
+    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))
@@ -38,6 +34,7 @@ def extra_repr(self):
 
 
 class Qwen4ExpTextGatedResidual(nn.Module):
+
     def __init__(self, config, use_combine: bool = True):
         super().__init__()
         self.hc_count = config.hc_count
@@ -45,14 +42,12 @@ def __init__(self, config, use_combine: bool = True):
         hc_hidden_size = self.hc_count * self.hidden_size
         # (transformers: Qwen4ExpTextRMSNorm(hc_hidden_size, group_size=self.hidden_size,
         #  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.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
         # mcore-specific: SP grad flag on the replicated weights.
         for param in self.parameters():
             setattr(param, 'sequence_parallel', config.sequence_parallel)
@@ -66,10 +61,9 @@ def forward(self, hyper_input: torch.Tensor) -> Tuple[torch.Tensor, ...]:
         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)
+        mixed_input = (input_mix_weight * hyper_input_normed.unflatten(-1,
+                                                                       (self.hc_count, self.hidden_size))).mean(dim=-2)
         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)
         return mixed_input, hyper_input, injection_weights
-
diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py
index 23bca0e7..09b34007 100644
--- a/src/mcore_bridge/model/modules/ple.py
+++ b/src/mcore_bridge/model/modules/ple.py
@@ -81,6 +81,7 @@ def _nth_prime_after(start: int, count: int) -> int:
 
 
 class Qwen4ExpTextNGramEmbedding(nn.Module):
+
     def __init__(self, config, ple_layer_index: int):
         super().__init__()
         self.ngram_size = config.ngram_size
@@ -124,10 +125,9 @@ def __init__(self, config, ple_layer_index: int):
             self.head_vocab_sizes.append(size)
             self.head_offsets.append(self.total_vocab_size)
             self.total_vocab_size += size
-        self.register_buffer('ngram_heads_vocab_sizes', torch.tensor(self.head_vocab_sizes, dtype=torch.long),
-                             persistent=True)
-        self.register_buffer('ngram_heads_offsets', torch.tensor(self.head_offsets, dtype=torch.long),
-                             persistent=True)
+        self.register_buffer(
+            'ngram_heads_vocab_sizes', torch.tensor(self.head_vocab_sizes, dtype=torch.long), persistent=True)
+        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).
@@ -227,15 +227,24 @@ def __init__(self, config, ple_layer_index: int, pg_collection=None):
         self.value_proj = nn.Linear(ple_embed_dim, self.hidden_size, bias=False, dtype=config.params_dtype)
         # 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(hc_hidden_size, group_size=self.hidden_size,
-                                                   eps=config.layernorm_epsilon, dtype=config.params_dtype,
-                                                   sequence_parallel=config.sequence_parallel)
-        self.norm_query = Qwen4ExpTextGroupedRMSNorm(hc_hidden_size, group_size=self.hidden_size,
-                                                     eps=config.layernorm_epsilon, dtype=config.params_dtype,
-                                                     sequence_parallel=config.sequence_parallel)
-        self.norm_conv = Qwen4ExpTextGroupedRMSNorm(hc_hidden_size, group_size=self.hidden_size,
-                                                    eps=config.layernorm_epsilon, dtype=config.params_dtype,
-                                                    sequence_parallel=config.sequence_parallel)
+        self.norm_key = Qwen4ExpTextGroupedRMSNorm(
+            hc_hidden_size,
+            group_size=self.hidden_size,
+            eps=config.layernorm_epsilon,
+            dtype=config.params_dtype,
+            sequence_parallel=config.sequence_parallel)
+        self.norm_query = Qwen4ExpTextGroupedRMSNorm(
+            hc_hidden_size,
+            group_size=self.hidden_size,
+            eps=config.layernorm_epsilon,
+            dtype=config.params_dtype,
+            sequence_parallel=config.sequence_parallel)
+        self.norm_conv = Qwen4ExpTextGroupedRMSNorm(
+            hc_hidden_size,
+            group_size=self.hidden_size,
+            eps=config.layernorm_epsilon,
+            dtype=config.params_dtype,
+            sequence_parallel=config.sequence_parallel)
         self.conv1d = nn.Conv1d(
             hc_hidden_size,
             hc_hidden_size,
@@ -298,10 +307,9 @@ def _normalize_cu_seqlens(cu: Optional[torch.Tensor], total: int) -> Optional[to
             elif int(cu[-1]) - int(cu[0]) == total:
                 cu = cu - cu[0]
         if not (cu.numel() > 0 and int(cu[-1]) == total):
-            raise ValueError(
-                f'PLE cannot align cu_seqlens (last={int(cu[-1]) if cu.numel() else None}, '
-                f'first={int(cu[0]) if cu.numel() else None}) with the gathered sequence '
-                f'length {total} under sequence parallelism.')
+            raise ValueError(f'PLE cannot align cu_seqlens (last={int(cu[-1]) if cu.numel() else None}, '
+                             f'first={int(cu[0]) if cu.numel() else None}) with the gathered sequence '
+                             f'length {total} under sequence parallelism.')
         return cu
 
     def forward(
@@ -318,9 +326,9 @@ def forward(
         (undoing SP first, then the CP zigzag), ``compute`` runs on the full
         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)
+        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
         if not (sp_on or cp_on):
             return self._forward_impl(hidden_states, input_ids, packed_seq_params)
@@ -342,18 +350,16 @@ def forward(
             if input_ids.shape[-1] != hidden_states.shape[0]:
                 input_ids = reconstruct_tensor_cp(input_ids, psp_for_cp, dim=1)
         elif input_ids.shape[-1] != hidden_states.shape[0]:
-            raise ValueError(
-                f'PLE input_ids length {input_ids.shape[-1]} does not match the gathered '
-                f'sequence length {hidden_states.shape[0]} under sequence parallelism; '
-                'gpt_model is expected to pass the full, unsharded input_ids.')
+            raise ValueError(f'PLE input_ids length {input_ids.shape[-1]} does not match the gathered '
+                             f'sequence length {hidden_states.shape[0]} under sequence parallelism; '
+                             'gpt_model is expected to pass the full, unsharded input_ids.')
 
         if thd:
             # SP keeps one global cu_seqlens copy per rank; normalize padded/offset
             # forms against the gathered total so cu indexes the full sequence.
             # copy.copy (not dataclasses.replace) preserves dynamically attached
             # fields such as `num_samples` that the data pipeline relies on.
-            cu = self._normalize_cu_seqlens(
-                getattr(packed_seq_params, 'cu_seqlens_q', None), hidden_states.shape[0])
+            cu = self._normalize_cu_seqlens(getattr(packed_seq_params, 'cu_seqlens_q', None), hidden_states.shape[0])
             psp = copy.copy(packed_seq_params)
             psp.cu_seqlens_q = cu
             packed_seq_params = psp
diff --git a/src/mcore_bridge/model/modules/qsa_indexer.py b/src/mcore_bridge/model/modules/qsa_indexer.py
index 967fa506..193bcbbf 100644
--- a/src/mcore_bridge/model/modules/qsa_indexer.py
+++ b/src/mcore_bridge/model/modules/qsa_indexer.py
@@ -3,7 +3,9 @@
 import torch
 from torch import nn
 
+
 class Qwen4ExpTextRMSNorm(nn.Module):
+
     def __init__(self, dim: int, eps: float = 1e-6, dtype=None, sequence_parallel: bool = False):
         super().__init__()
         self.eps = eps
@@ -41,10 +43,14 @@ def __init__(self, config):
             bias=False,
             dtype=config.params_dtype)
         self.q_layernorm = Qwen4ExpTextRMSNorm(
-            self.index_head_dim, eps=config.layernorm_epsilon, dtype=config.params_dtype,
+            self.index_head_dim,
+            eps=config.layernorm_epsilon,
+            dtype=config.params_dtype,
             sequence_parallel=config.sequence_parallel)
         self.k_layernorm = Qwen4ExpTextRMSNorm(
-            self.index_head_dim, eps=config.layernorm_epsilon, dtype=config.params_dtype,
+            self.index_head_dim,
+            eps=config.layernorm_epsilon,
+            dtype=config.params_dtype,
             sequence_parallel=config.sequence_parallel)
         setattr(self.index_qk_proj.weight, 'sequence_parallel', config.sequence_parallel)
 
@@ -123,14 +129,14 @@ def apply_rope(t, cos_, sin_):
         scores = torch.relu(scores).sum(dim=2) / math.sqrt(self.index_head_dim)  # [b, s, nb]
 
         # ---- restrict to blocks fully inside the causal prefix ----
-        n_blocks = (torch.arange(s, device=device) + 1) // R                      # [s]
+        n_blocks = (torch.arange(s, device=device) + 1) // R  # [s]
         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
+        top_blocks = scores.topk(k, dim=-1).indices  # [b, s, k]
+        keep = top_blocks < n_blocks[None, :, None]  # drop the -inf padding slots
         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)
 
@@ -145,4 +151,4 @@ def apply_rope(t, cos_, sin_):
         # invariant explicit so a future layout change fails loudly instead of
         # silently attending to the future)
         allowed &= (pos[None, :] <= pos[:, None])[None]
-        return ~allowed.unsqueeze(1)                       # True == masked out
+        return ~allowed.unsqueeze(1)  # True == masked out

From a69f5cc8bc8cf09f9a65733aa171a8f9085a9ae3 Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Wed, 26 Aug 2026 22:35:52 +0800
Subject: [PATCH 06/10] lint

---
 src/mcore_bridge/model/gpts/qwen4_exp.py | 1 -
 src/mcore_bridge/model/modules/ple.py    | 4 ++--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py
index 47ffe0a1..7ce61f87 100644
--- a/src/mcore_bridge/model/gpts/qwen4_exp.py
+++ b/src/mcore_bridge/model/gpts/qwen4_exp.py
@@ -434,7 +434,6 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None):
         moe_pattern = self._get_moe_layer_pattern()
         layer_specs = []
         for layer_idx, is_linear_attention in enumerate(config.linear_attention_freq):
-            from copy import deepcopy
             layer_spec = deepcopy(moe_spec if moe_pattern[layer_idx] else dense_spec)
             if is_linear_attention:
                 layer_spec.submodules.self_attention = deepcopy(gdn_spec)
diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py
index 09b34007..1121bd2b 100644
--- a/src/mcore_bridge/model/modules/ple.py
+++ b/src/mcore_bridge/model/modules/ple.py
@@ -255,9 +255,9 @@ def __init__(self, config, ple_layer_index: int, pg_collection=None):
             dtype=config.params_dtype,
         )
         nn.init.zeros_(self.conv1d.weight)
+        replicated_prefixes = ('key_proj', 'value_proj', 'norm_key', 'norm_query', 'norm_conv', 'conv1d')
         for name, param in self.named_parameters():
-            if name.startswith(('key_proj', 'value_proj')) or name.endswith('norm.weight') or name.startswith(
-                ('norm_key', 'norm_query', 'norm_conv')) or name.startswith('conv1d'):
+            if name.startswith(replicated_prefixes) or name.endswith('norm.weight'):
                 # Replicated across TP; reduce grads across TP when SP is on.
                 setattr(param, 'sequence_parallel', config.sequence_parallel)
 

From ec98088228f5375f87e589c04d281e0a43effa4c Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Thu, 3 Sep 2026 13:35:10 +0800
Subject: [PATCH 07/10] update

---
 src/mcore_bridge/config/model_config.py       |   1 -
 src/mcore_bridge/config/parser.py             |   2 -
 src/mcore_bridge/model/gpts/qwen3_next.py     |   6 +-
 src/mcore_bridge/model/gpts/qwen4_exp.py      | 286 ++++----
 src/mcore_bridge/model/mm_gpts/qwen3_5.py     |   6 +-
 src/mcore_bridge/model/modules/__init__.py    |   1 +
 .../model/modules/hyper_connection_gated.py   |  61 +-
 .../model/modules/kernels/__init__.py         |  10 +
 .../model/modules/kernels/ple_kernels.py      | 463 +++++++++++++
 .../modules/kernels/qsa_block_sparse_attn.py  | 616 ++++++++++++++++++
 .../model/modules/kernels/qsa_kernels.py      | 286 ++++++++
 src/mcore_bridge/model/modules/ple.py         | 219 ++++++-
 src/mcore_bridge/model/modules/qsa_indexer.py | 279 +++++++-
 tests/test_qsa_indexer.py                     |  67 ++
 tests/test_qwen4_exp_units.py                 | 507 ++++++++++++++
 15 files changed, 2620 insertions(+), 190 deletions(-)
 create mode 100644 src/mcore_bridge/model/modules/kernels/__init__.py
 create mode 100644 src/mcore_bridge/model/modules/kernels/ple_kernels.py
 create mode 100644 src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py
 create mode 100644 src/mcore_bridge/model/modules/kernels/qsa_kernels.py
 create mode 100644 tests/test_qsa_indexer.py
 create mode 100644 tests/test_qwen4_exp_units.py

diff --git a/src/mcore_bridge/config/model_config.py b/src/mcore_bridge/config/model_config.py
index c5231404..121caea4 100644
--- a/src/mcore_bridge/config/model_config.py
+++ b/src/mcore_bridge/config/model_config.py
@@ -207,7 +207,6 @@ class ModelConfig(TransformerConfig):
     ngram_vocab_size_base: Optional[int] = None
     make_ngram_vocab_size_divisible_by: Optional[int] = None
     split_ngram_parts: Optional[int] = None
-    ple_seed: Optional[int] = None
     eos_token_id: Optional[int] = None
     indexer_n_heads: Optional[int] = None
     indexer_kv_heads: Optional[int] = None
diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py
index ef554700..53f589d2 100644
--- a/src/mcore_bridge/config/parser.py
+++ b/src/mcore_bridge/config/parser.py
@@ -276,8 +276,6 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
         res['linear_attention_freq'] = f"[{','.join(linear_pattern)}]"
         if res.get('num_moe_experts'):
             res['moe_layer_freq'] = f"[{','.join(['1'] * num_layers)}]"
-        # seed is hardcoded in transformers, not in config
-        res['ple_seed'] = int(getattr(text_config, 'seed', 1234))
         eos_token_id = getattr(text_config, 'eos_token_id', None)
         if eos_token_id is not None:
             res['eos_token_id'] = eos_token_id
diff --git a/src/mcore_bridge/model/gpts/qwen3_next.py b/src/mcore_bridge/model/gpts/qwen3_next.py
index 69df8f96..6fcf1ea5 100644
--- a/src/mcore_bridge/model/gpts/qwen3_next.py
+++ b/src/mcore_bridge/model/gpts/qwen3_next.py
@@ -470,10 +470,10 @@ def forward(self, hidden_states: torch.Tensor, **kwargs):
         # Note: for packed inputs, we do not perform padding_free unpadding.
         # Doing so would allow different sequences to see each other; for efficiency we keep this implementation.
         if thd_format:
+            max_seqlen_q = int(packed_seq_params.max_seqlen_q)
             new_hidden_states = hidden_states.new_zeros(
-                (packed_seq_params.num_samples, packed_seq_params.max_seqlen_q.item(), hidden_states.shape[-1]))
-            attention_mask = hidden_states.new_zeros(
-                (packed_seq_params.num_samples, packed_seq_params.max_seqlen_q.item()), dtype=torch.bool)
+                (packed_seq_params.num_samples, max_seqlen_q, hidden_states.shape[-1]))
+            attention_mask = hidden_states.new_zeros((packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool)
             cu_seqlens_q = packed_seq_params.cu_seqlens_q
             for i in range(packed_seq_params.num_samples):
                 start, end = cu_seqlens_q[i], cu_seqlens_q[i + 1]
diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py
index 7ce61f87..e62e6609 100644
--- a/src/mcore_bridge/model/gpts/qwen4_exp.py
+++ b/src/mcore_bridge/model/gpts/qwen4_exp.py
@@ -5,21 +5,24 @@
 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.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
 from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
+from megatron.core.packed_seq_params import PackedSeqParams
+
 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 +45,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 +60,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 +89,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 +136,142 @@ 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 +279,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 +339,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 +355,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 +370,20 @@ 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 +392,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 +435,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/mm_gpts/qwen3_5.py b/src/mcore_bridge/model/mm_gpts/qwen3_5.py
index 8ba53783..b673bf73 100644
--- a/src/mcore_bridge/model/mm_gpts/qwen3_5.py
+++ b/src/mcore_bridge/model/mm_gpts/qwen3_5.py
@@ -40,10 +40,10 @@ def forward(self, hidden_states: torch.Tensor, **kwargs):
         # Note: for packed inputs, we do not perform padding_free unpadding.
         # Doing so would allow different sequences to see each other; for efficiency we keep this implementation.
         if thd_format:
+            max_seqlen_q = int(packed_seq_params.max_seqlen_q)
             new_hidden_states = hidden_states.new_zeros(
-                (packed_seq_params.num_samples, packed_seq_params.max_seqlen_q.item(), hidden_states.shape[-1]))
-            attention_mask = hidden_states.new_zeros(
-                (packed_seq_params.num_samples, packed_seq_params.max_seqlen_q.item()), dtype=torch.bool)
+                (packed_seq_params.num_samples, max_seqlen_q, hidden_states.shape[-1]))
+            attention_mask = hidden_states.new_zeros((packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool)
             cu_seqlens_q = packed_seq_params.cu_seqlens_q
             for i in range(packed_seq_params.num_samples):
                 start, end = cu_seqlens_q[i], cu_seqlens_q[i + 1]
diff --git a/src/mcore_bridge/model/modules/__init__.py b/src/mcore_bridge/model/modules/__init__.py
index 7f5f2a2e..d8d613c1 100644
--- a/src/mcore_bridge/model/modules/__init__.py
+++ b/src/mcore_bridge/model/modules/__init__.py
@@ -9,6 +9,7 @@
 from .multi_latent_attention import MLASelfAttention
 from .ple import Qwen4ExpTextNGramEmbedding, Qwen4ExpTextPLELayer
 from .qsa_indexer import QSAIndexer
+from .kernels import QSASparseCoreAttention, qsa_sparse_supported
 from .topk_router import TopKRouter
 from .transformer_block import TransformerBlock
 from .transformer_layer import TransformerLayer
diff --git a/src/mcore_bridge/model/modules/hyper_connection_gated.py b/src/mcore_bridge/model/modules/hyper_connection_gated.py
index 78d373e7..5c94b58f 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 00000000..21391734
--- /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 00000000..47e1d3f3
--- /dev/null
+++ b/src/mcore_bridge/model/modules/kernels/ple_kernels.py
@@ -0,0 +1,463 @@
+# 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/CUDA, 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.
+    """
+    if not HAVE_TRITON or not torch.cuda.is_available():
+        return None
+    if host_table.dtype != torch.bfloat16 or ids.device.type != 'cuda':
+        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 not hc_state.is_cuda:
+        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 00000000..bff3882c
--- /dev/null
+++ b/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py
@@ -0,0 +1,616 @@
+# 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 00000000..1a6bf1ea
--- /dev/null
+++ b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py
@@ -0,0 +1,286 @@
+# 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 not q.is_cuda:
+        raise RuntimeError(
+            'QSA sparse attention requires triton and CUDA tensors '
+            f'(HAVE_TRITON={HAVE_TRITON}, q.is_cuda={q.is_cuda}). 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 1121bd2b..74dd75b8 100644
--- a/src/mcore_bridge/model/modules/ple.py
+++ b/src/mcore_bridge/model/modules/ple.py
@@ -3,20 +3,37 @@
 import math
 import torch
 import torch.nn.functional as F
+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 megatron.core import parallel_state
+
+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
 _SPLITMIX_M1 = 0xBF58476D1CE4E5B9
 _SPLITMIX_M2 = 0x94D049BB133111EB
 _PRIME_1 = 10007
+# Mirrors Qwen4ExpTextConfig.seed's default (transformers still reads it off the
+# config; no released Qwen4-Exp checkpoint actually sets it).
+_PLE_SEED = 1234
+
+
+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:
@@ -99,19 +116,21 @@ 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)
-        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}).')
+        eos_token_id = config.eos_token_id
+        split_ngram_parts = config.split_ngram_parts
+        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 matches
+        # Qwen4ExpTextConfig.seed's default; no released checkpoint sets it, and 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))
+                                               _PLE_SEED)
         self.register_buffer('layer_multipliers', torch.tensor(multipliers, dtype=torch.long), persistent=True)
 
         # Per-head prime table sizes/offsets (checkpoint-persistent), named as
@@ -130,13 +149,82 @@ 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.is_cuda:
+            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 193bcbbf..8a9d38eb 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,152 @@ 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 00000000..22f9b3c4
--- /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 00000000..54527531
--- /dev/null
+++ b/tests/test_qwen4_exp_units.py
@@ -0,0 +1,507 @@
+# 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)'

From 84f9635f4bce7e61868f6a3693ac87ef7f85eea1 Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Thu, 3 Sep 2026 18:15:29 +0800
Subject: [PATCH 08/10] lint

---
 src/mcore_bridge/model/gpts/qwen4_exp.py      |  69 ++++-----
 src/mcore_bridge/model/modules/__init__.py    |   2 +-
 .../model/modules/kernels/ple_kernels.py      | 146 +++++++++++++-----
 .../modules/kernels/qsa_block_sparse_attn.py  |  69 ++++-----
 .../model/modules/kernels/qsa_kernels.py      |  69 +++++----
 src/mcore_bridge/model/modules/ple.py         |   5 +-
 src/mcore_bridge/model/modules/qsa_indexer.py |   1 -
 tests/test_qwen4_exp_units.py                 |  17 +-
 8 files changed, 212 insertions(+), 166 deletions(-)

diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py
index e62e6609..5075dc28 100644
--- a/src/mcore_bridge/model/gpts/qwen4_exp.py
+++ b/src/mcore_bridge/model/gpts/qwen4_exp.py
@@ -7,14 +7,13 @@
 from copy import deepcopy
 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
 from megatron.core.transformer.transformer_block import TransformerBlockSubmodules
-from megatron.core.packed_seq_params import PackedSeqParams
-
 from transformers.utils import is_torch_npu_available
 from typing import List, Optional
 
@@ -163,42 +162,35 @@ def _qsa_select(self, hidden_states, attn_kwargs, position_ids=None):
         # 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.')
+            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.')
+            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.')
+            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)
+            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)
+            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):
+    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)
+            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:
             # TE's packed CP partition (thd_get_partitioned_indices) requires
@@ -222,10 +214,9 @@ def _qsa_select_indices_thd(self, hidden_states, rotary_pos_emb, packed_seq_para
         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.')
+                    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:
@@ -236,19 +227,17 @@ def _qsa_select_indices_thd(self, hidden_states, rotary_pos_emb, packed_seq_para
             # 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.')
+            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.')
+            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)
@@ -265,12 +254,10 @@ def _qsa_select_mask(self, hidden_states, attn_kwargs):
             return None
         rotary_pos_emb = attn_kwargs.get('rotary_pos_emb')
         if rotary_pos_emb is None:
-            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.')
+            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)
+            hidden_states = gather_from_sequence_parallel_region(hidden_states, tensor_parallel_output_grad=False)
         return indexer.selection_as_mask(hidden_states, rotary_pos_emb)
 
 
@@ -374,10 +361,8 @@ def _set_layer_ple(self, mg_layer, hf_state_dict, to_mcore: bool):
         # 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)
+        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, buf)
diff --git a/src/mcore_bridge/model/modules/__init__.py b/src/mcore_bridge/model/modules/__init__.py
index d8d613c1..fdba46b0 100644
--- a/src/mcore_bridge/model/modules/__init__.py
+++ b/src/mcore_bridge/model/modules/__init__.py
@@ -5,11 +5,11 @@
 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
 from .qsa_indexer import QSAIndexer
-from .kernels import QSASparseCoreAttention, qsa_sparse_supported
 from .topk_router import TopKRouter
 from .transformer_block import TransformerBlock
 from .transformer_layer import TransformerLayer
diff --git a/src/mcore_bridge/model/modules/kernels/ple_kernels.py b/src/mcore_bridge/model/modules/kernels/ple_kernels.py
index 47e1d3f3..9255f83f 100644
--- a/src/mcore_bridge/model/modules/kernels/ple_kernels.py
+++ b/src/mcore_bridge/model/modules/kernels/ple_kernels.py
@@ -18,7 +18,6 @@
 except Exception:  # pragma: no cover - triton absent
     HAVE_TRITON = False
 
-
 if HAVE_TRITON:
 
     @triton.jit
@@ -72,7 +71,7 @@ def gather_ple_rows(host_table, ids, row_start, row_end, out=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(),)](
+        _gather_ple_rows_from_pinned[(flat.numel(), )](
             host_table.data_ptr(),
             flat.contiguous(),
             out.view(-1, embedding_dim),
@@ -91,9 +90,12 @@ 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]
+        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]
+        gate_ptr,
+        rstdk_ptr,
+        rstdq_ptr,  # fp32 out [T, N]
         T,
         N: tl.constexpr,
         C: tl.constexpr,
@@ -136,11 +138,19 @@ def _ple_gate_fwd_kernel(
     @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]
+        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)
+        dwk_partial_ptr,
+        dwq_partial_ptr,  # fp32 out [T, N*C] (host sums over T)
         T,
         N: tl.constexpr,
         C: tl.constexpr,
@@ -196,7 +206,10 @@ def _ple_gate_bwd_kernel(
 
     @triton.jit(do_not_specialize=['T'])
     def _ple_norm_fwd_kernel(
-        x_ptr, w_ptr, out_ptr, rstd_ptr,
+        x_ptr,
+        w_ptr,
+        out_ptr,
+        rstd_ptr,
         T,
         N: tl.constexpr,
         C: tl.constexpr,
@@ -220,7 +233,11 @@ def _ple_norm_fwd_kernel(
 
     @triton.jit(do_not_specialize=['T'])
     def _ple_norm_bwd_kernel(
-        x_ptr, w_ptr, rstd_ptr, dout_ptr, dx_ptr,
+        x_ptr,
+        w_ptr,
+        rstd_ptr,
+        dout_ptr,
+        dx_ptr,
         T,
         N: tl.constexpr,
         C: tl.constexpr,
@@ -359,15 +376,27 @@ def forward(ctx, hc_state, key, value, wk, wq, wc, conv_w, n, eps, dilation, seq
             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)
+                _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)
+                _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()
@@ -376,19 +405,17 @@ def forward(ctx, hc_state, key, value, wk, wq, wc, conv_w, n, eps, dilation, seq
             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)
+                    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.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
+            (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
@@ -403,14 +430,26 @@ def backward(ctx, dout):
             _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)
+                _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)
+                _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)
@@ -418,16 +457,27 @@ def backward(ctx, dout):
             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)
+                    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)
+                _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)
@@ -436,10 +486,26 @@ def backward(ctx, dout):
             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)
+                _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)
@@ -448,8 +514,8 @@ def backward(ctx, dout):
             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):
+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``.
 
@@ -457,7 +523,5 @@ def ple_gate_conv_triton(hc_state, key, value, norm_key_w, norm_query_w, norm_co
     """
     if not HAVE_TRITON or not hc_state.is_cuda:
         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)
+    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
index bff3882c..f0b8586b 100644
--- a/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py
+++ b/src/mcore_bridge/model/modules/kernels/qsa_block_sparse_attn.py
@@ -96,8 +96,8 @@ def _qsa_bs_fwd_kernel(
     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)
+    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
@@ -128,13 +128,13 @@ def _qsa_bs_fwd_kernel(
             other=0.0,
         )
         s = tl.dot(q, tl.trans(k_tile)) * scale
-        s = tl.where(ok, s, float("-inf"))
+        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)
+        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))
+        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]
 
@@ -153,7 +153,7 @@ def _qsa_bs_fwd_kernel(
         out,
         mask=q_mask[:, None],
     )
-    lse = tl.where(m_i == float("-inf"), float("-inf"), m_i + tl.log(l_safe))
+    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)
 
 
@@ -203,12 +203,11 @@ def _qsa_bs_dq_kernel(
 
     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
-    )
+        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")
+    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)
@@ -233,20 +232,16 @@ def _qsa_bs_dq_kernel(
             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]
-        )
+        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
-        )
+            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
-        )
+            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])
@@ -311,11 +306,9 @@ def _qsa_bs_dkdv_kernel(
     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
-    )
+        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
-    )
+        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)
 
@@ -335,13 +328,11 @@ def _qsa_bs_dkdv_kernel(
             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]
-        )
+        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
@@ -357,9 +348,9 @@ def _qsa_bs_dkdv_kernel(
             )
             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"))
+            okh = ok & (lse[:, None] != float('-inf'))
 
-            lse_safe = tl.where(lse == float("-inf"), 0.0, lse)
+            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)
@@ -441,6 +432,7 @@ def compact(mat):
 
 
 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
@@ -602,9 +594,12 @@ def qsa_block_sparse_attention_triton(
     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:
+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)
diff --git a/src/mcore_bridge/model/modules/kernels/qsa_kernels.py b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py
index 1a6bf1ea..fc017a4b 100644
--- a/src/mcore_bridge/model/modules/kernels/qsa_kernels.py
+++ b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py
@@ -57,8 +57,7 @@ def _cp_query_global_positions(seq_len: int, cp_size: int, cp_rank: int, device)
     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:
+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;
@@ -89,7 +88,6 @@ def _cp_gathered_to_logical_order(seq_len: int, cp_size: int, device) -> torch.T
     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).
 
@@ -102,10 +100,9 @@ def qsa_sparse_attention_thd(q, k, v, indices, scale, block_size):
     note in qsa_block_sparse_attn.py.
     """
     if not HAVE_TRITON or not q.is_cuda:
-        raise RuntimeError(
-            'QSA sparse attention requires triton and CUDA tensors '
-            f'(HAVE_TRITON={HAVE_TRITON}, q.is_cuda={q.is_cuda}). This path is only '
-            'selected for packing (thd) or CP>1, where no dense fallback is correct.')
+        raise RuntimeError('QSA sparse attention requires triton and CUDA tensors '
+                           f'(HAVE_TRITON={HAVE_TRITON}, q.is_cuda={q.is_cuda}). 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]:
@@ -113,8 +110,7 @@ def qsa_sparse_attention_thd(q, k, v, indices, scale, block_size):
         # 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]}.')
+        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)
 
 
@@ -141,10 +137,9 @@ def qsa_sparse_attention(q, k, v, indices, scale, block_size):
     # 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.')
+        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:])
@@ -184,12 +179,18 @@ def __init__(self, core_attention, config, softmax_scale=None):
         # 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):
+            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
@@ -204,8 +205,14 @@ def forward(self, query, key, value, attention_mask, attn_mask_type=None,
                 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)
+            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
@@ -217,13 +224,13 @@ def _forward_cp(self, query, key, value, indices, scale, packed_seq_params):
         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)
+            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)])
+            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]
@@ -237,8 +244,7 @@ def _forward_cp(self, query, key, value, indices, scale, packed_seq_params):
         kv_reorder_t = kv_reorder
 
         def _gather_full(t):
-            g = gather_from_sequence_parallel_region(
-                t, tensor_parallel_output_grad=True, group=cp_group)
+            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)
@@ -251,9 +257,8 @@ def _gather_full(t):
         # 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)
+            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]
@@ -265,8 +270,7 @@ def _gather_full(t):
         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 = 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)
 
@@ -283,4 +287,3 @@ def _scatter_q_to_full(q, k, v, indices, rows):
         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 ddad1c25..e9ba94ee 100644
--- a/src/mcore_bridge/model/modules/ple.py
+++ b/src/mcore_bridge/model/modules/ple.py
@@ -3,6 +3,7 @@
 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,
@@ -10,8 +11,6 @@
 from torch import nn
 from typing import List, Optional
 
-from megatron.core import parallel_state
-
 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
@@ -118,7 +117,7 @@ def __init__(self, config, ple_layer_index: int):
         ple_seed = getattr(config, 'ple_seed', None)
         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 "
+                             '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 '
diff --git a/src/mcore_bridge/model/modules/qsa_indexer.py b/src/mcore_bridge/model/modules/qsa_indexer.py
index 8a9d38eb..080e28a4 100644
--- a/src/mcore_bridge/model/modules/qsa_indexer.py
+++ b/src/mcore_bridge/model/modules/qsa_indexer.py
@@ -372,4 +372,3 @@ def apply_rope(t, cos_, sin_):
         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_qwen4_exp_units.py b/tests/test_qwen4_exp_units.py
index 54527531..c7166b0f 100644
--- a/tests/test_qwen4_exp_units.py
+++ b/tests/test_qwen4_exp_units.py
@@ -1,6 +1,5 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import math
-
 import pytest
 import torch
 import torch.nn.functional as F
@@ -155,7 +154,7 @@ def test_indexer_matches_reference():
                 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"
+    assert ok, 'numerical check failed; see printed relL1/diff above'
 
 
 def test_indexer_is_forward_only():
@@ -241,7 +240,7 @@ def test_gated_residual_equivalence():
             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"
+    assert ok, 'numerical check failed; see printed relL1/diff above'
 
 
 # --------------------------------------------------------------------------
@@ -250,6 +249,7 @@ def test_gated_residual_equivalence():
 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)
@@ -262,12 +262,13 @@ def test_host_lookup_reference():
     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"
+    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
@@ -284,7 +285,7 @@ def test_host_lookup_tp_partition():
     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"
+    assert ok, 'numerical check failed; see printed relL1/diff above'
 
 
 # --------------------------------------------------------------------------
@@ -375,7 +376,7 @@ def test_ple_fused_kernel():
         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"
+    assert ok, 'numerical check failed; see printed relL1/diff above'
 
 
 def test_indexer_indices_match_reference():
@@ -420,7 +421,7 @@ def test_indexer_indices_match_reference():
             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"
+    assert ok, 'numerical check failed; see printed relL1/diff above'
 
 
 def test_indexer_indices_packed():
@@ -470,7 +471,7 @@ def test_indexer_indices_packed():
         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"
+    assert ok, 'numerical check failed; see printed relL1/diff above'
 
 
 def test_indexer_mrope_batch():

From 0216c4e16373c75692b1234ee5c753b1a75c554a Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Thu, 3 Sep 2026 23:04:33 +0800
Subject: [PATCH 09/10] relax kernel

---
 .../model/modules/kernels/ple_kernels.py          | 15 ++++++++++-----
 .../model/modules/kernels/qsa_kernels.py          |  6 +++---
 src/mcore_bridge/model/modules/ple.py             |  2 +-
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/src/mcore_bridge/model/modules/kernels/ple_kernels.py b/src/mcore_bridge/model/modules/kernels/ple_kernels.py
index 9255f83f..ef1f0c6f 100644
--- a/src/mcore_bridge/model/modules/kernels/ple_kernels.py
+++ b/src/mcore_bridge/model/modules/kernels/ple_kernels.py
@@ -51,8 +51,8 @@ def _gather_ple_rows_from_pinned(
 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/CUDA, or the table
-    is not bf16), so the caller can fall back to the torch path.
+    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.
@@ -60,9 +60,14 @@ def gather_ple_rows(host_table, ids, row_start, row_end, out=None):
         row_start / row_end: this rank's global row range.
         out: optional preallocated ``[*ids.shape, embedding_dim]`` bf16 device tensor.
     """
-    if not HAVE_TRITON or not torch.cuda.is_available():
+    # 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 or ids.device.type != 'cuda':
+    if host_table.dtype != torch.bfloat16:
         return None
     embedding_dim = host_table.shape[-1]
 
@@ -521,7 +526,7 @@ def ple_gate_conv_triton(hc_state, key, value, norm_key_w, norm_query_w, norm_co
 
     Returns ``None`` when the fast path is unavailable so callers fall back.
     """
-    if not HAVE_TRITON or not hc_state.is_cuda:
+    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_kernels.py b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py
index fc017a4b..95957175 100644
--- a/src/mcore_bridge/model/modules/kernels/qsa_kernels.py
+++ b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py
@@ -99,9 +99,9 @@ def qsa_sparse_attention_thd(q, k, v, indices, scale, block_size):
     different value silently changes which keys are attended -- see the contract
     note in qsa_block_sparse_attn.py.
     """
-    if not HAVE_TRITON or not q.is_cuda:
-        raise RuntimeError('QSA sparse attention requires triton and CUDA tensors '
-                           f'(HAVE_TRITON={HAVE_TRITON}, q.is_cuda={q.is_cuda}). This path is only '
+    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]}.')
diff --git a/src/mcore_bridge/model/modules/ple.py b/src/mcore_bridge/model/modules/ple.py
index e9ba94ee..e97a1a0b 100644
--- a/src/mcore_bridge/model/modules/ple.py
+++ b/src/mcore_bridge/model/modules/ple.py
@@ -416,7 +416,7 @@ 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
-        if use_ple_fused_kernel() and embeddings.is_cuda:
+        if use_ple_fused_kernel() and embeddings.device.type != 'cpu':
             fused = self._compute_fused(hidden_states, embeddings)
             if fused is not None:
                 return fused

From 4ff37ab74b97948f4c05bf5a33ce79327c058aa9 Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Sun, 6 Sep 2026 00:11:05 +0800
Subject: [PATCH 10/10] bump transformers

---
 requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index 6b9b856c..bebc374d 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