diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/dflash.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/dflash.py new file mode 100644 index 0000000..b413094 --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/dflash.py @@ -0,0 +1,18 @@ +"""Mirrors ``tilelang_ops/asc/dflash.py`` (confidence_head_fused_tl, dflash_input_store_tl). + +DFlash (MTP drafting) only: the draft module and the decoder-layer DFlash input store exist only with +DEEPINFER_USE_NEXTN_DECODE>0. Out of scope on Neuron basic mode (NEURON_PORT.md "Scope"): the factories +keep the asc signatures and raise NotImplementedError. +""" + +from typing import Any + +from ._unsupported import unsupported + + +def confidence_head_fused_tl(n_embed: int, markov_rank: int, has_markov_emb: bool) -> Any: + unsupported("confidence_head_fused_tl") + + +def dflash_input_store_tl(n_embed: int, hc_mult: int) -> Any: + unsupported("dflash_input_store_tl") diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/dflash_copy.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/dflash_copy.py new file mode 100644 index 0000000..0984e30 --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/dflash_copy.py @@ -0,0 +1,13 @@ +"""Mirrors ``tilelang_ops/asc/dflash_copy.py`` (copy_dflash_decode_input_tl). + +Only used by worker/dflash_cuda_graph_runner.py (DFlash MTP drafting under CUDA graphs). Out of scope on +Neuron basic mode (NEURON_PORT.md "Scope"): the factory keeps the asc signature and raises. +""" + +from typing import Any + +from ._unsupported import unsupported + + +def copy_dflash_decode_input_tl(has_driver_inputs: bool) -> Any: + unsupported("copy_dflash_decode_input_tl") diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py new file mode 100644 index 0000000..ea4e01a --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_sample.py @@ -0,0 +1,145 @@ +"""Mirrors ``tilelang_ops/asc/fused_sample.py`` (fused_sample_tl) in pure torch. + +Implemented: temperature log-softmax, top-p/top-k filtering with the Ascend boundary semantics, and fresh +sampling (``do_sample``), which reuses the random_sampling two-level sampler on ``exp(logprobs)`` with +seed ``seq_seeds + seed_offset`` and RNG case 0 -- so fused samples bit-match the standalone sampler. + +Top-p/top-k follows the Ascend kernel (not the sort-based ``apply_top_p_top_k_ref``): + * probs ``p = exp((x - max) / T - log(sum))``; probs below 2**-32 get weight 0 and are never kept; + * targets: mass ``P = max(top_p, 2**-32)`` when ``top_p < 1``; count ``K = min(max(top_k, 1), #eligible)`` + when ``0 <= top_k < vocab``; + * the boundary ``lo`` is the largest fp32 value with ``M(lo) >= P`` or ``C(lo) >= K`` (M/C = mass/count of + weights ``>= lo``), found by bisection on the fp32 bit pattern (the kernel uses 16 quartile rounds); + * keep weights ``> lo`` plus the first ``n`` ties ``== lo`` in token order, with + ``n = min(ceil((P - M(>lo)) / lo), K - C(>lo))`` (each term only for an active target); + * kept logprobs are ``(x - max) / T - log(sum) - log(kept_mass)``; others are ``-inf``. +Unfiltered rows use ``(x - max) / T - log(sum)``. The LSE is computed over the whole row instead of the +kernel's chunked online merge, so logprobs may differ from Ascend in the last ulp. + +Reject sampling (``do_reject_sampling``) is MTP/DSpark verification only and is out of scope on Neuron +basic mode (NEURON_PORT.md "Scope"): the factory raises for it. +""" + +import struct +from typing import Any + +import torch + +from ._unsupported import unsupported +from .random_sampling import two_level_sample + +MIN_PROB = 2.0**-32 +_SEARCH_HI = 1.0 + 1e-3 +_SEARCH_ROUNDS = 30 # bit-pattern distance between MIN_PROB and 1 + 1e-3 is < 2**29 + + +def _f32_bits(x: float) -> int: + return struct.unpack(" torch.Tensor: + """Temperature log-softmax with Ascend top-p/top-k filtering. Returns fp32 [Q, vocab]; logits untouched.""" + num_tokens, vocab_size = logits.shape + x = logits.float() + inv_t = (1.0 / temperatures.float()).unsqueeze(1) + row_max = x.amax(dim=1, keepdim=True) + scaled = (x - row_max) * inv_t + log_sum = torch.log(torch.exp(scaled).sum(dim=1, keepdim=True)) + plain = scaled - log_sum # unfiltered rows (asc: -(log_sum - scaled)) + + top_p = top_p.float() + top_k = top_k.to(torch.int32) + p_active = top_p < 1.0 + k_active = (top_k >= 0) & (top_k < vocab_size) + filter_active = p_active | k_active + + normalized = plain # asc: (x - max) / T + lse0 with lse0 = -log(sum) + prob = torch.exp(normalized) + eligible = prob >= MIN_PROB + w = torch.where(eligible, prob, torch.zeros_like(prob)) + target_mass = top_p.clamp(min=MIN_PROB) + n_eligible = eligible.sum(dim=1, dtype=torch.int32) + target_count = torch.minimum( + torch.where(k_active, top_k.clamp(min=1), torch.full_like(top_k, vocab_size)), n_eligible + ) + + # Bisection on fp32 bit patterns (positive floats are ordered like their bits). + # Invariant: lo satisfies an active target, hi does not; m_hi/c_hi are mass/count of w >= hi. + lo = torch.full((num_tokens,), _LO_BITS, dtype=torch.int32, device=x.device) + hi = torch.full((num_tokens,), _HI_BITS, dtype=torch.int32, device=x.device) + m_hi = torch.zeros((num_tokens,), dtype=torch.float32, device=x.device) + c_hi = torch.zeros((num_tokens,), dtype=torch.int32, device=x.device) + zero = torch.zeros((), dtype=torch.float32, device=x.device) + for _ in range(_SEARCH_ROUNDS): + mid = lo + (hi - lo) // 2 + ge = w >= mid.view(torch.float32).unsqueeze(1) + mass = torch.where(ge, w, zero).sum(dim=1) + count = ge.sum(dim=1, dtype=torch.int32) + hit = (p_active & (mass >= target_mass)) | (k_active & (count >= target_count)) + lo = torch.where(hit, mid, lo) + hi = torch.where(hit, hi, mid) + m_hi = torch.where(hit, m_hi, mass) + c_hi = torch.where(hit, c_hi, count) + + boundary = lo.view(torch.float32) + n_keep = torch.full((num_tokens,), vocab_size, dtype=torch.int32, device=x.device) + p_quota = torch.ceil(((target_mass - m_hi) / boundary).clamp(min=0.0, max=float(vocab_size))).to(torch.int32) + n_keep = torch.where(p_active, p_quota, n_keep) + n_keep = torch.where(k_active, torch.minimum(n_keep, target_count - c_hi), n_keep) + + b = boundary.unsqueeze(1) + tie = w == b + tie_rank = tie.to(torch.int32).cumsum(dim=1, dtype=torch.int32) + keep = (w > b) | (tie & (tie_rank <= n_keep.unsqueeze(1))) + kept_mass = torch.where(keep, w, zero).sum(dim=1, keepdim=True) + filtered = torch.where(keep, normalized - torch.log(kept_mass), torch.full_like(normalized, float("-inf"))) + return torch.where(filter_active.unsqueeze(1), filtered, plain) + + +def fused_sample_tl( + vocab_size: int, + do_sample: bool, + do_reject_sampling: bool, +) -> Any: + """Factory mirroring asc ``fused_sample_tl``; the kernel takes the asc argument list.""" + assert not do_reject_sampling or do_sample, "reject sampling requires sampling" + if do_reject_sampling: + unsupported("fused_sample_tl(do_reject_sampling=True)") + + def _fused_sample( + logits: torch.Tensor, + temperatures: torch.Tensor, + top_p: torch.Tensor, + top_k: torch.Tensor, + sample_positions: torch.Tensor | None, + seq_seeds: torch.Tensor | None, + is_verify: torch.Tensor | None, + logprobs: torch.Tensor, + candidate_token_ids: torch.Tensor | None, + last_nextn_token_logprobs: torch.Tensor | None, + mtp_info_read_rows: torch.Tensor | None, + last_next_n_token_ids: torch.Tensor | None, + accept_result_verify: torch.Tensor | None, + draft_selected_logprobs: torch.Tensor | None, + accept_probs: torch.Tensor | None, + seed_offset: int, + lse_workspace: torch.Tensor | None, + ) -> None: + assert logits.shape[1] == vocab_size and logits.stride(1) == 1 + logprobs.copy_(compute_logprobs(logits, temperatures, top_p, top_k)) + if do_sample: + seed = seq_seeds.view(torch.int64) if seq_seeds.dtype == torch.uint64 else seq_seeds.to(torch.int64) + seed = seed + int(seed_offset) # uint64 wraparound == int64 wraparound on the same bits + # Exponentiate the written logprobs so the weights equal exp(logprobs_out) bit for bit. + candidate_token_ids.copy_(two_level_sample(torch.exp(logprobs), seed, sample_positions, 0, vocab_size)) + + return _fused_sample diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/random_sampling.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/random_sampling.py new file mode 100644 index 0000000..4a54f49 --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/random_sampling.py @@ -0,0 +1,161 @@ +"""Mirrors ``tilelang_ops/asc/random_sampling.py`` (random_sampling_tl) in pure torch. + +Same algorithm as the Ascend two-level fixed-point sampler: + * weights are quantized to fixed point ``trunc(p * 2**31)`` (saturated to int32), and folded into + ``ACC_SLOTS`` sums by ``col % ACC_SLOTS``; + * one 64-bit random number ``(rand_high << 32) | rand_low`` is drawn from the counter-based RNG + initialised with ``rng_init(seed, 3 * position + case)`` and reduced as ``r = rand % total``; + * level 1 picks the slot whose inclusive CDF first exceeds ``r``; level 2 picks the column + ``slot + m * ACC_SLOTS`` whose inclusive CDF first exceeds the remainder. + +RNG: ``T.rng_init`` / ``T.rng_rand`` are modelled as cuRAND ``curandStatePhilox4_32_10_t`` +(``curand_init(seed, subsequence=seq, offset=0)``; the first two ``curand()`` calls return output words +x and y of Philox4x32-10 at counter ``{0, 0, lo(seq), hi(seq)}``, key ``{lo(seed), hi(seed)}``). This is +what tilelang's CUDA backend emits; the Ascend intrinsic source is not available here, so bit-equality +with Ascend hardware assumes it implements the same generator. Unsigned 64/32-bit arithmetic is +emulated with int64 tensors (no uint64 tensor ops needed). +""" + +from typing import Any + +import torch + +_MASK32 = 0xFFFFFFFF +_PHILOX_M0 = 0xD2511F53 +_PHILOX_M1 = 0xCD9E8D57 +_PHILOX_W0 = 0x9E3779B9 +_PHILOX_W1 = 0xBB67AE85 +_INT32_MAX = (1 << 31) - 1 +SCALE_BITS = 31 + + +def _mulhilo32(a: torch.Tensor, m: int) -> tuple[torch.Tensor, torch.Tensor]: + """(hi, lo) 32-bit halves of a * m for int64 ``a`` in [0, 2**32) and a 32-bit constant ``m``. + + Splits ``m`` into 16-bit halves so every intermediate stays below 2**49 (no int64 overflow).""" + m_hi, m_lo = m >> 16, m & 0xFFFF + p_lo = a * m_lo + t = a * m_hi + (p_lo >> 16) + return t >> 16, ((t & 0xFFFF) << 16) | (p_lo & 0xFFFF) + + +def philox4x32_10( + c0: torch.Tensor, c1: torch.Tensor, c2: torch.Tensor, c3: torch.Tensor, k0: torch.Tensor, k1: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Philox4x32-10 block function (cuRAND / Random123 round and key schedule) on int64 lanes.""" + for i in range(10): + if i: + k0 = (k0 + _PHILOX_W0) & _MASK32 + k1 = (k1 + _PHILOX_W1) & _MASK32 + hi0, lo0 = _mulhilo32(c0, _PHILOX_M0) + hi1, lo1 = _mulhilo32(c2, _PHILOX_M1) + c0, c1, c2, c3 = hi1 ^ c1 ^ k0, lo1, hi0 ^ c3 ^ k1, lo0 + return c0, c1, c2, c3 + + +def rng_first_two(seed: torch.Tensor, seq: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """First two 32-bit outputs of ``rng_init(seed, seq)``; both returned as int64 in [0, 2**32). + + seed: int64 tensor holding uint64 bits. seq: int64 tensor holding the (sign-extended) subsequence.""" + k0 = seed & _MASK32 + k1 = (seed >> 32) & _MASK32 + c2 = seq & _MASK32 + c3 = (seq >> 32) & _MASK32 + zero = torch.zeros_like(c2) + x, y, _, _ = philox4x32_10(zero, zero, c2, c3, k0, k1) + return x, y + + +def _wrap_int32(x: torch.Tensor) -> torch.Tensor: + """int64 tensor -> value of the same bits truncated to int32, sign-extended back to int64.""" + return ((x + (1 << 31)) & _MASK32) - (1 << 31) + + +def rng_sequence(positions: torch.Tensor, sample_case: Any) -> torch.Tensor: + """``position * 3 + case`` computed in int32 (as in the kernel), sign-extended to int64.""" + return _wrap_int32(positions.to(torch.int64) * 3 + sample_case) + + +def acc_slots_for(vocab_size: int) -> int: + """ACC_SLOTS used by asc random_sampling_tl / fused_sample_tl for this vocab size.""" + if vocab_size >= 8192: + return 4096 + bd = max(64, 1 << max(0, (vocab_size - 1).bit_length())) + return min(bd, 2048) + + +def quantize_weights(weights: torch.Tensor) -> torch.Tensor: + """trunc(w * 2**31) as int32, saturated at INT32_MAX (vcvt ROUND_Z, sat=True).""" + w = weights.float() + sat = torch.full_like(w, _INT32_MAX, dtype=torch.int32) + # Below 1.0 the scaled value is exact in fp32 and < 2**31, so the int32 cast is a plain trunc. + below_one = w.clamp(min=0.0, max=1.0 - 2.0**-24) + return torch.where(w >= 1.0, sat, (below_one * float(1 << SCALE_BITS)).to(torch.int32)) + + +def two_level_sample( + weights: torch.Tensor, + seed: torch.Tensor, + positions: torch.Tensor, + sample_case: Any, + vocab_size: int, +) -> torch.Tensor: + """Ascend two-level fixed-point sampling for every row. Returns int32 [num_tokens]. + + weights: [num_tokens, vocab_size] non-negative fp32, row sums <= 1 (may be strided rows). + seed: [num_tokens] int64 holding uint64 bits. positions: [num_tokens] int32. + sample_case: int or int64 tensor (0 fresh, 1 nextn/residual).""" + num_tokens = weights.shape[0] + device = weights.device + acc_slots = acc_slots_for(vocab_size) + n_slots = (vocab_size + acc_slots - 1) // acc_slots + fx = quantize_weights(weights[:, :vocab_size]) + pad = n_slots * acc_slots - vocab_size + if pad: + fx = torch.nn.functional.pad(fx, (0, pad)) + fx = fx.view(num_tokens, n_slots, acc_slots) + + # Level 1: per-slot fixed-point sums and their inclusive CDF. + acc_cdf = fx.sum(dim=1, dtype=torch.int64).cumsum(dim=1) # [T, ACC_SLOTS] + total = acc_cdf[:, -1] + # NEURON: the kernel device_asserts total > 0; here an all-zero row samples column 0. + safe_total = total.clamp(min=1) + + lo, hi = rng_first_two(seed, rng_sequence(positions, sample_case)) + # (hi * 2**32 + lo) % total without uint64: all partial products stay below 2**63. + r = ((hi % safe_total) * ((1 << 32) % safe_total) + lo) % safe_total + + offset = (acc_cdf <= r.unsqueeze(1)).sum(dim=1) # first slot with cdf > r + offset = torch.where(offset < acc_slots, offset, torch.zeros_like(offset)) + prev = torch.where( + offset > 0, + acc_cdf.gather(1, (offset - 1).clamp(min=0).unsqueeze(1)).squeeze(1), + torch.zeros_like(r), + ) + r2 = r - prev + + # Level 2: the columns offset + m * ACC_SLOTS of the chosen slot. + idx = offset.view(num_tokens, 1, 1).expand(num_tokens, n_slots, 1) + slot_cdf = fx.gather(2, idx).squeeze(2).to(torch.int64).cumsum(dim=1) # [T, n_slots] + sel = (slot_cdf <= r2.unsqueeze(1)).sum(dim=1) + sel = torch.where(sel < n_slots, sel, torch.zeros_like(sel)) + col = offset + sel * acc_slots + col = torch.where(col < vocab_size, col, torch.zeros_like(col)) + return col.to(torch.int32).to(device) + + +def random_sampling_tl(vocab_size: int, is_nextn_predict: bool) -> Any: + """Factory mirroring asc ``random_sampling_tl``; kernel(probs, seq_seeds, sample_positions, out).""" + sample_case = 1 if is_nextn_predict else 0 + + def _random_sampling( + probs: torch.Tensor, + seq_seeds: torch.Tensor, + sample_positions: torch.Tensor, + out: torch.Tensor, + ) -> None: + assert probs.shape[1] == vocab_size and probs.stride(1) == 1 + seed = seq_seeds.view(torch.int64) if seq_seeds.dtype == torch.uint64 else seq_seeds.to(torch.int64) + out.copy_(two_level_sample(probs, seed, sample_positions, sample_case, vocab_size)) + + return _random_sampling diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/sampler.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/sampler.py new file mode 100644 index 0000000..6738aae --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/sampler.py @@ -0,0 +1,122 @@ +"""Mirrors ``tilelang_ops/asc/sampler.py`` in pure torch. + +Implemented (used by plain sampling, including nextn=0 decode): prepare_verification_metadata_tl, +expand_sampling_params_tl, map_lookup_tl. + +MTP/DSpark verification only, out of scope on Neuron basic mode (NEURON_PORT.md "Scope"; the factories keep +the asc signatures and raise): remap_draft_topk_tl, get_accept_result_tl, finalize_verification_tl. +""" + +from typing import Any, Literal + +import torch + +from ..common import DRAFT_POS_HASH_MUL +from ._unsupported import unsupported + + +def prepare_verification_metadata_tl(num_steps: int) -> Any: + """kernel(ftp, seqlen_q, row_to_req, local_step, is_verify); rows are packed per request from ftp.""" + + def _prepare_verification_metadata( + ftp: torch.Tensor, + seqlen_q: torch.Tensor, + row_to_req: torch.Tensor, + local_step: torch.Tensor, + is_verify: torch.Tensor, + ) -> None: + total_q = row_to_req.shape[0] + req_cnt = ftp.shape[0] + device = ftp.device + # row -> request without searchsorted/host sync: mark each request start, then prefix-count. + starts = torch.zeros(total_q + 1, dtype=torch.int32, device=device) + starts.scatter_add_(0, ftp.to(torch.int64).clamp(0, total_q), torch.ones(req_cnt, dtype=torch.int32, device=device)) + req = starts[:total_q].cumsum(0, dtype=torch.int32) - 1 + req_idx = req.clamp(min=0).to(torch.int64) + step = torch.arange(total_q, dtype=torch.int32, device=device) - ftp.to(torch.int32)[req_idx] + row_to_req.copy_(req) + local_step.copy_(step) + is_verify.copy_(step + 1 < seqlen_q.to(torch.int32)[req_idx]) + + return _prepare_verification_metadata + + +def expand_sampling_params_tl() -> Any: + """kernel(temperatures, top_ps, top_ks, sample_positions, seq_seeds, row_to_req, local_step, + temperatures_q, top_ps_q, top_ks_q, sample_positions_q, seq_seeds_q).""" + + def _expand_sampling_params( + temperatures: torch.Tensor, + top_ps: torch.Tensor, + top_ks: torch.Tensor, + sample_positions: torch.Tensor, + seq_seeds: torch.Tensor, + row_to_req: torch.Tensor, + local_step: torch.Tensor, + temperatures_q: torch.Tensor, + top_ps_q: torch.Tensor, + top_ks_q: torch.Tensor, + sample_positions_q: torch.Tensor, + seq_seeds_q: torch.Tensor, + ) -> None: + req = row_to_req.to(torch.int64) + temperatures_q.copy_(temperatures[req]) + top_ps_q.copy_(top_ps[req]) + top_ks_q.copy_(top_ks[req]) + # int32 arithmetic with wraparound, as in the kernel. + sample_positions_q.copy_(sample_positions.to(torch.int32)[req] + local_step.to(torch.int32) * DRAFT_POS_HASH_MUL) + seq_seeds_q.view(torch.int64).copy_(seq_seeds.view(torch.int64)[req]) + + return _expand_sampling_params + + +def map_lookup_tl( + force_topk: int, + mode: Literal["by_key", "by_index"] = "by_key", + value_dtype: torch.dtype = torch.float32, +) -> Any: + """kernel(key, value, query, out) with out/query [rows, 1]. + + by_index: ``value[row, query]`` or the out-of-bounds value (-inf float / -1 int). + by_key: the value whose key equals the query (keys are unique top-k indices; with duplicate keys this + returns their max), -inf when missing.""" + assert mode in ("by_key", "by_index") + assert mode != "by_key" or value_dtype.is_floating_point, "by_key requires floating-point values" + oob_value = float("-inf") if value_dtype.is_floating_point else -1 + + def _map_lookup( + key: torch.Tensor | None, + value: torch.Tensor, + query: torch.Tensor, + out: torch.Tensor, + ) -> None: + q = query[:, 0] + if mode == "by_index": + q64 = q.to(torch.int64) + in_bounds = (q64 >= 0) & (q64 < force_topk) + got = value[:, :force_topk].gather(1, torch.where(in_bounds, q64, torch.zeros_like(q64)).unsqueeze(1)).squeeze(1) + result = torch.where(in_bounds, got, torch.full_like(got, oob_value)) + else: + match = key[:, :force_topk] == q.unsqueeze(1) + v = value[:, :force_topk] + result = torch.where(match, v, torch.full_like(v, oob_value)).amax(dim=1) + out[:, 0].copy_(result) + + return _map_lookup + + +def remap_draft_topk_tl(force_topk: int, vocab_size: int) -> Any: + unsupported("remap_draft_topk_tl") + + +def get_accept_result_tl(has_golden: bool) -> Any: + unsupported("get_accept_result_tl") + + +def finalize_verification_tl( + has_sum_smaller_probs: bool, + has_confidence: bool, + num_steps: int, + num_stop_tokens: int, +) -> Any: + unsupported("finalize_verification_tl") diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/schedule_mtp_varlen.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/schedule_mtp_varlen.py new file mode 100644 index 0000000..d063d97 --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/schedule_mtp_varlen.py @@ -0,0 +1,14 @@ +"""Mirrors ``tilelang_ops/asc/schedule_mtp_varlen.py`` (compute_seqlens_q_by_confidence_rates_tl). + +Confidence-scheduled MTP varlen (DEEPINFER_ENABLE_MTP_VARLEN, zero-overhead scheduling path in +worker/worker.py). Out of scope on Neuron basic mode (NEURON_PORT.md "Scope"): the factory keeps the asc +signature and raises. +""" + +from typing import Any + +from ._unsupported import unsupported + + +def compute_seqlens_q_by_confidence_rates_tl(nextn: int) -> Any: + unsupported("compute_seqlens_q_by_confidence_rates_tl") diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/store_mtp_info.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/store_mtp_info.py new file mode 100644 index 0000000..142b62f --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/store_mtp_info.py @@ -0,0 +1,13 @@ +"""Mirrors ``tilelang_ops/asc/store_mtp_info.py`` (store_mtp_info_tl). + +MTP info pool writes; MTPInfoPool is only built by the DFlash draft module (MTP drafting). Out of scope on +Neuron basic mode (NEURON_PORT.md "Scope"): the factory keeps the asc signature and raises. +""" + +from typing import Any + +from ._unsupported import unsupported + + +def store_mtp_info_tl(nextn: int, vocab_size: int, has_index: bool, has_confidence: bool) -> Any: + unsupported("store_mtp_info_tl") diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/structure_output.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/structure_output.py new file mode 100644 index 0000000..56687a9 --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/structure_output.py @@ -0,0 +1,54 @@ +"""Mirrors ``tilelang_ops/asc/structure_output.py`` (apply_token_bitmask_by_step_inplace_tl) in pure torch. + +For every selected request, the logits row (``indices[i]``, or ``ftp[indices[i]] + step`` when ``has_ftp``; +skipped when that row lies outside the request) gets ``-inf`` in the first ``vocab_size`` columns whose bit +``bitmask[req, col // 32] >> (col % 32) & 1`` is 0. In place, no host sync: rows are resolved with a +scatter into a per-row request map, so selected rows must be distinct (as in the kernel). +""" + +from typing import Any + +import torch + +_BITS_PER_BLOCK = 32 + + +def apply_token_bitmask_by_step_inplace_tl(vocab_size: int, has_ftp: bool) -> Any: + """Factory mirroring asc ``apply_token_bitmask_by_step_inplace_tl``; kernel(logits, bitmask, indices, ftp, step).""" + num_words = (vocab_size + _BITS_PER_BLOCK - 1) // _BITS_PER_BLOCK + + def _apply( + logits: torch.Tensor, + bitmask: torch.Tensor, + indices: torch.Tensor, + ftp: torch.Tensor | None, + step: int, + ) -> None: + num_rows = logits.shape[0] + device = logits.device + req = indices.to(torch.int64) + if has_ftp: + first = ftp.to(torch.int64) + num_requests = first.shape[0] + rows = first[req] + int(step) + nxt = req + 1 + end = torch.where( + nxt < num_requests, + first[nxt.clamp(max=num_requests - 1)], + torch.full_like(nxt, num_rows), + ) + valid = rows < end + else: + rows = req + valid = torch.ones_like(req, dtype=torch.bool) + # row -> request map; invalid entries go to a dummy slot at num_rows. + row_req = torch.full((num_rows + 1,), -1, dtype=torch.int64, device=device) + row_req.scatter_(0, torch.where(valid, rows, torch.full_like(rows, num_rows)), req) + row_req = row_req[:num_rows] + selected = row_req >= 0 + words = bitmask[row_req.clamp(min=0)][:, :num_words].to(torch.int32) + shifts = torch.arange(_BITS_PER_BLOCK, dtype=torch.int32, device=device) + bits = ((words.unsqueeze(-1) >> shifts) & 1).reshape(num_rows, num_words * _BITS_PER_BLOCK)[:, :vocab_size] + logits[:, :vocab_size].masked_fill_((bits == 0) & selected.unsqueeze(1), float("-inf")) + + return _apply diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/sum_smaller_probs.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/sum_smaller_probs.py new file mode 100644 index 0000000..e4eec81 --- /dev/null +++ b/deepinfer/model_executor/layers/tilelang_ops/neuron/sum_smaller_probs.py @@ -0,0 +1,27 @@ +"""Mirrors ``tilelang_ops/asc/sum_smaller_probs.py`` (sum_smaller_probs_tl) in pure torch. + +Per row: sum of ``exp(logprob)`` over tokens with ``logprob < pivot`` or ``logprob == pivot`` at a column +``<= sample``, where ``pivot = logprobs[row, sample]``; fp32 accumulation, rows with ``sample < 0`` give 0. +""" + +from typing import Any + +import torch + + +def sum_smaller_probs_tl(vocab_size: int) -> Any: + """Factory mirroring asc ``sum_smaller_probs_tl``; kernel(logprobs, sample_results, output).""" + + def _sum_smaller_probs(logprobs: torch.Tensor, sample_results: torch.Tensor, output: torch.Tensor) -> None: + assert logprobs.shape[1] == vocab_size and logprobs.stride(1) == 1 + sample = sample_results.to(torch.int64) + invalid = sample < 0 + col = sample.clamp(min=0).unsqueeze(1) + lp = logprobs.float() + pivot = lp.gather(1, col) + cols = torch.arange(vocab_size, device=lp.device).unsqueeze(0) + pred = (lp < pivot) | ((lp == pivot) & (cols <= col)) + out = torch.where(pred, torch.exp(lp), torch.zeros((), dtype=torch.float32, device=lp.device)).sum(dim=1) + output.copy_(torch.where(invalid, torch.zeros_like(out), out)) + + return _sum_smaller_probs diff --git a/deepinfer/model_executor/models/v4_modules/engram.py b/deepinfer/model_executor/models/v4_modules/engram.py index 123f2ac..1f40794 100644 --- a/deepinfer/model_executor/models/v4_modules/engram.py +++ b/deepinfer/model_executor/models/v4_modules/engram.py @@ -97,7 +97,11 @@ class NgramEmbedding(nn.Module): assert self.sf is not None local_sf = _pack_ue8m0_to_int(self.sf) sf = local_sf.new_empty(self.engram_group.size() * self.num_local_embeddings, local_sf.size(1)) - dist.all_gather_into_tensor(sf, local_sf, group=self.engram_group) + if PLATFORM.is_neuron: + # Gather raw bytes: gloo (host mode) has no int16/int32 collectives for this dtype. + dist.all_gather_into_tensor(sf.view(torch.uint8), local_sf.contiguous().view(torch.uint8), group=self.engram_group) + else: + dist.all_gather_into_tensor(sf, local_sf, group=self.engram_group) return sf def release_storage(self) -> None: diff --git a/deepinfer/neuron_ops/transfer_to_neuron.py b/deepinfer/neuron_ops/transfer_to_neuron.py index 775c16b..ee48d44 100644 --- a/deepinfer/neuron_ops/transfer_to_neuron.py +++ b/deepinfer/neuron_ops/transfer_to_neuron.py @@ -111,6 +111,9 @@ _FACTORIES = ( def _patch_factories() -> None: + # `with torch.device(...)` redirects calls whose func is in this lru_cached set; build it from the + # original C functions before they are wrapped, or the context manager stops applying to them. + torch.utils._device._device_constructors() for name in _FACTORIES: if hasattr(torch, name): setattr(torch, name, _wrap_device_kwarg(getattr(torch, name))) diff --git a/deepinfer/neuron_ops/vendor/deep_ep/__init__.py b/deepinfer/neuron_ops/vendor/deep_ep/__init__.py new file mode 100644 index 0000000..d431fb5 --- /dev/null +++ b/deepinfer/neuron_ops/vendor/deep_ep/__init__.py @@ -0,0 +1,44 @@ +"""Pure-torch drop-in for DeepSeek's `deep_ep` (DeepEP-V2), following the Ascend variant. + +Upstream (`vendor/DeepEP-Ascend`) is a JIT-compiled HCCL/URMA extension over registered +symmetric memory. On Trainium there is no such runtime in this basic port, so every buffer +is implemented with `torch.distributed` collectives on the process group DeepInfer passes in +(torch-neuronx's ``neuron`` backend on device, ``gloo`` for host correctness runs), and the +asynchronous event objects become synchronous no-ops with the same interface. + +Implemented: + EPBuffer / EPHandle expanded-layout MoE dispatch & combine (all-to-all) + EngramBuffer sharded n-gram embedding lookup (all-to-all of indices + rows) + BucketBuffer / registered-storage sessions + batched all-gather (CP only) + BucketSession + EventHandle / EventOverlap, BufferAllocator, alignment and domain-size helpers + +Ignored-for-compatibility knobs: `num_sms`, `num_qps`, `qp_depth`, `sl_idx`, hybrid mode, +`previous_event`, `allocate_on_comm_stream`, huge pages. Raising stubs (basic mode): +`EPBuffer.lb_*`, `BucketBuffer.reduce_scatter` / `all_reduce`. +""" + +from . import comm +from .comm import ( + destroy_all_managed_hccl_comm, + get_logical_domain_size, + get_num_allocation_alignment, + get_num_rdma_alignment, + get_num_tma_alignment, + get_physical_domain_size, +) +from .buffers.allocator import BufferAllocator +from .buffers.base import BufferBase +from .buffers.bucket import BucketBuffer, BucketSession +from .buffers.engram import EngramBuffer +from .buffers.ep import EPBuffer, EPHandle, topk_idx_t +from .utils.event import EventHandle, EventOverlap + +__all__ = [ + 'BucketBuffer', 'BucketSession', 'BufferAllocator', 'BufferBase', 'EPBuffer', 'EPHandle', + 'EngramBuffer', 'EventHandle', 'EventOverlap', 'comm', 'destroy_all_managed_hccl_comm', + 'get_logical_domain_size', 'get_num_allocation_alignment', 'get_num_rdma_alignment', + 'get_num_tma_alignment', 'get_physical_domain_size', 'topk_idx_t', +] + +__version__ = '1.0.0' diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/__init__.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/allocator.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/allocator.py new file mode 100644 index 0000000..4c04e6d --- /dev/null +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/allocator.py @@ -0,0 +1,40 @@ +"""Mirrors `deep_ep.buffers.allocator` of DeepEP-Ascend (device-agnostic).""" + +from __future__ import annotations + +from typing import List, Tuple + +import torch + +from ..comm import get_num_allocation_alignment +from ..utils.math import align + + +class BufferAllocator: + """Plan contiguous, aligned tensors before allocating their backing storage.""" + + def __init__(self): + self.alignment = get_num_allocation_alignment() + self.allocations: List[Tuple[torch.Tensor, int]] = [] + self.num_bytes = 0 + self.materialized = False + + def allocate(self, shape, dtype: torch.dtype) -> torch.Tensor: + """Append an aligned tensor allocation and return its meta placeholder.""" + assert not self.materialized, 'A materialized allocation plan cannot be changed' + tensor = torch.empty(shape, dtype=dtype, device='meta') + offset = align(self.num_bytes, self.alignment) + self.allocations.append((tensor, offset)) + self.num_bytes = align(offset + tensor.nbytes, self.alignment) + return tensor + + def materialize(self, storage: torch.Tensor) -> None: + assert not self.materialized, 'An allocation plan can only be materialized once' + # NEURON: no `npu` device check; the storage follows the caller's device. + assert storage.is_contiguous() and storage.dtype == torch.uint8 + assert storage.numel() == self.num_bytes + + for tensor, offset in self.allocations: + value = storage.narrow(0, offset, tensor.nbytes).view(tensor.dtype).view(tensor.shape) + torch.utils.swap_tensors(tensor, value) + self.materialized = True diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/base.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/base.py new file mode 100644 index 0000000..9e0bad0 --- /dev/null +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/base.py @@ -0,0 +1,28 @@ +"""Mirrors `deep_ep.buffers.base` of DeepEP-Ascend.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class BufferBase(ABC): + """Common lifecycle for communication buffers.""" + + def __init__(self, explicitly_destroy: bool = False): + self.explicitly_destroy = explicitly_destroy + self.destroyed = False + + @property + def main_context(self): + """Upstream returns the native communication context; this port has none.""" + return None + + def print_memory_usage(self) -> None: + print(f'{type(self).__name__}: {getattr(self, "num_bytes", 0)} bytes ' + f'(collective-based port, no registered symmetric memory)') + + @abstractmethod + def destroy(self) -> None: + """Destroy the buffer. Requires ``explicitly_destroy=True`` at construction.""" + assert self.explicitly_destroy + self.destroyed = True diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/bucket.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/bucket.py new file mode 100644 index 0000000..8b6a683 --- /dev/null +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/bucket.py @@ -0,0 +1,280 @@ +"""Mirrors `deep_ep.buffers.bucket` (`BucketBuffer` / `BucketSession`) of DeepEP-Ascend. + +Only used by context parallelism (`model_executor/parallel_utils/cp.py`), which DeepInfer +builds exclusively for `cp_size > 1` (`parallel_utils/userbuffer.py:103-105` returns before +constructing a `CPManager` at CP=1). CP>1 is off in basic mode, so this is the minimal +version: a plain storage tensor plus a session that carves aligned regions out of it, and a +batched `all_gather` implemented with `dist.all_gather_into_tensor`. `reduce_scatter` / +`all_reduce` raise. + +The session keeps upstream's calling convention: inside `with buffer.session() as session`, +`buffer.all_gather(...)` accepts unregistered sources, and `wait()` returns 1-D views of the +gathered data over the registered storage (what `CPManager.all_gather` re-views). +""" + +from __future__ import annotations + +import inspect +import math +from functools import partial +from typing import Any, Optional, Sequence, Union + +import torch +import torch.distributed as dist + +from .. import comm +from .allocator import BufferAllocator +from .base import BufferBase +from ..utils.event import EventHandle, EventOverlap +from ..utils.math import align + + +def _resolve_device() -> torch.device: + """Device for the registered storage. + + DeepInfer constructs the CP buffer without passing any tensor, so follow the accelerator + DeepInfer itself uses: `torch.device("cuda")`, which `transfer_to_neuron` maps onto + `neuron` (or the host with `DEEPINFER_NEURON_DEVICE=cpu`). Fall back to CPU. + """ + try: + if torch.cuda.is_available(): + return torch.device('cuda') + except Exception: + pass + return torch.device('cpu') + + +class BucketSession: + """A session in which collective methods can accept unregistered tensors.""" + + def __init__(self, buffer: 'BucketBuffer'): + self.buffer = buffer + self.alignment = comm.get_num_rdma_alignment() + self.num_allocated_bytes = 0 + self.context_idx: Optional[int] = None + self.collective_methods = ('all_gather', 'reduce_scatter', 'all_reduce') + + def __enter__(self) -> 'BucketSession': + assert not any(name in self.buffer.__dict__ for name in self.collective_methods), \ + 'Nested sessions are not supported' + self.num_allocated_bytes = 0 + self.context_idx = None + for name in self.collective_methods: + setattr(self.buffer, name, partial(self._collective_func, name)) + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + for name in self.collective_methods: + delattr(self.buffer, name) + + def _call_original(self, name: str, *args, **kwargs): + return getattr(type(self.buffer), name)(self.buffer, *args, **kwargs) + + def _parse_args(self, name: str, *args, **kwargs) -> dict: + signature = inspect.signature(getattr(type(self.buffer), name)) + kwargs = signature.bind(self.buffer, *args, **kwargs).arguments + kwargs.pop('self') + return kwargs + + def _allocate(self, num_bytes: int, dtype: torch.dtype = torch.uint8) -> torch.Tensor: + assert num_bytes > 0 + allocated_bytes = align(num_bytes, self.alignment) + offset = self.num_allocated_bytes + if offset + allocated_bytes > self.buffer.num_bytes: + raise RuntimeError(f'BucketBuffer session out of storage, which has only {self.buffer.num_bytes}') + self.num_allocated_bytes += allocated_bytes + return self.buffer.storage.narrow(0, offset, num_bytes).view(dtype) + + def _check_group(self, group: Optional[dist.ProcessGroup]) -> None: + context_idx = self.buffer.get_context_idx(group) + assert self.context_idx in (None, context_idx), \ + 'All collectives in a session must use the same group' + self.context_idx = context_idx + + def allocate(self, + shapes: Union[Sequence[int], Sequence[Sequence[int]]], + dtype: torch.dtype = torch.float32, + group: Optional[dist.ProcessGroup] = None): + """Allocate tensors inside the registered storage.""" + group = self.buffer.groups[self.buffer.get_context_idx(group)] + self._check_group(group) + + is_batched = isinstance(shapes[0], (tuple, torch.Size, list)) + shapes = shapes if is_batched else (shapes,) + tensors = [self._allocate(math.prod(shape) * dtype.itemsize, dtype).view(shape) for shape in shapes] + return tuple(tensors) if is_batched else tensors[0] + + def _collective_func(self, name: str, *args, **kwargs) -> EventOverlap: + kwargs = self._parse_args(name, *args, **kwargs) + group = self.buffer.groups[self.buffer.get_context_idx(kwargs.get('group'))] + self._check_group(group) + rank_idx, num_ranks = group.rank(), group.size() + if name != 'all_gather': + raise NotImplementedError(f'BucketSession.{name}: not supported on Neuron (basic mode)') + + srcs = kwargs.pop('srcs') + dsts = kwargs.pop('dsts', None) + dst_exists = dsts is not None + is_packed, srcs, dsts = BucketBuffer._pack_tensors(srcs, dsts if dst_exists else srcs) + + storage_begin = self.buffer.storage.data_ptr() + storage_end = storage_begin + self.buffer.num_bytes + src_shards = [] + out_buffers = [] + for src, dst in zip(srcs, dsts, strict=True): + if dst_exists: + assert src.dtype == dst.dtype and dst.is_contiguous() + assert dst.nbytes == src.nbytes * num_ranks + num_buffer_bytes = src.nbytes * num_ranks + is_src_registered = storage_begin <= src.data_ptr() < storage_end + is_dst_registered = dst_exists and storage_begin <= dst.data_ptr() < storage_end + if is_src_registered: + assert src.is_contiguous(), 'all_gather registered source must be contiguous' + offset = src.data_ptr() - storage_begin - rank_idx * src.nbytes + buffer = self.buffer.storage.narrow(0, offset, num_buffer_bytes).view(src.dtype) + elif is_dst_registered: + offset = dst.data_ptr() - storage_begin + buffer = self.buffer.storage.narrow(0, offset, num_buffer_bytes).view(src.dtype) + else: + buffer = self._allocate(num_buffer_bytes, src.dtype) + shard = buffer.view(num_ranks, -1)[rank_idx] + if not is_src_registered: + shard.copy_(src.reshape(-1)) + src_shards.append(shard) + out_buffers.append(buffer.view(-1)) + + handle = self._call_original('all_gather', src_shards, group=group) + external_dsts = dsts if dst_exists else None + + hook = handle.hook_after_wait + handle.hook_after_wait = None + + def hook_after_wait(): + gathered = hook() + gathered = [gathered] if isinstance(gathered, torch.Tensor) else list(gathered) + if external_dsts is not None: + for dst, out in zip(external_dsts, gathered, strict=True): + if dst.data_ptr() != out.data_ptr(): + dst.view(-1).copy_(out.view(dst.dtype).view(-1)) + results = external_dsts + else: + results = gathered + return results[0] if is_packed else results + + handle.register_hook_after_wait(hook_after_wait) + return handle + + +class BucketBuffer(BufferBase): + """Contiguous storage plus batched all-gather over `torch.distributed`.""" + + get_comm_stream = staticmethod(comm.get_comm_stream) + + @staticmethod + def _pack_tensors(*tensors): + is_packed = isinstance(tensors[0], torch.Tensor) + tensors = tuple([] if item is None else [item] if isinstance(item, torch.Tensor) else list(item) + for item in tensors) + assert all(len(item) in (0, len(tensors[0])) for item in tensors) + return is_packed, *tensors + + def __init__(self, + group: Union[dist.ProcessGroup, Sequence[dist.ProcessGroup]], + allocation_plan_or_num_bytes: Union[BufferAllocator, int], + sl_idx: Optional[int] = None, + num_gpu_timeout_secs: int = 100, + explicitly_destroy: bool = False, + *, use_huge1g_page: bool = True, + device: Optional[torch.device] = None): + allocation_plan = allocation_plan_or_num_bytes if isinstance(allocation_plan_or_num_bytes, BufferAllocator) else None + num_bytes = allocation_plan_or_num_bytes if allocation_plan is None else allocation_plan.num_bytes + if allocation_plan is not None: + assert not allocation_plan.materialized + assert isinstance(num_bytes, int) and num_bytes > 0 + self.groups = [group] if isinstance(group, dist.ProcessGroup) else list(group) + assert self.groups and len(self.groups) <= 8 + assert len(set(self.groups)) == len(self.groups), 'Duplicate communication groups' + assert num_bytes % comm.get_num_allocation_alignment() == 0 + self.num_bytes = num_bytes + super().__init__(explicitly_destroy) + self.storage = torch.zeros(num_bytes, dtype=torch.uint8, + device=_resolve_device() if device is None else device) + for item in self.groups: + if item.size() > 1: + dist.barrier(item) + if allocation_plan is not None: + allocation_plan.materialize(self.storage) + + def destroy(self) -> None: + super().destroy() + self.storage = None + + def get_theoretical_num_sms(self, op: str, group: Optional[dist.ProcessGroup] = None) -> int: + self.get_context_idx(group) + if op == 'all_gather': + return 0 + if op in ('reduce_scatter', 'all_reduce'): + return 1 + raise NotImplementedError(op) + + def get_context_idx(self, group: Optional[dist.ProcessGroup]) -> int: + assert len(self.groups) == 1 or group is not None, \ + '`group` is required when the buffer has multiple groups' + return 0 if group is None else self.groups.index(group) + + def session(self) -> BucketSession: + """Open a session in which collectives accept unregistered tensors.""" + return BucketSession(self) + + def barrier(self, + wait_comm_stream: bool = True, + with_cpu_sync: bool = False, + sequential: bool = True, + group: Optional[dist.ProcessGroup] = None) -> None: + context_idx = self.get_context_idx(group) + comm.barrier(self, wait_comm_stream, with_cpu_sync, sequential, context_idx) + + def all_gather(self, + srcs: Union[torch.Tensor, Sequence[torch.Tensor]], + dsts: Optional[Union[torch.Tensor, Sequence[torch.Tensor]]] = None, + group: Optional[dist.ProcessGroup] = None, + num_sms: Optional[int] = None) -> EventOverlap: + """Batched all-gather; `wait()` returns the gathered 1-D view(s) with `srcs`' dtype.""" + assert not self.destroyed + is_packed, srcs, dsts = self._pack_tensors(srcs, dsts) + group = self.groups[self.get_context_idx(group)] + num_ranks, rank_idx = group.size(), group.rank() + + storage_begin = self.storage.data_ptr() + storage_end = storage_begin + self.num_bytes + outputs = [] + for idx, src in enumerate(srcs): + assert src.is_contiguous(), 'all_gather source must be contiguous' + if dsts: + dst = dsts[idx] + assert dst.is_contiguous() and dst.dtype == src.dtype + assert dst.nbytes == src.nbytes * num_ranks + out = dst.view(-1) + else: + # In-place: the source must be this rank's shard of a registered region. + offset = src.data_ptr() - storage_begin - rank_idx * src.nbytes + assert storage_begin <= src.data_ptr() < storage_end and offset >= 0, \ + 'in-place all_gather requires a registered source' + out = self.storage.narrow(0, offset, src.nbytes * num_ranks).view(src.dtype) + comm.all_gather_into_tensor(out, src.reshape(-1), group) + outputs.append(out.view(-1)) + + handle = EventOverlap(EventHandle(), extra_tensors=tuple(srcs) + tuple(dsts) + (self.storage,)) + handle.register_hook_after_wait(lambda: outputs[0] if is_packed else outputs) + setattr(handle, 'num_sms', 0) + return handle + + def reduce_scatter(self, srcs, comm_precision: str = 'fp32', scale: float = 1.0, + group: Optional[dist.ProcessGroup] = None, + num_sms: Optional[int] = None) -> EventOverlap: + raise NotImplementedError('BucketBuffer.reduce_scatter: not supported on Neuron (basic mode)') + + def all_reduce(self, srcs, scale: float = 1.0, + group: Optional[dist.ProcessGroup] = None, + num_sms: Optional[int] = None) -> EventOverlap: + raise NotImplementedError('BucketBuffer.all_reduce: not supported on Neuron (basic mode)') diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py new file mode 100644 index 0000000..a8a3fe5 --- /dev/null +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/engram.py @@ -0,0 +1,264 @@ +"""Mirrors `deep_ep.buffers.engram` (`EngramBuffer`) of DeepEP-Ascend. + +What Engram needs from DeepEP: the n-gram embedding tables of every Engram layer are +**sharded over the Engram-parallel ranks** (`models/v4_modules/engram.py:107` keeps only +`num_local_embeddings` rows per rank), and every step each rank must read arbitrary rows of +the *global* table for its own tokens. Upstream serves those reads out of RDMA-registered +storage; this port exchanges the row indices and the rows themselves with all-to-all. + +Index space (confirmed by `tests/engram/test_engram.py:80-95`): the reference global table is +`all_gather(local_storage)`, so a global index `g` of layer `l` lives on rank +`g // num_entries_per_layer[l]`, row `g % num_entries_per_layer[l]`. + +Layout contract (`csrc/buffers/engram.hpp:139-200`): + * `write(storages, sfs)`: `storages[l]` is this rank's `[num_entries_l, hidden]` shard; + `sfs[l]` is the **replicated** `[num_ranks * num_entries_l, num_sf_packs]` table of which + only this rank's shard is kept. + * `fetch(indices)`: `indices` is int32 `[num_layers, num_tokens, num_entries_per_token]`; + the returned per-layer hook yields `[num_tokens, num_entries_per_token * hidden]` + (and, in FP8 mode, `[num_tokens, num_entries_per_token * num_sf_packs]`, column-major + with stride `(1, num_tokens)` when `use_tma_aligned_col_major_sf=True`). + +Basic mode: the request counts are host values, so `fetch` host-syncs once per layer. +""" + +from __future__ import annotations + +import functools +from typing import Callable, List, Optional, Tuple + +import torch +import torch.distributed as dist + +from .. import comm +from .base import BufferBase +from ..utils.indexing import byte_rows, empty_rows, put_rows, take_rows +from ..utils.math import align + + +def _accept_legacy_storage_size(init): + @functools.wraps(init) + def wrapper(self, *args, **kwargs): + if 'num_rdma_storage_size' in kwargs: + if 'num_rdma_storage_bytes' in kwargs: + raise TypeError('Specify only one of num_rdma_storage_bytes and num_rdma_storage_size') + kwargs['num_rdma_storage_bytes'] = kwargs.pop('num_rdma_storage_size') + return init(self, *args, **kwargs) + return wrapper + + +class EngramBuffer(BufferBase): + """Distributed n-gram embedding lookup over a sharded table.""" + + barrier = comm.barrier + get_comm_stream = staticmethod(comm.get_comm_stream) + get_physical_domain_size = comm.get_physical_domain_size + get_logical_domain_size = comm.get_logical_domain_size + + @_accept_legacy_storage_size + def __init__(self, + group: dist.ProcessGroup, + num_gpu_bytes: int, + num_rdma_storage_bytes: int, + use_cpu_rdma_storage: bool = False, + num_allocated_qps: int = 0, + qp_depth: int = 0, + restrict_rd_atomic: bool = False, + allow_hybrid_mode: bool = True, + sl_idx: Optional[int] = None, + num_gpu_timeout_secs: int = 100, + explicitly_destroy: bool = False, + *, use_huge1g_page: bool = True): + """Create the buffer. QP / service-level / CPU-storage knobs are accepted and ignored. + + Unlike Ascend (which rejects `use_cpu_rdma_storage=True`), this port accepts it: + DeepInfer passes it whenever `DEEPINFER_ENGRAM_FORCE_GPU_STORAGE` is off + (`parallel_utils/userbuffer.py:147`). There is no RDMA storage here; the table shard + simply stays on whatever device `write()` is given. + """ + alignment = comm.get_num_allocation_alignment() + assert num_gpu_bytes > 0 and num_gpu_bytes % alignment == 0 + assert num_rdma_storage_bytes > 0 and num_rdma_storage_bytes % alignment == 0 + self.group = group + self.rank_idx = group.rank() + self.num_ranks = group.size() + self.num_gpu_bytes = num_gpu_bytes + self.num_bytes = num_gpu_bytes + self.num_rdma_storage_bytes = num_rdma_storage_bytes + self.use_cpu_rdma_storage = use_cpu_rdma_storage + self.allow_hybrid_mode = allow_hybrid_mode + self.num_allocated_qps = 0 + super().__init__(explicitly_destroy) + + self.num_entries_per_layer: List[int] = [] + self.hidden = 0 + self.num_max_tokens = 0 + self.num_entries_per_token = 0 + self.num_sf_packs = 0 + self.dtype: Optional[torch.dtype] = None + self._storages: List[torch.Tensor] = [] + self._sf_shards: List[torch.Tensor] = [] + + def destroy(self) -> None: + super().destroy() + self._storages = [] + self._sf_shards = [] + + # --- Sizing --------------------------------------------------------------------- + + @staticmethod + def get_storage_size_hint(group: dist.ProcessGroup, + num_entries_per_layer: List[int], hidden: int, + num_max_tokens: int, num_entries_per_token: int, + dtype: torch.dtype = torch.bfloat16, + num_sf_packs: int = 0) -> Tuple[int, int]: + """Allocation-aligned (receive + SF, per-rank storage) sizes; same formula as Ascend.""" + assert num_entries_per_layer and all(num_entries > 0 for num_entries in num_entries_per_layer) + assert hidden > 0 and num_max_tokens > 0 and num_entries_per_token > 0 and num_sf_packs >= 0 + assert dtype in (torch.bfloat16, torch.float8_e4m3fn) + assert (dtype == torch.float8_e4m3fn) == (num_sf_packs > 0) + alignment = comm.get_num_allocation_alignment() + num_recv_bytes = len(num_entries_per_layer) * num_max_tokens * num_entries_per_token * hidden * dtype.itemsize + num_sf_bytes = sum(num_entries_per_layer) * num_sf_packs * torch.int16.itemsize + num_storage_bytes = sum(num_entries_per_layer) * hidden * dtype.itemsize + return align(num_recv_bytes + num_sf_bytes, alignment), align(num_storage_bytes, alignment) + + @staticmethod + def get_theoretical_config(group: dist.ProcessGroup, + num_layers: int, + num_max_tokens: int, + num_entries_per_token: int, + num_max_read_requests_per_qp: int = 16) -> Tuple[int, int, bool]: + """`(num_allocated_qps, qp_depth, restrict_rd_atomic)`; no QPs on Neuron.""" + return 0, 0, False + + def set_config(self, num_entries_per_layer: List[int], hidden: int, + num_max_tokens: int, num_entries_per_token: int, + dtype: torch.dtype = torch.bfloat16, + num_sf_packs: int = 0) -> None: + """Configure the local shard geometry and the maximum fetch geometry.""" + assert not self.destroyed + self.get_storage_size_hint(self.group, num_entries_per_layer, hidden, + num_max_tokens, num_entries_per_token, dtype, num_sf_packs) + self.num_entries_per_layer = list(num_entries_per_layer) + self.hidden = hidden + self.num_max_tokens = num_max_tokens + self.num_entries_per_token = num_entries_per_token + self.num_sf_packs = num_sf_packs + self.dtype = dtype + + # --- Table upload --------------------------------------------------------------- + + def write(self, storages: List[torch.Tensor], + sfs: Optional[List[torch.Tensor]] = None) -> None: + """Copy this rank's table shard (and its slice of the replicated SF table) in.""" + assert not self.destroyed + num_layers = len(self.num_entries_per_layer) + assert num_layers > 0, 'set_config must be called before write' + assert len(storages) == num_layers + assert (sfs is not None) == (self.num_sf_packs > 0) + + self._storages = [] + for layer_idx, storage in enumerate(storages): + assert storage.dim() == 2 and storage.is_contiguous() + assert storage.size(0) == self.num_entries_per_layer[layer_idx] + assert storage.size(1) == self.hidden + assert storage.dtype == self.dtype + self._storages.append(storage.clone()) + + self._sf_shards = [] + if sfs is not None: + for layer_idx, sf in enumerate(sfs): + num_entries = self.num_entries_per_layer[layer_idx] + assert sf.dim() == 2 and sf.is_contiguous() + assert sf.size(0) == self.num_ranks * num_entries and sf.size(1) == self.num_sf_packs + begin = self.rank_idx * num_entries + self._sf_shards.append(sf[begin:begin + num_entries].clone()) + + comm.barrier(self) + + # --- Fetch ---------------------------------------------------------------------- + + def _fetch_layer(self, layer_idx: int, indices: torch.Tensor, + use_tma_aligned_col_major_sf: bool): + """Gather `indices` ([num_tokens, num_entries_per_token], global ids) of one layer.""" + device = indices.device + num_tokens, num_entries_per_token = indices.shape + num_entries = self.num_entries_per_layer[layer_idx] + storage = self._storages[layer_idx] + sf_shard = self._sf_shards[layer_idx] if self._sf_shards else None + num_requests = num_tokens * num_entries_per_token + flat = indices.reshape(-1).to(torch.int64) + owner = torch.div(flat, num_entries, rounding_mode='floor') + local = flat - owner * num_entries + + if self.num_ranks == 1: + rows = take_rows(storage, local) + sf_rows = take_rows(sf_shard, local) if sf_shard is not None else None + else: + # Group the requests by owner rank (stable: ascending request position). + order = torch.argsort(owner, stable=True) + # NEURON: host sync — all-to-all split sizes must be host values. + send_counts = torch.zeros(self.num_ranks, dtype=torch.int64, device=device) + send_counts.index_add_(0, owner, torch.ones_like(owner)) + send_counts = send_counts.tolist() + recv_counts_t = torch.empty(self.num_ranks, dtype=torch.int64, device=device) + comm.all_to_all_single(recv_counts_t, + torch.tensor(send_counts, dtype=torch.int64, device=device), + None, None, self.group) + recv_counts = recv_counts_t.tolist() + num_served = int(sum(recv_counts)) + + send_idx = local[order] + served_idx = torch.empty((num_served,), dtype=torch.int64, device=device) + comm.all_to_all_single(served_idx, send_idx, recv_counts, send_counts, self.group) + served_idx = served_idx.clamp(0, max(num_entries - 1, 0)) + + served_rows = take_rows(storage, served_idx).contiguous() + answer_rows = empty_rows(num_requests, self.hidden, storage.dtype, device) + # Split sizes count rows: `all_to_all_single` splits along dim 0. + comm.all_to_all_single(byte_rows(answer_rows), byte_rows(served_rows), + send_counts, recv_counts, self.group) + rows = empty_rows(num_requests, self.hidden, storage.dtype, device) + put_rows(rows, order, answer_rows) + + sf_rows = None + if sf_shard is not None: + served_sf = take_rows(sf_shard, served_idx).contiguous() + answer_sf = empty_rows(num_requests, self.num_sf_packs, sf_shard.dtype, device) + comm.all_to_all_single(byte_rows(answer_sf), byte_rows(served_sf), + send_counts, recv_counts, self.group) + sf_rows = torch.empty_like(answer_sf) + sf_rows[order] = answer_sf + + fetched = rows.view(num_tokens, num_entries_per_token * self.hidden) + if sf_rows is None: + return fetched + sf_flat = sf_rows.view(num_tokens, num_entries_per_token * self.num_sf_packs) + if use_tma_aligned_col_major_sf: + # Column-major view: stride (1, num_tokens), no token-dimension padding + # (`as_strided` in csrc/buffers/engram.hpp:177). + num_sf_cols = sf_flat.size(1) + fetched_sf = torch.empty(num_sf_cols * num_tokens, dtype=sf_flat.dtype, + device=device).as_strided((num_tokens, num_sf_cols), + (1, num_tokens)) + fetched_sf.copy_(sf_flat) + else: + fetched_sf = sf_flat + return fetched, fetched_sf + + def fetch(self, indices: torch.Tensor, num_qps: int = 0, + use_tma_aligned_col_major_sf: bool = True) -> List[Callable]: + """Fetch all layers; returns one result hook per layer (work happens eagerly here).""" + assert not self.destroyed + assert indices.dim() == 3 and indices.is_contiguous() + assert indices.dtype == torch.int32 + num_layers, num_tokens, num_entries_per_token = indices.shape + assert num_layers == len(self.num_entries_per_layer) + assert num_tokens <= self.num_max_tokens + assert num_entries_per_token == self.num_entries_per_token + assert self._storages, 'write() must be called before fetch()' + + results = [self._fetch_layer(layer_idx, indices[layer_idx], use_tma_aligned_col_major_sf) + for layer_idx in range(num_layers)] + return [(lambda result=result: result) for result in results] diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/ep.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/ep.py new file mode 100644 index 0000000..04e8626 --- /dev/null +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/ep.py @@ -0,0 +1,612 @@ +"""Mirrors `deep_ep.buffers.ep` (`EPBuffer` / `EPHandle`) of DeepEP-Ascend. + +Expert-parallel dispatch/combine in the Ascend **expanded** layout (`do_expand=True`, +one slot per (token, top-k selection) grouped by local expert), implemented with plain +`torch.distributed` all-to-all on the process group DeepInfer passes in. + +Semantics follow `vendor/DeepEP-Ascend`: + * dispatch: `deep_ep/utils/refs.py:dispatch` (token selection/ordering, `-1` masking of + non-local experts) plus `csrc/buffers/ep.hpp:289-342` (expanded sizes, metadata layout). + * combine: `deep_ep/utils/refs.py:combine` with `reduce_in_local=True`, i.e. per-rank + fp32 reduction stored in bf16, then fp32 accumulation of biases + the per-rank partials + in ascending "master" top-k slot order (`csrc/impls/ep/combine_reduce_epilogue.hpp`). + +Basic mode (see NEURON_PORT.md "Scope"): correctness only. Every dispatch needs the +all-to-all split sizes, which are host values, so routing forces one device->host sync per +dispatch (marked `# NEURON: host sync`). Low-latency mode, overlap, `num_sms`/`num_qps` +knobs and the load-balancer (`lb_*`) entry points are accepted-and-ignored or raise. +""" + +from __future__ import annotations + +from typing import Optional, Tuple, Union + +import torch +import torch.distributed as dist + +from .. import comm +from .base import BufferBase +from ..utils.event import EventHandle, EventOverlap +from ..utils.indexing import byte_rows, empty_rows, put_rows, take_rows, zeros_rows +from ..utils.math import align, align_tensor, exclusive_cumsum + +# Upstream exports the index dtype through `deep_ep._C`. +topk_idx_t = torch.int64 + + +def _value_or(value, default): + return default if value is None else value + + +class EPHandle: + """Communication handle returned by `EPBuffer.dispatch`. + + Public attributes mirror the Ascend handle (`deep_ep/buffers/ep.py:21`): + do_expand, num_experts, expert_alignment, num_max_tokens_per_rank, num_sms, topk_idx, + psum_num_recv_tokens_per_rank, psum_num_recv_tokens_per_expert, + num_unaligned_recv_tokens_per_expert, num_recv_tokens_per_expert_list, + recv_src_metadata (`[num_recv_tokens, num_topk + 2]`: expanded slot per top-k, then + `src_token_global_idx`, then the master top-k slot), dst_buffer_slot_idx, dst_gsge_idx, + num_recv_tokens, num_expanded_tokens. + + `dst_buffer_slot_idx` / `dst_gsge_idx` describe Ascend's on-device send layout; this port + does not use a registered buffer, so they are returned as all `-1` placeholders with the + upstream shape/dtype (nothing in DeepInfer reads them). + """ + + def __init__(self, + do_expand: bool, + num_experts: int, expert_alignment: int, + num_max_tokens_per_rank: int, + num_sms: int, + topk_idx: torch.Tensor, + num_recv_tokens: int, + num_expanded_tokens: int, + num_recv_tokens_per_expert_list: list, + psum_num_recv_tokens_per_rank: torch.Tensor, + psum_num_recv_tokens_per_expert: torch.Tensor, + num_unaligned_recv_tokens_per_expert: torch.Tensor, + recv_src_metadata: torch.Tensor, + dst_buffer_slot_idx: torch.Tensor, + dst_gsge_idx: torch.Tensor): + assert topk_idx is not None + + self.do_expand = do_expand + self.num_experts = num_experts + self.expert_alignment = expert_alignment + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.num_sms = num_sms + self.psum_num_recv_tokens_per_rank = psum_num_recv_tokens_per_rank + self.psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert + self.num_unaligned_recv_tokens_per_expert = num_unaligned_recv_tokens_per_expert + self.num_recv_tokens_per_expert_list = num_recv_tokens_per_expert_list + self.recv_src_metadata = recv_src_metadata + self.dst_buffer_slot_idx = dst_buffer_slot_idx + self.dst_gsge_idx = dst_gsge_idx + + self.num_recv_tokens = num_recv_tokens + self.num_expanded_tokens = num_expanded_tokens + + # Routing plan reused by cached dispatch and by combine (private to this port). + self._plan: Optional['_RoutingPlan'] = None + + self._topk_idx = topk_idx + self._topk_idx_version = None if topk_idx.is_inference() else topk_idx._version + + @property + def topk_idx(self) -> torch.Tensor: + assert self._topk_idx_version is None or self._topk_idx._version == self._topk_idx_version, \ + '`topk_idx` must not be modified while the EP handle is in use' + return self._topk_idx + + def deterministic_sort(self, *args, **kwargs) -> None: + """No-op: this port's dispatch layout is already deterministic. + + Slots are assigned in (local expert, source rank, source token, top-k index) order, + so re-running dispatch with the same routing reproduces the same expanded layout. + """ + return None + + +class _RoutingPlan: + """Everything derived from `topk_idx` once, reused by cached dispatch and combine.""" + + __slots__ = ('num_tokens', 'num_topk', 'num_local_experts', 'expert_alignment', + 'send_counts', 'recv_counts', 'num_send_tokens', 'num_recv_tokens', + 'send_token_idx', 'send_row_rank', 'master_per_rank', + 'local_valid', 'slot', 'slot_rows', 'valid_flat', + 'num_expanded_tokens', 'counts_per_expert') + + def __init__(self, **kwargs): + for key, value in kwargs.items(): + setattr(self, key, value) + + +class EPBuffer(BufferBase): + """Expert-parallel dispatch/combine over `torch.distributed` collectives. + + No symmetric memory is registered: `num_bytes` is bookkeeping only, kept so DeepInfer's + buffer-size accounting (`parallel_utils/userbuffer.py:82`) keeps working. + """ + + barrier = comm.barrier + get_comm_stream = staticmethod(comm.get_comm_stream) + get_physical_domain_size = comm.get_physical_domain_size + get_logical_domain_size = comm.get_logical_domain_size + + def __init__(self, + group: dist.ProcessGroup, + num_bytes: Optional[int] = None, + num_max_tokens_per_rank: int = 0, + hidden: int = 0, + num_topk: int = 0, + use_fp8_dispatch: bool = False, + lb_allocation_plan_or_num_bytes=0, + deterministic: bool = False, + allow_hybrid_mode: bool = True, + allow_multiple_reduction: bool = True, + prefer_overlap_with_compute: bool = True, + sl_idx: Optional[int] = None, + num_allocated_qps: int = 0, + num_cpu_timeout_secs: int = 300, num_gpu_timeout_secs: int = 100, + explicitly_destroy: bool = False, + *, use_huge1g_page: bool = True): + assert allow_multiple_reduction, 'Neuron follows Ascend: allow_multiple_reduction=True' + self.group = group + self.rank_idx = group.rank() + self.num_ranks = group.size() + self.allow_hybrid_mode = allow_hybrid_mode + self.allow_multiple_reduction = allow_multiple_reduction + self.prefer_overlap_with_compute = prefer_overlap_with_compute + self.deterministic = deterministic + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.num_allocated_qps = 0 + self.num_lb_bytes = 0 + super().__init__(explicitly_destroy) + if num_bytes is None: + num_bytes = self.get_buffer_size_hint( + group, num_max_tokens_per_rank, hidden, num_topk, use_fp8_dispatch, + allow_hybrid_mode, allow_multiple_reduction) + self.num_bytes = num_bytes + self.num_scaleout_ranks, self.num_scaleup_ranks = self.get_logical_domain_size() + self.scaleout_rank_idx = self.rank_idx // self.num_scaleup_ranks + self.scaleup_rank_idx = self.rank_idx % self.num_scaleup_ranks + self.num_rdma_ranks, self.num_nvlink_ranks = self.get_physical_domain_size() + + # --- Lifecycle ----------------------------------------------------------------- + + def destroy(self) -> None: + super().destroy() + + # --- Static helpers ------------------------------------------------------------ + + @staticmethod + def get_buffer_size_hint(group: dist.ProcessGroup, + num_max_tokens_per_rank: int, hidden: int, + num_topk: int = 0, use_fp8_dispatch: bool = False, + allow_hybrid_mode: bool = True, + allow_multiple_reduction: bool = True) -> int: + """Recommended (2 MB-aligned) buffer size; mirrors `calculate_buffer_size`. + + This port runs on collectives and allocates its outputs per call, so the number is + only an upper bound used by DeepInfer's accounting. + """ + assert allow_multiple_reduction + num_ranks = group.size() + num_topk = 32 if num_topk == 0 else num_topk + elem_size = 1 if use_fp8_dispatch else 2 + num_sf_bytes = 2 * ((hidden + 31) // 32) if use_fp8_dispatch else 0 + per_token = hidden * elem_size + num_sf_bytes + num_topk * (8 + 4) + 64 + num_dispatch_bytes = num_ranks * num_max_tokens_per_rank * per_token + num_combine_bytes = (min(num_topk, num_ranks) * num_max_tokens_per_rank * + (hidden * 2 + num_topk * 4 + 64)) + return align(max(num_dispatch_bytes, num_combine_bytes), comm.get_num_allocation_alignment()) + + @staticmethod + def capture() -> EventHandle: + """Record an event on the current stream (synchronous no-op here).""" + return EventHandle() + + def get_theoretical_num_sms(self, num_experts: int, num_topk: int, + num_scaleout_topk: int = 0, + rdma_gbs: float = 0, nvlink_gbs: float = 0, + sm_read_gbs: float = 0, sm_write_gbs: float = 0) -> int: + """`num_sms` is ignored by this port; return a nonzero placeholder.""" + return 1 + + def get_theoretical_num_qps(self, num_sms: int) -> int: + return 0 + + # --- Routing ------------------------------------------------------------------- + + def _build_plan(self, topk_idx: torch.Tensor, num_experts: int, + num_max_tokens_per_rank: int, expert_alignment: int, + do_cpu_sync: bool): + device = topk_idx.device + num_ranks, rank_idx = self.num_ranks, self.rank_idx + num_tokens, num_topk = topk_idx.shape + num_local_experts = num_experts // num_ranks + assert num_experts % num_ranks == 0 + assert num_tokens <= num_max_tokens_per_rank + + valid = topk_idx >= 0 + dst_rank = torch.where(valid, torch.div(topk_idx, num_local_experts, rounding_mode='floor'), + torch.full_like(topk_idx, -1)) + rank_ids = torch.arange(num_ranks, device=device, dtype=topk_idx.dtype).view(num_ranks, 1, 1) + per_rank_sel = dst_rank.unsqueeze(0) == rank_ids # [R, T, K] + send_mask = per_rank_sel.any(-1) # [R, T] + + # Master top-k slot per (destination rank, token): the last selection landing on that + # rank owns the per-rank partial (`combine_reduce_epilogue.hpp` master_slot_mask). + k_range = torch.arange(num_topk, device=device, dtype=topk_idx.dtype).view(1, 1, num_topk) + master_per_rank = torch.where(per_rank_sel, k_range.expand(num_ranks, num_tokens, num_topk), + torch.full((1, 1, 1), -1, device=device, dtype=topk_idx.dtype) + ).amax(-1) # [R, T] + + # NEURON: host sync — all-to-all split sizes and output shapes are host values. + send_pairs = send_mask.nonzero() # [S, 2], (rank, token) + send_row_rank = send_pairs[:, 0] + send_token_idx = send_pairs[:, 1] + send_counts_t = send_mask.sum(1).to(torch.int64) + send_counts = send_counts_t.tolist() + recv_counts_t = torch.empty_like(send_counts_t) + comm.all_to_all_single(recv_counts_t, send_counts_t, None, None, self.group) + recv_counts = recv_counts_t.tolist() + num_send_tokens = int(sum(send_counts)) + num_recv_tokens = int(sum(recv_counts)) + + # Exchange the per-token routing: only the destination rank's experts stay visible + # (`refs.py:dispatch` masks the others to -1). + send_sel = per_rank_sel[send_row_rank, send_token_idx] # [S, K] + send_topk_idx = torch.where(send_sel, topk_idx[send_token_idx], + torch.full_like(send_sel, -1, dtype=topk_idx.dtype)) + send_src_idx = (send_token_idx + rank_idx * num_max_tokens_per_rank).to(torch.int64) + + recv_topk_idx = torch.empty((num_recv_tokens, num_topk), dtype=torch.int64, device=device) + comm.all_to_all_single(recv_topk_idx, send_topk_idx.to(torch.int64), + recv_counts, send_counts, self.group) + recv_src_idx = torch.empty((num_recv_tokens,), dtype=torch.int64, device=device) + comm.all_to_all_single(recv_src_idx, send_src_idx, recv_counts, send_counts, self.group) + + # Local expanded layout: group the received (token, top-k) pairs by local expert. + expert_base = rank_idx * num_local_experts + local_expert = recv_topk_idx - expert_base + local_valid = (recv_topk_idx >= expert_base) & (recv_topk_idx < expert_base + num_local_experts) + keys = torch.where(local_valid, local_expert, + torch.full_like(local_expert, num_local_experts)).reshape(-1) + counts_all = torch.zeros(num_local_experts + 1, dtype=torch.int64, device=device) + counts_all.index_add_(0, keys, torch.ones_like(keys)) + counts_per_expert = counts_all[:num_local_experts] + aligned_counts = align_tensor(counts_per_expert, expert_alignment) + expert_start = exclusive_cumsum(aligned_counts) + + # Stable sort by expert id => slots ordered by (expert, received row, top-k index). + perm = torch.argsort(keys, stable=True) + sorted_keys = keys[perm] + first_in_sorted = exclusive_cumsum(counts_all) + padded_start = torch.cat([expert_start, expert_start.new_zeros(1)]) + slot_sorted = padded_start[sorted_keys] + \ + (torch.arange(keys.numel(), device=device, dtype=torch.int64) - first_in_sorted[sorted_keys]) + slot_flat = torch.empty_like(slot_sorted) + slot_flat[perm] = slot_sorted + slot = torch.where(local_valid.reshape(-1), slot_flat, + torch.full_like(slot_flat, -1)).view(num_recv_tokens, num_topk) + + if do_cpu_sync: + # NEURON: host sync — exact expanded size. + num_expanded_tokens = int(aligned_counts.sum().item()) + else: + # Worst case, mirroring `csrc/buffers/ep.hpp:317-319`. + num_expanded_tokens = num_ranks * num_max_tokens_per_rank * min(num_topk, num_local_experts) + num_expanded_tokens += (expert_alignment - 1) * num_local_experts + num_expanded_tokens = align(num_expanded_tokens, expert_alignment) + + valid_flat = local_valid.reshape(-1).nonzero(as_tuple=True)[0] # NEURON: host sync + slot_rows = torch.div(valid_flat, num_topk, rounding_mode='floor') + + plan = _RoutingPlan( + num_tokens=num_tokens, num_topk=num_topk, + num_local_experts=num_local_experts, expert_alignment=expert_alignment, + send_counts=send_counts, recv_counts=recv_counts, + num_send_tokens=num_send_tokens, num_recv_tokens=num_recv_tokens, + send_token_idx=send_token_idx, send_row_rank=send_row_rank, + master_per_rank=master_per_rank, + local_valid=local_valid, slot=slot, slot_rows=slot_rows, valid_flat=valid_flat, + num_expanded_tokens=num_expanded_tokens, counts_per_expert=counts_per_expert) + return plan, recv_src_idx, recv_topk_idx, aligned_counts, expert_start + + # --- Dispatch ------------------------------------------------------------------ + + def _move_payload(self, plan: _RoutingPlan, + x: torch.Tensor, sf: Optional[torch.Tensor], + topk_weights: Optional[torch.Tensor], + num_expanded_tokens: int, + use_tma_aligned_col_major_sf: bool): + """All-to-all the token payload and scatter it into the expanded layout.""" + device = x.device + num_topk = plan.num_topk + hidden = x.size(1) + + send_x = take_rows(x, plan.send_token_idx).contiguous() + recv_x = empty_rows(plan.num_recv_tokens, hidden, x.dtype, device) + # Split sizes count rows: `all_to_all_single` splits along dim 0. + comm.all_to_all_single(byte_rows(recv_x), byte_rows(send_x), + plan.recv_counts, plan.send_counts, self.group) + + recv_sf = None + if sf is not None: + num_sf_packs = sf.size(1) + send_sf = take_rows(sf, plan.send_token_idx).contiguous() + recv_sf = empty_rows(plan.num_recv_tokens, num_sf_packs, sf.dtype, device) + comm.all_to_all_single(byte_rows(recv_sf), byte_rows(send_sf), + plan.recv_counts, plan.send_counts, self.group) + + recv_weights = None + if topk_weights is not None: + send_weights = topk_weights[plan.send_token_idx].contiguous().float() + recv_weights = torch.empty((plan.num_recv_tokens, num_topk), dtype=torch.float32, device=device) + comm.all_to_all_single(recv_weights, send_weights, + plan.recv_counts, plan.send_counts, self.group) + + # Expand: one row per valid (received token, top-k) pair, grouped by local expert. + # The whole buffer is zero-initialized, so alignment padding is always bitwise zero + # (Ascend only guarantees that with `do_zero_padding=True`). + slots = plan.slot.reshape(-1)[plan.valid_flat] + expanded_x = zeros_rows(num_expanded_tokens, hidden, x.dtype, device) + if slots.numel(): + put_rows(expanded_x, slots, take_rows(recv_x, plan.slot_rows)) + + expanded_sf = None + if recv_sf is not None: + num_sf_packs = recv_sf.size(1) + if use_tma_aligned_col_major_sf: + # Column-major SF: stride (1, num_expanded_tokens), no token-dim padding + # (`empty_strided` in csrc/buffers/ep.hpp:331; `as_strided` keeps the exact + # strides for degenerate zero-token cases too). + expanded_sf = torch.zeros(num_sf_packs * num_expanded_tokens, + dtype=recv_sf.dtype, device=device).as_strided( + (num_expanded_tokens, num_sf_packs), (1, num_expanded_tokens)) + else: + expanded_sf = torch.zeros((num_expanded_tokens, num_sf_packs), + dtype=recv_sf.dtype, device=device) + if slots.numel(): + expanded_sf[slots] = take_rows(recv_sf, plan.slot_rows) + + expanded_weights = None + if recv_weights is not None: + expanded_weights = torch.zeros((num_expanded_tokens,), dtype=torch.float32, device=device) + if slots.numel(): + expanded_weights[slots] = recv_weights.reshape(-1)[plan.valid_flat] + + return recv_x, expanded_x, expanded_sf, expanded_weights + + def dispatch(self, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + topk_idx: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + num_max_tokens_per_rank: Optional[int] = None, + expert_alignment: Optional[int] = None, + num_sms: int = 0, + num_qps: int = 0, + previous_event: Optional[EventHandle] = None, + async_with_compute_stream: bool = False, + allocate_on_comm_stream: bool = False, + handle: Optional[EPHandle] = None, + do_handle_copy: bool = False, + do_cpu_sync: Optional[bool] = None, + do_expand: bool = False, + do_zero_padding: bool = False, + use_tma_aligned_col_major_sf: bool = False, + defer_epilogue: bool = False): + """Dispatch tokens to the experts' owner ranks; see the Ascend docstring for arguments. + + Returns `(recv_x, None, recv_topk_weights, handle, event)`, or just the `EventOverlap` + when `defer_epilogue=True` (its `current_stream_wait()` yields the first four). + `num_sms` / `num_qps` / `previous_event` / `allocate_on_comm_stream` are ignored. + """ + assert not self.destroyed + assert do_expand, 'Neuron follows Ascend: dispatch only supports do_expand=True' + assert not do_handle_copy, '`do_handle_copy` must be False' + + x, sf = x if isinstance(x, tuple) else (x, None) + assert x.dim() == 2 and x.is_contiguous() + if sf is not None: + assert x.dtype == torch.float8_e4m3fn and sf.dim() == 2 and sf.size(0) == x.size(0) + else: + assert x.dtype == torch.bfloat16 + + is_cached_dispatch = handle is not None + if is_cached_dispatch: + assert topk_idx is None + assert handle.do_expand == do_expand + assert do_cpu_sync is None or not do_cpu_sync, 'Cannot do CPU sync with cached handle' + assert cumulative_local_expert_recv_stats is None + topk_idx = handle.topk_idx + num_max_tokens_per_rank = _value_or(num_max_tokens_per_rank, handle.num_max_tokens_per_rank) + num_experts = _value_or(num_experts, handle.num_experts) + expert_alignment = _value_or(expert_alignment, handle.expert_alignment) + num_sms = handle.num_sms + do_cpu_sync = False + assert (num_experts, expert_alignment, num_max_tokens_per_rank) == \ + (handle.num_experts, handle.expert_alignment, handle.num_max_tokens_per_rank) + elif num_sms == 0: + num_sms = self.get_theoretical_num_sms(num_experts, topk_idx.size(1)) + + num_max_tokens_per_rank = _value_or(num_max_tokens_per_rank, self.num_max_tokens_per_rank) + expert_alignment = _value_or(expert_alignment, 1) + do_cpu_sync = _value_or(do_cpu_sync, True) + assert topk_idx.dtype == torch.int64 and topk_idx.is_contiguous() + assert topk_idx.size(0) == x.size(0) + if topk_weights is not None: + assert topk_weights.shape == topk_idx.shape + + if is_cached_dispatch: + plan = handle._plan + assert plan is not None + assert plan.num_tokens == x.size(0), 'Cached dispatch requires an unchanged token count' + _, expanded_x, expanded_sf, expanded_weights = self._move_payload( + plan, x, sf, topk_weights, handle.num_expanded_tokens, use_tma_aligned_col_major_sf) + result_handle = handle + else: + plan, recv_src_idx, recv_topk_idx, aligned_counts, expert_start = self._build_plan( + topk_idx, num_experts, num_max_tokens_per_rank, expert_alignment, do_cpu_sync) + _, expanded_x, expanded_sf, expanded_weights = self._move_payload( + plan, x, sf, topk_weights, plan.num_expanded_tokens, use_tma_aligned_col_major_sf) + + device = x.device + num_topk = plan.num_topk + num_local_experts = plan.num_local_experts + + psum_per_rank = torch.tensor(plan.recv_counts, dtype=torch.int32, + device=device).cumsum(0).to(torch.int32) + psum_per_expert = (expert_start + plan.counts_per_expert).to(torch.int32) + num_unaligned = plan.counts_per_expert.to(torch.int32) + if cumulative_local_expert_recv_stats is not None: + assert cumulative_local_expert_recv_stats.shape == (num_local_experts,) + cumulative_local_expert_recv_stats.add_( + plan.counts_per_expert.to(cumulative_local_expert_recv_stats.dtype)) + + # Metadata: one row per received token, padded to the worst case without a CPU sync + # so `num_recv_tokens` matches Ascend (`csrc/buffers/ep.hpp:289`). + num_metadata_rows = plan.num_recv_tokens if do_cpu_sync else \ + self.num_ranks * num_max_tokens_per_rank + metadata = torch.full((num_metadata_rows, num_topk + 2), -1, + dtype=torch.int32, device=device) + rows = plan.num_recv_tokens + metadata[:rows, :num_topk] = plan.slot.to(torch.int32) + metadata[:rows, num_topk] = recv_src_idx.to(torch.int32) + k_range = torch.arange(num_topk, dtype=torch.int32, device=device).view(1, num_topk) + metadata[:rows, num_topk + 1] = torch.where( + plan.local_valid, k_range.expand(rows, num_topk), + torch.full((1, 1), -1, dtype=torch.int32, device=device)).amax(-1) + + placeholder = torch.full((plan.num_tokens, num_topk), -1, dtype=torch.int32, device=device) + result_handle = EPHandle( + do_expand, num_experts, expert_alignment, num_max_tokens_per_rank, num_sms, + topk_idx, + num_metadata_rows, + plan.num_expanded_tokens, + aligned_counts.tolist() if do_cpu_sync else [], # NEURON: host sync + psum_per_rank, psum_per_expert, num_unaligned, + metadata, placeholder, placeholder.clone()) + result_handle._plan = plan + + recv_x = (expanded_x, expanded_sf) if expanded_sf is not None else expanded_x + result = (recv_x, None, expanded_weights, result_handle) + + event_overlap = EventOverlap(EventHandle()) + if defer_epilogue: + assert async_with_compute_stream, 'defer_epilogue requires async_with_compute_stream=True' + event_overlap.register_hook_after_wait(lambda: result) + return event_overlap + return *result, event_overlap + + # --- Combine ------------------------------------------------------------------- + + @staticmethod + def _unpack_bias(bias) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + if isinstance(bias, torch.Tensor): + return bias, None + if isinstance(bias, tuple): + assert len(bias) == 2 + return bias + return None, None + + def combine(self, + x: torch.Tensor, + handle: EPHandle, + topk_weights: Optional[torch.Tensor] = None, + bias=None, + num_sms: int = 0, + num_qps: int = 0, + previous_event: Optional[EventHandle] = None, + async_with_compute_stream: bool = False, + allocate_on_comm_stream: bool = False, + defer_epilogue: bool = False): + """Reduce the expanded tokens back to their source ranks. + + Matches `refs.py:combine(..., reduce_in_local=True)`: each receiving rank sums its + slots for one source token in fp32 and stores bf16; the source rank then accumulates + the biases and the per-rank partials in fp32, in ascending master top-k slot order. + """ + assert not self.destroyed + assert handle is not None and handle.recv_src_metadata is not None + plan = handle._plan + assert plan is not None, 'combine requires a handle produced by this port' + assert x.dim() == 2 and x.is_contiguous() and x.dtype == torch.bfloat16 + + device = x.device + num_topk = plan.num_topk + num_tokens = plan.num_tokens + hidden = x.size(1) + bias_0, bias_1 = self._unpack_bias(bias) + + # Per-rank reduction of this rank's slots (fp32 accumulation, bf16 result). + slot = plan.slot # [N, K], -1 where unused + valid = slot >= 0 + safe_slot = slot.clamp(min=0) + partial = torch.zeros((plan.num_recv_tokens, hidden), dtype=torch.float32, device=device) + for k in range(num_topk): # sequential adds to match the reference accumulation order + rows = x[safe_slot[:, k]].float() + partial += torch.where(valid[:, k:k + 1], rows, torch.zeros_like(rows)) + partial = partial.to(torch.bfloat16) + + send_weights = None + if topk_weights is not None: + assert topk_weights.dim() == 1 and topk_weights.size(0) == x.size(0) + assert topk_weights.dtype == torch.float32 + gathered = topk_weights[safe_slot] + send_weights = torch.where(valid, gathered, torch.zeros_like(gathered)) + + # Return the partials to the source ranks (reverse split sizes of dispatch). + back_x = torch.empty((plan.num_send_tokens, hidden), dtype=torch.bfloat16, device=device) + comm.all_to_all_single(byte_rows(back_x), byte_rows(partial.contiguous()), + plan.send_counts, plan.recv_counts, self.group) + back_weights = None + if send_weights is not None: + back_weights = torch.empty((plan.num_send_tokens, num_topk), dtype=torch.float32, device=device) + comm.all_to_all_single(back_weights, send_weights.contiguous(), + plan.send_counts, plan.recv_counts, self.group) + + # Accumulate: biases first, then the partials in ascending master top-k slot order. + acc = torch.zeros((num_tokens, hidden), dtype=torch.float32, device=device) + for bias_tensor in (bias_0, bias_1): + if bias_tensor is not None: + assert bias_tensor.shape == (num_tokens, hidden) and bias_tensor.dtype == x.dtype + acc += bias_tensor.float() + row_master = plan.master_per_rank[plan.send_row_rank, plan.send_token_idx] + for k in range(num_topk): + selected = (row_master == k).nonzero(as_tuple=True)[0] # NEURON: host sync + if selected.numel(): + acc.index_add_(0, plan.send_token_idx[selected], back_x[selected].float()) + combined_x = acc.to(torch.bfloat16) + + combined_topk_weights = None + if back_weights is not None: + combined_topk_weights = torch.zeros((num_tokens, num_topk), dtype=torch.float32, device=device) + combined_topk_weights.index_add_(0, plan.send_token_idx, back_weights) + + result = (combined_x, combined_topk_weights) + event_overlap = EventOverlap(EventHandle()) + if defer_epilogue: + assert async_with_compute_stream, 'defer_epilogue requires async_with_compute_stream=True' + event_overlap.register_hook_after_wait(lambda: result) + return event_overlap + return *result, event_overlap + + # --- Load balancer (EPLB is off in basic mode) --------------------------------- + + def lb_get_theoretical_num_sms(self) -> int: + return 1 + + def lb_prefetch_weights(self, redundant_expert_weights, expert_weights, + redundancy_mapping, num_sms: int = 0, + previous_event: Optional[EventHandle] = None) -> EventOverlap: + raise NotImplementedError('EPBuffer.lb_prefetch_weights: not supported on Neuron (basic mode)') + + def lb_reduce_grads(self, redundant_expert_grads, expert_grads, + redundancy_mapping, num_sms: int = 0, + previous_event: Optional[EventHandle] = None) -> EventOverlap: + raise NotImplementedError('EPBuffer.lb_reduce_grads: not supported on Neuron (basic mode)') diff --git a/deepinfer/neuron_ops/vendor/deep_ep/comm.py b/deepinfer/neuron_ops/vendor/deep_ep/comm.py new file mode 100644 index 0000000..1158275 --- /dev/null +++ b/deepinfer/neuron_ops/vendor/deep_ep/comm.py @@ -0,0 +1,118 @@ +"""Mirrors `deep_ep.comm` of DeepEP-Ascend, implemented with `torch.distributed`. + +Upstream borrows / creates an HCCL communicator and exposes the physical topology of the +Ascend box. On Neuron there is no symmetric-memory runtime in this basic port: every buffer +runs plain `torch.distributed` collectives on the process group DeepInfer passes in (the +torch-neuronx ``neuron`` backend on device, ``gloo`` for host correctness runs). + +All ranks are reported as one "direct" (NVLink-equivalent) domain, which is what +`ParallelProcessGroup.is_within_nvl_domain` assumes on `PLATFORM.is_npu` anyway +(deepinfer/model_executor/parallel_utils/parallel_state.py:134). +""" + +from __future__ import annotations + +from typing import Any, List, Optional, Sequence, Tuple + +import torch +import torch.distributed as dist + +# Upstream exposes these through `deep_ep._C`; the values mirror the Ascend constants. +_NUM_ALLOCATION_ALIGNMENT_BYTES = 2 * 1024 * 1024 # kNumAllocationAlignmentBytes (2 MB) +_NUM_RDMA_ALIGNMENT_BYTES = 512 +_NUM_TMA_ALIGNMENT_BYTES = 16 + + +def get_num_allocation_alignment() -> int: + """Alignment (bytes) every registered buffer size must respect.""" + return _NUM_ALLOCATION_ALIGNMENT_BYTES + + +def get_num_rdma_alignment() -> int: + """Alignment (bytes) of one transfer inside a registered buffer.""" + return _NUM_RDMA_ALIGNMENT_BYTES + + +def get_num_tma_alignment() -> int: + return _NUM_TMA_ALIGNMENT_BYTES + + +def _group_size(group_or_buffer: Any) -> int: + group = getattr(group_or_buffer, 'group', group_or_buffer) + if isinstance(group, dist.ProcessGroup): + return group.size() + return 1 + + +def get_physical_domain_size(group_or_buffer: Any) -> Tuple[int, int]: + """Return (scale-out, direct-domain) sizes: one Trainium instance is a single domain.""" + return 1, _group_size(group_or_buffer) + + +def get_logical_domain_size(group_or_buffer: Any, + allow_hybrid_mode: Optional[bool] = None) -> Tuple[int, int]: + """Return the logical (scale-out, scale-up) domain sizes; identical to the physical ones.""" + return get_physical_domain_size(group_or_buffer) + + +def destroy_all_managed_hccl_comm() -> None: + """No communicators are managed by this port; kept for API compatibility.""" + return None + + +def get_comm_stream(buffer: Any = None) -> None: + """There is no separate communication stream in this port.""" + return None + + +def barrier(buffer: Any, wait_comm_stream: bool = True, + with_cpu_sync: bool = False, sequential: bool = True, + context_idx: Optional[int] = None) -> None: + """Barrier on the buffer's process group (`dist.barrier`).""" + group = getattr(buffer, 'group', None) + if group is None: + groups = getattr(buffer, 'groups', None) + if groups: + group = groups[0 if context_idx is None else context_idx] + if isinstance(group, dist.ProcessGroup) and group.size() > 1: + dist.barrier(group) + + +# --- Collective helpers (world-size-1 groups never touch `torch.distributed`) --- + + +def all_to_all_single(output: torch.Tensor, input_: torch.Tensor, + output_split_sizes: Optional[Sequence[int]], + input_split_sizes: Optional[Sequence[int]], + group: dist.ProcessGroup) -> torch.Tensor: + """`dist.all_to_all_single` with an exact single-rank fast path. + + Split sizes are host values; computing them from routing data forces a device->host + sync in the caller (see `EPBuffer.dispatch`). + """ + if group.size() == 1: + output.view(-1).copy_(input_.reshape(-1)[:output.numel()]) + return output + dist.all_to_all_single(output, input_.contiguous(), + None if output_split_sizes is None else list(output_split_sizes), + None if input_split_sizes is None else list(input_split_sizes), + group=group) + return output + + +def all_gather_into_tensor(output: torch.Tensor, input_: torch.Tensor, + group: dist.ProcessGroup) -> torch.Tensor: + """`dist.all_gather_into_tensor` with an exact single-rank fast path.""" + if group.size() == 1: + output.view(-1).copy_(input_.reshape(-1)) + return output + dist.all_gather_into_tensor(output, input_.contiguous(), group=group) + return output + + +def all_gather_object(obj: Any, group: dist.ProcessGroup) -> List[Any]: + if group.size() == 1: + return [obj] + gathered: List[Any] = [None] * group.size() + dist.all_gather_object(gathered, obj, group) + return gathered diff --git a/deepinfer/neuron_ops/vendor/deep_ep/utils/__init__.py b/deepinfer/neuron_ops/vendor/deep_ep/utils/__init__.py new file mode 100644 index 0000000..feaa29c --- /dev/null +++ b/deepinfer/neuron_ops/vendor/deep_ep/utils/__init__.py @@ -0,0 +1,5 @@ +"""Mirrors `deep_ep.utils` of DeepEP-Ascend (forward compatibility re-export).""" + +from .event import EventHandle, EventOverlap + +__all__ = ['EventHandle', 'EventOverlap'] diff --git a/deepinfer/neuron_ops/vendor/deep_ep/utils/event.py b/deepinfer/neuron_ops/vendor/deep_ep/utils/event.py new file mode 100644 index 0000000..5843f27 --- /dev/null +++ b/deepinfer/neuron_ops/vendor/deep_ep/utils/event.py @@ -0,0 +1,74 @@ +"""Mirrors `deep_ep.utils.event` / `deep_ep._C.EventHandle` of DeepEP-Ascend. + +On Neuron everything in this package runs synchronously on the calling thread, so the +event objects are no-ops that keep the upstream interface: `EventHandle.current_stream_wait()` +does nothing, and `EventOverlap.current_stream_wait()` just runs the registered epilogue hook. +""" + +from __future__ import annotations + +from typing import Any, Callable, Optional, Tuple + +import torch + + +class EventHandle: + """A recorded "event". Nothing is asynchronous here, so waiting is a no-op.""" + + def current_stream_wait(self) -> None: + return None + + # Upstream `_C.EventHandle` also exposes these; keep them for compatibility. + def synchronize(self) -> None: + return None + + def elapsed_time(self, other: 'EventHandle') -> float: + return 0.0 + + +class EventOverlap: + """Wait for an event and run a deferred epilogue once on the current stream. + + Attributes: + event: the recorded event (always synchronous on Neuron). + extra_tensors: tensors kept alive by this wrapper. + """ + + def __init__(self, event: Optional[EventHandle] = None, + extra_tensors: Optional[Tuple[torch.Tensor, ...]] = None) -> None: + self.event = event + self.extra_tensors = extra_tensors + self._release_handle_by_call = False + self.hook_after_wait: Optional[Callable] = None + + def current_stream_wait(self, release_handle: bool = False) -> Any: + assert self.event is not None + self.event.current_stream_wait() + + result = None + if self.hook_after_wait is not None: + result = self.hook_after_wait() + self.hook_after_wait = None + + if release_handle: + self.event = None + return result + + def wait(self) -> Any: + return self.current_stream_wait() + + def register_hook_after_wait(self, hook_after_wait: Callable) -> None: + assert self.hook_after_wait is None, 'A hook is already registered on this `EventOverlap`' + self.hook_after_wait = hook_after_wait + + def __call__(self, release_handle: bool = False) -> 'EventOverlap': + self._release_handle_by_call = release_handle + return self + + def __enter__(self) -> Any: + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + if self.event is not None: + self.current_stream_wait(release_handle=self._release_handle_by_call) + self._release_handle_by_call = False diff --git a/deepinfer/neuron_ops/vendor/deep_ep/utils/indexing.py b/deepinfer/neuron_ops/vendor/deep_ep/utils/indexing.py new file mode 100644 index 0000000..bab9c1d --- /dev/null +++ b/deepinfer/neuron_ops/vendor/deep_ep/utils/indexing.py @@ -0,0 +1,58 @@ +"""Row gather/scatter helpers that also work for narrow dtypes (FP8). + +FP8 (`torch.float8_e4m3fn`) tensors do not support every indexing / fill op on CPU or on +`neuron`, and DeepEP only ever moves FP8 rows around bit-exactly. These helpers therefore +route such dtypes through a `uint8` view of the row bytes; everything else is indexed +directly. Mirrors the "move bytes, never reinterpret" behaviour of the Ascend kernels. +""" + +from __future__ import annotations + +import torch + +_BYTE_PATH_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + + +def _needs_byte_path(dtype: torch.dtype) -> bool: + return dtype in _BYTE_PATH_DTYPES + + +def byte_rows(x: torch.Tensor) -> torch.Tensor: + """`[rows, cols]` -> `[rows, cols * itemsize]` uint8 view (requires contiguity). + + `view(torch.uint8)` already rescales the last dimension, which keeps this correct for + zero-row tensors (an extra `view(rows, -1)` would be ambiguous there). + """ + assert x.dim() == 2 and x.is_contiguous() + return x.view(torch.uint8) + + +def zeros_rows(num_rows: int, num_cols: int, dtype: torch.dtype, device) -> torch.Tensor: + """Zero-filled `[num_rows, num_cols]` tensor, FP8-safe.""" + if _needs_byte_path(dtype): + return torch.zeros((num_rows, num_cols * dtype.itemsize), dtype=torch.uint8, + device=device).view(dtype) + return torch.zeros((num_rows, num_cols), dtype=dtype, device=device) + + +def empty_rows(num_rows: int, num_cols: int, dtype: torch.dtype, device) -> torch.Tensor: + """Uninitialized `[num_rows, num_cols]` tensor, FP8-safe.""" + if _needs_byte_path(dtype): + return torch.empty((num_rows, num_cols * dtype.itemsize), dtype=torch.uint8, + device=device).view(dtype) + return torch.empty((num_rows, num_cols), dtype=dtype, device=device) + + +def take_rows(src: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """`src[index]` on the row axis, FP8-safe.""" + if not _needs_byte_path(src.dtype): + return src[index] + return byte_rows(src.contiguous())[index].contiguous().view(src.dtype) + + +def put_rows(dst: torch.Tensor, index: torch.Tensor, src: torch.Tensor) -> None: + """`dst[index] = src` on the row axis, FP8-safe (`dst` must be contiguous for FP8).""" + if not _needs_byte_path(dst.dtype): + dst[index] = src + return + byte_rows(dst)[index] = byte_rows(src.contiguous()) diff --git a/deepinfer/neuron_ops/vendor/deep_ep/utils/math.py b/deepinfer/neuron_ops/vendor/deep_ep/utils/math.py new file mode 100644 index 0000000..bca245d --- /dev/null +++ b/deepinfer/neuron_ops/vendor/deep_ep/utils/math.py @@ -0,0 +1,83 @@ +"""Mirrors `deep_ep.utils.math` of DeepEP-Ascend (pure-torch helpers, no device assumptions).""" + +from __future__ import annotations + +from typing import Tuple + +import torch + + +def calc_diff(x: torch.Tensor, y: torch.Tensor) -> float: + x, y = x.double() + 1, y.double() + 1 + denominator = (x * x + y * y).sum() + sim = 2 * (x * y).sum() / denominator + return (1 - sim).item() + + +def safe_div(a, b) -> float: + try: + return a / b + except ZeroDivisionError: + if a == 0: + return 0 + raise + + +def ceil_div(x: int, y: int) -> int: + return (x + y - 1) // y + + +def align(x: int, y: int) -> int: + return ceil_div(x, y) * y + + +def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 + m, n = x.shape + aligned_n = align(n, 128) + x_padded = torch.nn.functional.pad(x, (0, aligned_n - n), mode='constant', value=0) + x_padded_view = x_padded.view(m, -1, 128) + x_amax = x_padded_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) + return (x_padded_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( + m, aligned_n)[:, :n].contiguous(), (x_amax / 448.0).view(m, -1) + + +def per_token_cast_back(x_fp8: torch.Tensor, x_scales: torch.Tensor) -> torch.Tensor: + if x_fp8.numel() == 0: + return x_fp8.to(torch.bfloat16) + + assert x_fp8.dim() == 2 + m, n = x_fp8.shape + aligned_n = align(n, 128) + x_fp8_padded = torch.nn.functional.pad(x_fp8, (0, aligned_n - n), mode='constant', value=0) + if x_scales.dtype == torch.int: + x_scales = x_scales.view(dtype=torch.uint8).to(torch.int) << 23 + x_scales = x_scales.view(dtype=torch.float) + x_fp32_padded = x_fp8_padded.to(torch.float32).view(x_fp8.shape[0], -1, 128) + x_scales = x_scales.view(x_fp8.shape[0], -1, 1) + return (x_fp32_padded * x_scales).view(x_fp8_padded.shape).to(torch.bfloat16)[:, :n].contiguous() + + +def count_bytes(*tensors) -> int: + total = 0 + for t in tensors: + if isinstance(t, (tuple, list)): + total += count_bytes(*t) + elif t is not None: + total += t.numel() * t.element_size() + return total + + +def exclusive_cumsum(x: torch.Tensor) -> torch.Tensor: + """Exclusive prefix sum of a 1-D tensor, same dtype/device.""" + out = torch.zeros_like(x) + if x.numel() > 1: + out[1:] = torch.cumsum(x[:-1], dim=0) + return out + + +def align_tensor(x: torch.Tensor, alignment: int) -> torch.Tensor: + """Element-wise `align()` for integer tensors.""" + if alignment <= 1: + return x + return ((x + (alignment - 1)) // alignment) * alignment diff --git a/deepinfer/neuron_ops/vendor/flash_attn/__init__.py b/deepinfer/neuron_ops/vendor/flash_attn/__init__.py new file mode 100644 index 0000000..3620140 --- /dev/null +++ b/deepinfer/neuron_ops/vendor/flash_attn/__init__.py @@ -0,0 +1 @@ +"""Minimal ``flash_attn`` stand-in for the ViT encoder (mirrors deepinfer/npu_ops/vit_attn.py).""" diff --git a/deepinfer/neuron_ops/vendor/flash_attn/cute/__init__.py b/deepinfer/neuron_ops/vendor/flash_attn/cute/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/deepinfer/neuron_ops/vendor/flash_attn/cute/interface.py b/deepinfer/neuron_ops/vendor/flash_attn/cute/interface.py new file mode 100644 index 0000000..f0ed16c --- /dev/null +++ b/deepinfer/neuron_ops/vendor/flash_attn/cute/interface.py @@ -0,0 +1,34 @@ +"""``flash_attn.cute.interface._flash_attn_fwd`` for the ViT: non-causal varlen attention. + +Same contract as the Ascend replacement (deepinfer/npu_ops/vit_attn.py): q/k/v are TND +([total_tokens, heads, head_dim]); ``cu_seqlens_q`` is a host list of cumulative segment ends +(InputMetadataVL.vit_attn_cu_seqlens on the is_npu path); each segment attends only to itself. +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.nn.functional as F + + +def _flash_attn_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: list[int] | None = None, + **kwargs: Any, +) -> tuple[torch.Tensor, None]: + ends = list(cu_seqlens_q) if cu_seqlens_q is not None else [q.shape[0]] + scale = 1.0 / math.sqrt(q.shape[-1]) + out = torch.empty_like(q) + start = 0 + for end in ends: + if end > start: + # [T, H, D] -> [H, T, D] + qs, ks, vs = (t[start:end].transpose(0, 1).float() for t in (q, k, v)) + out[start:end] = F.scaled_dot_product_attention(qs, ks, vs, scale=scale).transpose(0, 1).to(q.dtype) + start = end + return out, None