Neuron basic mode — part 2: deep_ep, sampling kernels, fixes

- vendor/deep_ep: EPBuffer dispatch/combine with torch.distributed uneven
  all_to_all (host-synced split sizes), EngramBuffer sharded lookup,
  BucketBuffer all_gather (CP-only reduce paths stubbed); verified against
  the DeepEP-Ascend references at world size 1/2/4.
- tilelang_ops/neuron sampling: random_sampling (asc two-level fixed-point
  sampler, Philox), fused_sample (logprobs / do_sample), sampler metadata,
  sum_smaller_probs, structured-output bitmask; DSpark/MTP/varlen kernels
  raise NotImplementedError. All 42 asc *_tl factories now have a Neuron
  counterpart.
- vendor/flash_attn: SDPA varlen stand-in for the ViT (mirrors npu_ops).
- transfer_to_neuron: keep `with torch.device(...)` working on patched
  factories; fix manual_seed recursion in host mode.
- engram: gather the packed ue8m0 sf as bytes on Neuron.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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")

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")

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("<i", struct.pack("<f", x))[0]
+
+
+_LO_BITS = _f32_bits(MIN_PROB)
+_HI_BITS = _f32_bits(_SEARCH_HI)
+
+
+def compute_logprobs(
+    logits: torch.Tensor,
+    temperatures: torch.Tensor,
+    top_p: torch.Tensor,
+    top_k: torch.Tensor,
+) -> 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

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

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")

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")

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")

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

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

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:

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)))

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'

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

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

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)')

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]

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)')

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

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']

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

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())

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

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)."""

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