diff --git a/deepinfer/environ.py b/deepinfer/environ.py index e3a55bb..3fdec73 100644 --- a/deepinfer/environ.py +++ b/deepinfer/environ.py @@ -153,6 +153,19 @@ class envs: DEEPINFER_ENGRAM_PAD_TOKEN = EnvVar.int(default=2) DEEPINFER_ENGRAM_FORCE_GPU_STORAGE = EnvVar(lambda v: _bool(v) or envs.DEEPINFER_ACCELERATOR == "ascend") + @EnvVar + def DEEPINFER_ENGRAM_MIN_PREFILLING_TOKENS(v: str | None) -> int: + """Lower bound of the Engram buffer's prefill max_num_tokens (per CP rank, 0 disables it). + + It only raises the Engram buffer (and therefore the number of RDMA QPs DeepEP allocates for it), + leaving the EP / MoE / CP buffers untouched. Same unit as max_num_prefilling_tokens, i.e. a + per-rank value already divided by the context parallel size. + """ + tokens = _int(v) + if tokens < 0: + raise ValueError(f"expected a non-negative number of tokens, got {tokens}") + return tokens + # -- Sampling -- DEEPINFER_SAMPLER_FORCE_TOPK = EnvVar.int() DEEPINFER_DECODE_TOKEN_BLACKLIST = EnvVar.str() diff --git a/deepinfer/model_executor/models/v4_modules/engram.py b/deepinfer/model_executor/models/v4_modules/engram.py index 123f2ac..c08d1e4 100644 --- a/deepinfer/model_executor/models/v4_modules/engram.py +++ b/deepinfer/model_executor/models/v4_modules/engram.py @@ -91,7 +91,11 @@ class NgramEmbedding(nn.Module): assert self.weight.shape[0] == self.num_local_embeddings, (self.weight.shape, self.num_local_embeddings) self.num_entries_per_token = len(source_embedding.head_vocab_sizes) - self.max_num_tokens = cfg.max_num_decoding_tokens if cfg.engine_type == EngineType.COMPLETION else cfg.max_num_prefilling_tokens + self.max_num_tokens = ( + cfg.max_num_decoding_tokens + if cfg.engine_type == EngineType.COMPLETION + else max(cfg.max_num_prefilling_tokens, envs.DEEPINFER_ENGRAM_MIN_PREFILLING_TOKENS) + ) def gather_sf(self) -> torch.Tensor: assert self.sf is not None