From d0842730c226dab6a89f2ea3bb5a22617fe98c73 Mon Sep 17 00:00:00 2001 From: ehooon Date: Thu, 3 Sep 2026 10:03:52 +0800 Subject: [PATCH] fix(qwen4_exp): fall back QSA to full attention under sequence parallelism With sequence_parallel=True and TP>1, the QSA indexer consumes SP-sharded hidden states (s/tp) while TE core attention runs on the all-gathered full sequence, so the [b, 1, s, s] selection mask mismatches the q/k length and the unfused path crashes on a logical_or shape error. Fall back QSA layers to full attention in this configuration, with a one-time warning that training beyond the indexer budget differs from sparse inference. --- src/mcore_bridge/model/gpts/qwen4_exp.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index 7ce61f8..5a40c3f 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -151,6 +151,16 @@ def _qsa_select_mask(self, hidden_states, attn_kwargs): 'back to full attention -- training will differ from sparse inference beyond the ' 'indexer budget.') return None + if getattr(self.config, 'sequence_parallel', False) and self.config.tensor_model_parallel_size > 1: + # Under SP the indexer would see SP-sharded hidden states (s/tp) while TE core + # attention runs on the all-gathered full sequence, so the [b, 1, s, s] selection + # mask would not match the q/k length (logical_or shape error in the unfused path). + self._warn_qsa_fallback_once( + 'sequence_parallel is enabled: the QSA indexer consumes SP-sharded hidden states ' + 'while core attention runs on the full sequence. 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