| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147 |
- --- a/vllm/model_executor/layers/sparse_attn_indexer.py
- +++ b/vllm/model_executor/layers/sparse_attn_indexer.py
- @@ -12,6 +12,9 @@
- from vllm.config import CUDAGraphMode
- from vllm.forward_context import get_forward_context
- from vllm.model_executor.custom_op import CustomOp
- +from vllm.logger import init_logger
- +
- +logger = init_logger(__name__)
- from vllm.platforms import current_platform
- from vllm.triton_utils import HAS_TRITON, tl, triton
- from vllm.utils.import_utils import has_deep_gemm
- @@ -1238,49 +1241,73 @@
- dcp_rank,
- cp_kv_cache_interleave_size,
- ) = _get_dcp_warmup_params()
- - _prewarm_b12x_paged_indexer_prefill(
- - q_quant=q_quant,
- - weights=weights,
- - kv_cache=kv_cache,
- - topk_tokens=int(topk_tokens),
- - profile_q_rows=profile_q_rows,
- - profile_k_rows=warmup_k_rows,
- - page_table_k_rows=profile_k_rows,
- - )
- - if dcp_world_size > 1:
- - dcp_profile_k_rows = _get_dcp_local_k_rows(
- - profile_k_rows,
- - dcp_world_size,
- - dcp_rank,
- - cp_kv_cache_interleave_size,
- - )
- - dcp_warmup_k_rows = _get_b12x_paged_indexer_profile_warmup_k_rows(
- - dcp_profile_k_rows
- + # glm52-gb10 patch (tony's NF3 fix pattern): on sm_121a some
- + # CuteDSL prewarm shapes fail atom_tma_partition ("Operation
- + # creation failed"). Prewarm is a precompile optimization only —
- + # tolerate per-variant failures and let runtime shapes compile
- + # on demand.
- + try:
- + _prewarm_b12x_paged_indexer_prefill(
- + q_quant=q_quant,
- + weights=weights,
- + kv_cache=kv_cache,
- + topk_tokens=int(topk_tokens),
- + profile_q_rows=profile_q_rows,
- + profile_k_rows=warmup_k_rows,
- + page_table_k_rows=profile_k_rows,
- )
- - if dcp_warmup_k_rows != warmup_k_rows:
- - _prewarm_b12x_paged_indexer_prefill(
- - q_quant=q_quant,
- - weights=weights,
- - kv_cache=kv_cache,
- - topk_tokens=int(topk_tokens),
- - profile_q_rows=profile_q_rows,
- - profile_k_rows=dcp_warmup_k_rows,
- - page_table_k_rows=dcp_profile_k_rows,
- + if dcp_world_size > 1:
- + dcp_profile_k_rows = _get_dcp_local_k_rows(
- + profile_k_rows,
- + dcp_world_size,
- + dcp_rank,
- + cp_kv_cache_interleave_size,
- )
- - _prewarm_b12x_dcp_topk_merge(
- - q_rows=profile_q_rows,
- - topk_tokens=int(topk_tokens),
- - dcp_world_size=dcp_world_size,
- - dcp_rank=dcp_rank,
- - cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
- - device=q_quant.device,
- - )
- - _prewarm_b12x_contiguous_prefill_variants(
- - q_quant=q_quant,
- - weights=weights,
- - topk_tokens=int(topk_tokens),
- - profile_q_rows=profile_q_rows,
- - )
- + dcp_warmup_k_rows = _get_b12x_paged_indexer_profile_warmup_k_rows(
- + dcp_profile_k_rows
- + )
- + if dcp_warmup_k_rows != warmup_k_rows:
- + _prewarm_b12x_paged_indexer_prefill(
- + q_quant=q_quant,
- + weights=weights,
- + kv_cache=kv_cache,
- + topk_tokens=int(topk_tokens),
- + profile_q_rows=profile_q_rows,
- + profile_k_rows=dcp_warmup_k_rows,
- + page_table_k_rows=dcp_profile_k_rows,
- + )
- + except Exception as _e:
- + logger.warning(
- + "glm52-gb10: b12x paged-indexer prewarm failed on this "
- + "platform (%s); continuing without precompiled variants.",
- + _e,
- + )
- + try:
- + _prewarm_b12x_dcp_topk_merge(
- + q_rows=profile_q_rows,
- + topk_tokens=int(topk_tokens),
- + dcp_world_size=dcp_world_size,
- + dcp_rank=dcp_rank,
- + cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
- + device=q_quant.device,
- + )
- + except Exception as _e:
- + logger.warning(
- + "glm52-gb10: b12x dcp topk-merge prewarm failed (%s); "
- + "continuing.", _e,
- + )
- + try:
- + _prewarm_b12x_contiguous_prefill_variants(
- + q_quant=q_quant,
- + weights=weights,
- + topk_tokens=int(topk_tokens),
- + profile_q_rows=profile_q_rows,
- + )
- + except Exception as _e:
- + logger.warning(
- + "glm52-gb10: b12x contiguous-prefill prewarm failed (%s); "
- + "continuing.", _e,
- + )
- else:
- # Reserve workspace for indexer during profiling run.
- current_workspace_manager().get_simultaneous(
- @@ -1830,6 +1857,19 @@
- self.max_model_len = max_model_len
- self.max_total_seq_len = max_total_seq_len
- self.topk_indices_buffer = topk_indices_buffer
- + if topk_scores_buffer is None and topk_indices_buffer is not None:
- + try:
- + from vllm.distributed.parallel_state import get_dcp_group
- +
- + dcp_world_size = get_dcp_group().world_size
- + except Exception:
- + dcp_world_size = 1
- + if dcp_world_size > 1 and use_b12x_sparse_indexer():
- + topk_scores_buffer = torch.empty(
- + topk_indices_buffer.shape,
- + dtype=torch.float32,
- + device=topk_indices_buffer.device,
- + )
- self.topk_scores_buffer = topk_scores_buffer
- self.skip_k_cache_insert = skip_k_cache_insert
- self.use_fp4_cache = use_fp4_cache
|
