diff --git a/deepinfer/engine/llm_engine.py b/deepinfer/engine/llm_engine.py index 271b35c..acc295c 100644 --- a/deepinfer/engine/llm_engine.py +++ b/deepinfer/engine/llm_engine.py @@ -394,15 +394,18 @@ class LLMEngine: ) # Neuron: the follower claims core cp_rank of CUDA_VISIBLE_DEVICES (deepinfer/platform.py); the spawned # process copies this environment at start. - saved = os.environ.get("DEEPINFER_NEURON_LOCAL_RANK") - os.environ["DEEPINFER_NEURON_LOCAL_RANK"] = str(cp_rank) + # Runtime profiling (NEURON_RT_INSPECT_*) stays on the driver: 64 ranks would overwrite one output dir. + overrides = {"DEEPINFER_NEURON_LOCAL_RANK": str(cp_rank), "NEURON_RT_INSPECT_ENABLE": "0"} + saved = {k: os.environ.get(k) for k in overrides} + os.environ.update(overrides) try: worker.start() finally: - if saved is None: - os.environ.pop("DEEPINFER_NEURON_LOCAL_RANK", None) - else: - os.environ["DEEPINFER_NEURON_LOCAL_RANK"] = saved + for k, v in saved.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v logger.info("_init_workers: non-driver workers started") driver_kwargs = { diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py index 8045ad0..93ba4c3 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py @@ -126,6 +126,27 @@ def _plain_logprobs_and_greedy(logits, temperatures): # noqa: ANN001, ANN202 _COMPILED: dict = {} +_FLAGS: dict = {} + + +def reset_host_flags() -> None: + _FLAGS.clear() + + +def _host_flags(temperatures, top_p, top_k, vocab: int) -> tuple[bool, bool]: # noqa: ANN001 + """(any row filters by top-p / top-k, every row is greedy) for the step's sampling parameters. Reading them is a + device sync, and a DSpark step samples the same parameter tensors six times (verification, then each draft), + so the flags are kept per tensor until the step ends (reset_host_flags: a later step's tensors may reuse the + address).""" + from deepinfer.sampling_params import SAMPLING_EPS + + key = tuple((t.data_ptr(), t.shape[0]) for t in (temperatures, top_p, top_k)) + flags = _FLAGS.get(key) + if flags is None: + t_h, p_h, k_h = temperatures.cpu(), top_p.cpu(), top_k.cpu() + flags = (bool(((p_h < 1.0) | ((k_h >= 0) & (k_h < vocab))).any()), bool((t_h <= SAMPLING_EPS).all())) + _FLAGS[key] = flags + return flags def _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_size: int): # noqa: ANN001, ANN202 @@ -151,17 +172,15 @@ def _sample(logits, temperatures, top_p, top_k, seed, sample_positions, vocab_si _COMPILED[fn] = torch.compile(fn, backend=envs.DEEPINFER_NEURON_COMPILE_BACKEND, dynamic=False) return _COMPILED[fn] - # Host copies of the [rows] parameters (they come from the host; the sync waits for the logits anyway). - t_h, p_h, k_h = temperatures.cpu(), top_p.cpu(), top_k.cpu() # The top-p / top-k bisection is 32 rounds over the vocabulary (~0.5 s at 129280 even compiled): only rows that # filter need it, and the default parameters filter nothing. - filtering = bool(((p_h < 1.0) | ((k_h >= 0) & (k_h < logits.shape[1]))).any()) + filtering, all_greedy = _host_flags(temperatures, top_p, top_k, logits.shape[1]) with hide_cuda(): if filtering: logprobs, greedy = compiled(_logprobs_and_greedy)(logits, temperatures, top_p, top_k) else: logprobs, greedy = compiled(_plain_logprobs_and_greedy)(logits, temperatures) - if bool((t_h <= SAMPLING_EPS).all()): + if all_greedy: return logprobs, greedy return logprobs, two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size) @@ -174,7 +193,10 @@ def _greedy_verify( Sampled (non-greedy) verification is not ported (its int64 RNG reject draw).""" from deepinfer.sampling_params import SAMPLING_EPS - if not bool((temperatures.cpu() <= SAMPLING_EPS).all()): + key = ("greedy", temperatures.data_ptr(), temperatures.shape[0]) + if key not in _FLAGS: + _FLAGS[key] = bool((temperatures.cpu() <= SAMPLING_EPS).all()) + if not _FLAGS[key]: unsupported("fused_sample_tl(do_reject_sampling=True) for non-greedy rows") rows = read_rows.long() draft = draft_ids.index_select(0, rows) diff --git a/deepinfer/model_executor/layers/vocab_parallel_embedding.py b/deepinfer/model_executor/layers/vocab_parallel_embedding.py index f2a57a3..61ca149 100644 --- a/deepinfer/model_executor/layers/vocab_parallel_embedding.py +++ b/deepinfer/model_executor/layers/vocab_parallel_embedding.py @@ -3,6 +3,7 @@ import torch +from deepinfer import PLATFORM from deepinfer.model_executor.layers.tilelang_ops.gather import gather from deepinfer.model_executor.utils import reconstruct_param @@ -23,6 +24,10 @@ class VocabParallelEmbedding(torch.nn.Module): if self.weight.device != input_.device: # Host-resident table (Neuron, see DeepseekForCausalLM.load_model): move the looked-up rows only. return gather(self.weight, input_.cpu()).to(input_.device) + if PLATFORM.is_neuron and getattr(self, "ids_in_range", False) and not torch.compiler.is_compiling(): + # Eagerly the zero-filling gather is ~5 NEFF launches; callers whose ids are always in range (DSpark's + # Markov embedding of sampled draft tokens) take one embedding lookup. + return torch.nn.functional.embedding(input_.long().clamp(0, self.weight.shape[0] - 1), self.weight) output = gather(self.weight, input_) return output diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index ced2df3..50a9fc3 100644 --- a/deepinfer/worker/model_runner.py +++ b/deepinfer/worker/model_runner.py @@ -444,6 +444,10 @@ class ModelRunner: if self.nextn: kwargs["kv_caches"] = kv_caches + if PLATFORM.is_neuron: + from deepinfer.model_executor.layers.tilelang_ops.neuron.fused_sample import reset_host_flags + + reset_host_flags() # sampling-parameter flags are read once a step (see _host_flags) profiler = _maybe_profile_sample(self) from deepinfer.neuron_ops.op_census import census