[Refactor] Fix AttentionMaskBuilder singleton and remove redundant pcp_prefill_mask (#4870)
## What this PR does / why we need it? This PR fixes the `AttentionMaskBuilder` singleton initialization issue introduced in PR #4779 and removes the unused `pcp_prefill_mask` field. ### Background After PR #4779 made `AttentionMaskBuilder` a singleton with `@singleton` decorator, the class constructor now requires a `device` parameter. However, two initialization sites were still using the old parameterless constructor, causing failures. ### Changes 1. **Fix singleton initialization** - Fixed `AttentionMaskBuilder()` → `AttentionMaskBuilder(self.device)` in `AscendMLAMetadataBuilder.__init__()` - Fixed `AttentionMaskBuilder()` → `AttentionMaskBuilder(self.device)` in `AscendAttentionMetadataBuilder.__init__()` 2. **Remove unused field** - Removed `pcp_prefill_mask` field from `AscendPrefillContextParallelMetadata` (never used in codebase) - Updated related test assertions ### Related - Issue #5463 - PR #4779 (Unify all mask generation methods) - PR #5389 (Make AttentionMaskBuilder singleton) ## Does this PR introduce _any_ user-facing change? No. This is an internal refactoring. ## How was this patch tested? - ✅ Local testing: No linter errors - ✅ Unit tests for attention modules verified - ⏳ CI pipeline Signed-off-by: lico67373 <918688502@qq.com> Co-authored-by: weijinqian0 <1184188277@qq.com>
This commit is contained in:
@@ -13,6 +13,10 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
import torch
|
||||
from vllm.distributed import get_pcp_group
|
||||
|
||||
from vllm_ascend.platform import ModelConfig
|
||||
from vllm_ascend.utils import singleton
|
||||
|
||||
|
||||
def _generate_attn_mask(max_seq_len, dtype):
|
||||
@@ -29,6 +33,7 @@ def _generate_attn_mask(max_seq_len, dtype):
|
||||
return attn_mask
|
||||
|
||||
|
||||
@singleton
|
||||
class AttentionMaskBuilder:
|
||||
|
||||
def __init__(self, device: torch.device):
|
||||
@@ -82,4 +87,16 @@ class AttentionMaskBuilder:
|
||||
triu_mask = torch.triu(mask, diagonal=1).to(self.device)
|
||||
tril_mask = torch.tril(mask, -sliding_window).to(self.device)
|
||||
self.swa_mask = triu_mask + tril_mask
|
||||
return self.swa_mask
|
||||
return self.swa_mask
|
||||
|
||||
def get_attention_mask(self, model_config: ModelConfig):
|
||||
if model_config.runner_type == "pooling":
|
||||
return self.get_attn_mask(2048, torch.bool)
|
||||
|
||||
return self.get_splitfuse_attn_mask()
|
||||
|
||||
def get_final_mla_mask(self, model_config: ModelConfig):
|
||||
if get_pcp_group().world_size > 1:
|
||||
return self.get_pcp_mla_mask(model_config.dtype)
|
||||
# Prefill stages use 512x512 mask with appropriate dtype
|
||||
return self.get_mla_mask(model_config.dtype)
|
||||
Reference in New Issue
Block a user