From 0e36d808b5996856c2a8e9801c19e68762e3bc1f Mon Sep 17 00:00:00 2001 From: zhaozhuoliang Date: Tue, 27 Oct 2026 20:37:03 +0800 Subject: [PATCH] perf: decouple the Engram buffer size from max_num_batched_tokens The number of RDMA QPs DeepEP allocates for the Engram buffer is derived from its max_num_tokens, which in turn comes from max_num_batched_tokens (model_loader.py). Raising the QP count therefore required raising the global batched-token budget, which also inflates the EP/MoE/CP buffers. Add DEEPINFER_ENGRAM_MIN_PREFILLING_TOKENS as a lower bound applied only to the Engram buffer, so a deployment can give Engram more QPs without resizing the other buffers. Default 0 keeps the current behavior. --- deepinfer/environ.py | 13 +++++++++++++ .../model_executor/models/v4_modules/engram.py | 6 +++++- 2 files changed, 18 insertions(+), 1 deletion(-) 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