diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 543012f3..e49f651c 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -42,6 +42,8 @@ class EngineConfig: # CPU MoE backend (--moe-backend cpu): number of CPU worker threads computing # the decode experts. 0 = auto (physical cores). Ignored by other backends. moe_cpu_threads: int = 0 + # skip the startup prefill warmup forwards; first requests then pay the kernel module loads (and the full JIT compile on a cold start) + prefill_warmup: bool = True # Hybrid CPU/GPU decode (--moe-backend offload only): which MoE layers decode on # the CPU executor instead of the GPU offload/PCIe path. Spec is an explicit id # list ("3,7,11"), a count ("8" -> 8 layers evenly strided across depth), or a diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cd6505d2..ac3d3616 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -10,6 +10,7 @@ from freetoken.attention import AttnType, attention_backend_info, create_attention_backend from freetoken.core import Batch, Context, Req, set_global_ctx from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info +from freetoken.env import ENV from freetoken.gpu_select import gpu_identity from freetoken.layers import set_rope_device from freetoken.models import create_model, load_weight @@ -420,8 +421,7 @@ def __init__(self, config: EngineConfig): dummy_req=self.dummy_req, moe_offload_cache=self.moe_offload_cache, ) - if config.attention_backend.split(",")[0] == "triton": - # Prefill runs on the first comma part; warm its autotune cache. + if config.prefill_warmup: self._warmup_prefill() def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup: @@ -934,6 +934,40 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput: copy_done_event.record(self.stream) return ForwardOutput(next_tokens_gpu, next_tokens_cpu, copy_done_event) + def _warmup_prefill_lens(self) -> list[int]: + """Prefill lengths that cross every size bucket the in-repo Triton prefill kernels specialize on.""" + cap = min(self.max_seq_len, self.config.max_forward_len) + if ENV.WARMUP_MAX_LEN.value > 0: + cap = min(cap, ENV.WARMUP_MAX_LEN.value) + mc = self.config.model_config + # 16/32/64 BLOCK_M ladders (nvfp4_linear, fused MoE) plus an odd twin each: triton re-specializes int args on % 16 == 0 + lens = {9, 16, 17, 32, 33, 48, 65, 80} + if mc.is_moe: + top_k = max(1, mc.num_experts_per_tok) + # moe_align leaves its single-CTA path at numel = tokens * top_k > 1024 + lens.add(1024 // top_k + 1) + import triton + + from freetoken.kernel.backend import is_sgl_kernel_installed + + if not is_sgl_kernel_installed(): + # the single-CTA path keys on (next_pow2(numel), num_warps); walk each cell once + seen = set() + for tokens in range(2, 1024 // top_k + 1): + numel = tokens * top_k + cell = (triton.next_power_of_2(numel), triton.next_power_of_2(min(16, max(2, numel // 32)))) + if cell not in seen: + seen.add(cell) + lens.add(tokens) + if cap >= 4096: + # act_and_mul flips BLOCK_D at M >= 4096; also GDN chunk_delta_h NT bucket 1 (NT > 32) + lens.add(4096) + elif mc.has_linear_attention and cap >= 2049: + lens.add(2049) + if mc.has_linear_attention and cap >= 8193: + lens.add(8193) # chunk_delta_h NT bucket 2 (NT > 128) + return sorted(n for n in lens if 2 <= n <= cap) + @torch.inference_mode() def _warmup_prefill(self) -> None: """Compile the Triton prefill path before the first real request. @@ -943,18 +977,18 @@ def _warmup_prefill(self) -> None: restore it afterwards so padded decode graph replay keeps using the dedicated dummy KV slot. """ - if self.max_seq_len < 2: + if AttnType.DSV4 in _required_attn_types(self.config.model_config): + # DSV4 prefill reads pool.full_loc_map, which only the scheduler's paged allocation fills + logger.info_rank0("Prefill warmup skipped for DSV4.") return - - warmup_lens = [min(80, self.max_seq_len)] - if self.max_seq_len >= 128: - warmup_lens.append(128) - warmup_lens = sorted({length for length in warmup_lens if length >= 2}) + warmup_lens = self._warmup_prefill_lens() if not warmup_lens: return + logger.info_rank0(f"Prefill warmup over {len(warmup_lens)} lengths: {warmup_lens}") dummy_row = self.page_table[self.dummy_req.table_idx] dummy_slot = int(dummy_row[0].item()) + done: list[int] = [] started = torch.cuda.Event(enable_timing=True) ended = torch.cuda.Event(enable_timing=True) started.record(self.stream) @@ -972,6 +1006,9 @@ def _warmup_prefill(self) -> None: sampling_params=None, # type: ignore[arg-type] cache_handle=None, # type: ignore[arg-type] ) + if self.linear_state_pool is not None: + # like dummy_req: GDN state writes land in the padding sink, not a real slot + warm_req.linear_slot_idx = self.linear_state_pool.padding_slot batch = Batch(reqs=[warm_req], phase="prefill") batch.padded_reqs = batch.reqs batch.input_ids = torch.zeros(length, dtype=torch.int32, device=self.device) @@ -980,14 +1017,23 @@ def _warmup_prefill(self) -> None: self.attn_backend.prepare_metadata(batch) with self.ctx.forward_batch(batch): self.model.forward() + # surface async faults here, at the failing length, not past the except + torch.cuda.synchronize(self.device) + done.append(length) + except Exception as exc: + if self.config.tp_info.size > 1: + raise + logger.warning_rank0(f"Prefill warmup stopped after {done}: {type(exc).__name__}: {exc}") finally: dummy_row.fill_(dummy_slot) if self.moe_offload_cache is not None: self.moe_offload_cache.reset() + if len(done) != len(warmup_lens): + return ended.record(self.stream) torch.cuda.synchronize(self.device) logger.info_rank0( - f"Prefill warmup complete for lengths {warmup_lens} " + f"Prefill warmup complete for lengths {done} " f"in {started.elapsed_time(ended) / 1000.0:.3f} s" ) diff --git a/python/freetoken/env.py b/python/freetoken/env.py index 6878dc4a..c8eb1e75 100644 --- a/python/freetoken/env.py +++ b/python/freetoken/env.py @@ -75,6 +75,8 @@ class EnvClassSingleton: # fp32 matches the Qwen3.x configs (mamba_ssm_dtype); fp16/bf16 halves the GDN state # pool at some precision cost on the long recurrence (mirrors SGLang's mamba_ssm_dtype). MAMBA_SSM_DTYPE = EnvStr("float32") + # cap the longest prefill warmup forward; 0 = the engine's own chunk cap + WARMUP_MAX_LEN = EnvInt(0) def __new__(cls): # single instance diff --git a/python/freetoken/kernel/fla/l2norm.py b/python/freetoken/kernel/fla/l2norm.py index b7153120..f6ebd6cd 100644 --- a/python/freetoken/kernel/fla/l2norm.py +++ b/python/freetoken/kernel/fla/l2norm.py @@ -51,13 +51,12 @@ def l2norm_fwd_kernel1( # ], # key=["D", "NB"], # ) -@triton.jit +@triton.jit(do_not_specialize=["T"]) def l2norm_fwd_kernel( x, y, eps, - NB: tl.constexpr, - T: tl.constexpr, + T, D: tl.constexpr, BT: tl.constexpr, BD: tl.constexpr, @@ -91,7 +90,6 @@ def l2norm_fwd( raise RuntimeError("This layer doesn't support feature dim >= 64KB.") if D <= 512: - NB = triton.cdiv(T, 2048) def grid(meta): return (triton.cdiv(T, meta["BT"]),) @@ -100,7 +98,6 @@ def grid(meta): x, y, eps, - NB=NB, T=T, D=D, BD=BD, diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index a71b6819..899ac2ac 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -571,6 +571,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--skip-prefill-warmup", + action="store_false", + dest="prefill_warmup", + default=ServerArgs.prefill_warmup, + help=( + "Skip the startup prefill warmup forwards. The first request at each " + "new size then pays the Triton compile/load cost mid-request." + ), + ) + parser.add_argument( "--disable-moe-prefill-overlap", action="store_false",