mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +08:00
b2d94e16b2
- fused_sample reject mode for greedy rows (accept iff the draft is the argmax), get_accept_result / finalize_verification (static final rows), store_mtp_info (static pool writes, slot 0 kept), confidence_head_fused, dflash_input_store, get_fused_dflash_get_input: the frontend refs with the kernels' argument lists. - CP followers run the draft forward (its MoE has collectives); the driver's accepted offsets reach them through one all-reduce on Neuron (no eager broadcast), the request first rows through the graph metadata. - run_deepinfer_tiny.py --nextn: free-running greedy with drafts. Host tiny, nextn=5, eager: short 9/9 golden tokens; long matches up to its known near tie at step 5. - bench17: the two-level MoE all-reduce does not move TPOT (433 ms). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
24 lines
880 B
Python
24 lines
880 B
Python
"""Mirrors ``tilelang_ops/asc/dflash.py`` (confidence_head_fused_tl, dflash_input_store_tl) for DSpark drafting:
|
|
the frontend pure-torch references with the kernels' argument lists."""
|
|
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
|
|
def confidence_head_fused_tl(n_embed: int, markov_rank: int, has_markov_emb: bool) -> Any:
|
|
from ..dflash import confidence_head_fused_ref
|
|
|
|
def _confidence(hidden, markov_emb, w, out, temperature): # noqa: ANN001
|
|
confidence_head_fused_ref(hidden, w.view(-1, 1), temperature, markov_emb if has_markov_emb else None, out=out)
|
|
|
|
return _confidence
|
|
|
|
|
|
def dflash_input_store_tl(n_embed: int, hc_mult: int) -> Any:
|
|
def _store(hidden_states, out): # noqa: ANN001
|
|
h = hidden_states.float().mean(dim=-2) if hc_mult > 1 else hidden_states.reshape(-1, n_embed)
|
|
out.copy_(h.to(out.dtype).reshape(out.shape))
|
|
|
|
return _store
|