diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py index 5b103acd3802..60a2b18e0397 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py @@ -15,8 +15,10 @@ round-trips (required by the autotuner cache). * ``trtllm::cute_dsl_megamoe_nvfp4_blackwell`` is the registered torch custom op that the ``MegaMoECuteDsl`` backend calls from - ``run_moe``. It runs ``AutoTuner.choose_one`` once per call to pick - the best tactic and forwards to the runner. + ``run_moe``. ``tactic_autotune=True`` (bench-only, enabled via the + ``MEGAMOE_TACTIC_AUTOTUNE=1`` env var read by the backend) runs + ``AutoTuner.choose_one`` per call; the default bypasses the AutoTuner + and uses the deterministic token-bucket heuristic tactic. The backend never instantiates :class:`Sm100MegaMoENvfp4Runner` directly; this mirrors how ``CuteDslFusedMoE`` only consumes @@ -29,7 +31,8 @@ from __future__ import annotations import dataclasses -import time +import functools +import weakref from typing import Any, List, Optional, Tuple import torch @@ -37,6 +40,7 @@ from tensorrt_llm.logger import logger from ..._utils import get_sm_version +from ...math_utils import ceil_div, pad_up from ..autotuner import ( AutoTuner, ConstraintSpec, @@ -49,13 +53,21 @@ from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..utils import get_last_power_of_2_num_tokens_buckets, last_positive_power_of_2 + +def _import_megamoe_kernel(): + """Lazy import so non-SM100 / no-cutlass-dsl envs can still import this module.""" + from ..cute_dsl_kernels.mega_moe_nvfp4 import import_kernel + from ..cute_dsl_kernels.mega_moe_nvfp4.token_comm import CombineFormat + + return import_kernel(), CombineFormat + + __all__ = [ - "DEFAULT_MEGAMOE_TACTIC", "IS_MEGAMOE_OP_AVAILABLE", "MEGAMOE_OP_UNAVAILABLE_REASON", + "default_megamoe_tactic", "enumerate_megamoe_candidate_tactics", "megamoe_activation_sf_bytes_per_row", - "resolve_megamoe_group_hint", "validate_megamoe_tactic", ] @@ -69,58 +81,198 @@ IS_MEGAMOE_OP_AVAILABLE: bool = False MEGAMOE_OP_UNAVAILABLE_REASON: Optional[str] = None +# (local_ws_ptr, shared_ws_ptr, world_size) keys whose in-kernel barrier +# (counter, signal) was fenced-reset. Module-global because the runner is +# rebuilt every op call; the fence must fire ONCE per key (see forward()). +_MEGAMOE_FENCED_KEYS: set = set() + +# fkey -> (weakref(local UntypedStorage), weakref(shared UntypedStorage)) +# recorded when the fence ran. Fence keys are RAW pointers, so a new +# allocation on a recycled VA (ABA) is only detectable via these storage refs +# (identity-stable, unlike the provider's throwaway per-call views). +_MEGAMOE_FENCED_KEY_WS_REFS: dict = {} + + +def _megamoe_prune_fenced_keys(dead_keys) -> None: + """Drop ``dead_keys`` from ``_MEGAMOE_FENCED_KEYS`` and their recorded + workspace storage weakrefs (single helper so the two cannot drift).""" + _MEGAMOE_FENCED_KEYS.difference_update(dead_keys) + for _key in dead_keys: + _MEGAMOE_FENCED_KEY_WS_REFS.pop(_key, None) + + +def _megamoe_fence_key_live(fkey: Tuple[int, int, int]) -> bool: + """True iff ``fkey`` was fenced AND both recorded workspace storages are + still alive. + + A dead weakref means this pointer is a NEW allocation on a recycled VA: + skipping its fence would pair a virgin barrier side with the other side's + persisted phase (dispatch barrier spins forever). A new workspace + re-fences on EVERY rank (dead-ref prune or plain set miss), so the + collective fence cannot desync. Host-only; capture-safe. + """ + if fkey not in _MEGAMOE_FENCED_KEYS: + return False + refs = _MEGAMOE_FENCED_KEY_WS_REFS.get(fkey) + if refs is not None and refs[0]() is not None and refs[1]() is not None: + return True + _megamoe_prune_fenced_keys({fkey}) + return False + + +# Shared-workspace pointers whose pairings were replayed under CUDA-graph +# capture. A LATER eager fence must never zero such a shared side: captured +# pairings cannot re-fence -> phase-decoupled barrier -> silent all-rank hang. +_MEGAMOE_CAPTURED_SHARED_PTRS: set = set() + +# True once ANY MegaMoE launch was captured into a CUDA graph. Before that +# the whole local-workspace cache is safely evictable; afterwards only +# tuning-latched entries may go. See ``release_megamoe_profiling_scratch``. +_MEGAMOE_GRAPH_CAPTURE_SEEN: bool = False + # --------------------------------------------------------------------------- -# Tactic representation +# Tactic representation (v3: 8-tuple perf knobs) # --------------------------------------------------------------------------- # # A tactic is a tuple of JSON-friendly primitives (lists / ints / bools / -# strings) so it round-trips through ``json.dumps``/``json.loads`` *and* -# ``eval(repr(tactic))`` — both are required by ``TunableRunner`` cache -# serialization. Order matches the kernel constructor kwargs. +# strings / nested tuples) so it round-trips through ``json.dumps`` *and* +# ``eval(repr(tactic))`` (required by ``TunableRunner`` cache serialization). +# Order matches the kernel constructor kwargs. +# +# (mma_tiler_mnk, # list[int] of length 3 (M decides use_2cta) +# cluster_shape_mnk, # list[int] of length 3 +# group_hint, # int, TUNED (>=512); NOT max_active_clusters +# load_balance_mode, # str: "static" | "atomic_counter" +# token_back_mode, # "epi_warps" | "standalone_warps" | "reuse_dispatch_warps" +# use_bulk_fc2_store, # bool (bulk store valid only with epi_warps) +# flag_batch, # int >= 1 (standalone_warps requires == 1) +# epi_flag_batch) # (int, int) fc1/fc2 done-counter publish batch # -# (mma_tiler_mnk, # list[int] of length 3 -# cluster_shape_mnk, # list[int] of length 3 -# use_2cta_instrs, # bool -# resolved_group_hint, # int (always resolved before cache lookup) -# load_balance_mode, # str: "static" | "atomic_counter" -# use_bf16_redg) # bool: form A (False) vs form B (True) +# Derived / out-of-tuple: +# use_2cta_instrs = (mma_tiler_mnk[0] == 256), derived in _build_kernel. +# in_kernel_fc2_reduce -- functional config in runner unique_id, NOT a perf knob. # # Tuple wrapping makes the tactic hashable, which AutoTuner needs for the -# tactics cache. Lists nested inside the tuple are reconstructed from -# JSON intact. +# tactics cache. + +_TACTIC_LEN = 8 + +# Kernel-side ceiling: ``flag_batch`` is hard-checked ``[1, 32]`` and +# ``epi_flag_batch`` entries are SILENTLY clamped to ``[1, 32]``; reject > 32 +# here so a hand-supplied tactic fails fast instead of silently running at 32. +_FLAG_BATCH_MAX = 32 + + +def _unpack_tactic(tactic: Tuple) -> Tuple: + """Return the tactic's 8 fields in canonical order -- the single source of + truth for the field layout; every consumer unpacks through here. Plain + positional unpack, no validation (see :func:`validate_megamoe_tactic`).""" + ( + mma_tiler, + cluster_shape, + group_hint, + load_balance_mode, + token_back_mode, + use_bulk_fc2_store, + flag_batch, + epi_flag_batch, + ) = tactic + return ( + mma_tiler, + cluster_shape, + group_hint, + load_balance_mode, + token_back_mode, + use_bulk_fc2_store, + flag_batch, + epi_flag_batch, + ) -DEFAULT_MEGAMOE_TACTIC: Tuple[List[int], List[int], bool, int, str, bool] = ( - [128, 128, 256], - [1, 1, 1], - False, - 1, # placeholder; the launcher always resolves group_hint first - "static", - False, -) +def default_megamoe_tactic(num_tokens: int) -> Tuple: + """Deterministic token-bucket fallback tactic (autotune disabled / + cache miss / tactic=-1); never profiled by the autotuner.""" + if num_tokens <= 1024: + # decode winner: N128 only helps epi_warps + bulk. + return ([256, 128, 256], [2, 1, 1], 512, "static", "epi_warps", True, 1, (1, 1)) + if num_tokens <= 8192: + return ([256, 256, 256], [2, 1, 1], 512, "static", "epi_warps", True, 4, (1, 1)) + # prefill: atomic_counter only helps the very large tail (>=16384). + return ( + [256, 256, 256], + [2, 1, 1], + 512, + "atomic_counter" if num_tokens >= 16384 else "static", + "reuse_dispatch_warps", + False, + 8, + (2, 4), + ) + + +# --------------------------------------------------------------------------- +# Curated tuning space (~36 candidates per token bucket), filtered through +# ``validate_megamoe_tactic``. MUST be identical on every EP rank: the MERGE +# tuning sweep runs the candidate list in lockstep across ranks, so there is +# deliberately NO runtime knob selecting a different space. +# --------------------------------------------------------------------------- +# token_back_mode -> the only legal use_bulk_fc2_store (bulk binds to epi_warps). +_TOKEN_BACK_STORE_BINDING: dict = { + "epi_warps": True, + "reuse_dispatch_warps": False, + "standalone_warps": False, +} -# Candidate tactic geometries derived from the upstream functional test -# matrix ``moe_nvfp4_swapab/run_mega_tests.sh`` (M01..M20). Each entry is -# ``(mma_tiler_mnk, cluster_shape_mnk, use_2cta_instrs)``. Other tactic -# fields (load_balance_mode, use_bf16_redg) are intentionally constrained -# here. Expanding either axis needs the backend buffer contract and tests to -# move with it. -_RUN_MEGA_TESTS_CANDIDATE_GEOMETRIES: Tuple[ - Tuple[Tuple[int, int, int], Tuple[int, int, int], bool], ... -] = ( - ((128, 128, 256), (1, 1, 1), False), - ((256, 256, 256), (2, 1, 1), True), - ((256, 256, 256), (4, 1, 1), True), +# (mma_tiler_mnk, cluster_shape_mnk) geometries. use_2cta is derived (M==256). +# N64, M128/1-CTA and cluster_m4 never won a bucket upstream. +_GEOMETRIES: Tuple[Tuple[Tuple[int, int, int], Tuple[int, int, int]], ...] = ( + ((256, 256, 256), (2, 1, 1)), + ((256, 128, 256), (2, 1, 1)), ) +# token_back modes scanned (standalone_warps dropped: never a sole winner). +_TOKEN_BACK_MODES: Tuple[str, ...] = ("epi_warps", "reuse_dispatch_warps") + # Load-balance modes supported by the integrated fused FC12 path (see # ImplDesc.__post_init__ in fc1_fc2_fuse_sched.py). # ``clc`` is intentionally excluded -- it routes through a separate # scheduler class not wired through the fused FC12 kernel here. _LOAD_BALANCE_MODE_CANDIDATES: Tuple[str, ...] = ("static", "atomic_counter") +# Quantized combine (fp8/fp4) and form-B reject the bulk fc2 store at kernel +# construction; this is their deterministic default and sizing-probe tactic. +_MEGAMOE_NONBULK_STANDALONE_TACTIC: Tuple = ( + [256, 256, 256], + [2, 1, 1], + 512, + "static", + "standalone_warps", + False, + 1, + (2, 4), +) + +# group_hint saturates ~512; omission (-> ~74) is worst. +_GROUP_HINTS: Tuple[int, ...] = (512, 1024) + +_FLAG_BATCHES: Tuple[int, ...] = (1, 4, 8) + +# epi_flag_batch is fixed per token bucket, NOT a free scan axis; see +# ``_epi_flag_batch_for_tokens``. +_EPI_FLAG_BATCH_SMALL: Tuple[int, int] = (1, 1) +_EPI_FLAG_BATCH_LARGE: Tuple[int, int] = (2, 4) +_EPI_FLAG_BATCH_TOKEN_THRESHOLD = 8192 + + +def _epi_flag_batch_for_tokens(num_tokens: int) -> Tuple[int, int]: + """Token-bucket epi_flag_batch: ``<=8192 -> (1,1)``, ``>8192 -> (2,4)``.""" + if num_tokens > _EPI_FLAG_BATCH_TOKEN_THRESHOLD: + return _EPI_FLAG_BATCH_LARGE + return _EPI_FLAG_BATCH_SMALL + + # Kernel-construction knobs locked until the backend owns the corresponding # runtime buffer / scheduler contracts. _LOCKED_KERNEL_KWARGS = { @@ -151,11 +303,10 @@ def megamoe_activation_sf_bytes_per_row(hidden_size: int) -> int: if hidden_size <= 0 or hidden_size % 32 != 0: raise ValueError(f"hidden_size must be a positive multiple of 32, got {hidden_size}") # ceil(hidden / 16) rounded up to multiples of 4 FP8 columns - # (= `round_up(ceil(hidden_size / scaling_vector_size), 4)`), matching + # (= `pad_up(ceil_div(hidden_size, scaling_vector_size=16), 4)`), matching # the kernel's TMA load width and the ``can_implement`` hidden_size # alignment rule. - sf_cols = (hidden_size + 15) // 16 - return ((sf_cols + 3) // 4) * 4 + return pad_up(ceil_div(hidden_size, 16), 4) def validate_megamoe_tactic(tactic: Tuple) -> None: @@ -170,13 +321,21 @@ def validate_megamoe_tactic(tactic: Tuple) -> None: SupportedMmaTileN, ) - if (not isinstance(tactic, tuple)) or len(tactic) != 6: + if (not isinstance(tactic, tuple)) or len(tactic) != _TACTIC_LEN: raise ValueError( - f"MegaMoE tactic must be a 6-tuple, got {type(tactic).__name__}={tactic!r}" + f"MegaMoE tactic must be an {_TACTIC_LEN}-tuple, got " + f"{type(tactic).__name__} len={len(tactic) if isinstance(tactic, tuple) else 'NA'}={tactic!r}" ) - (mma_tiler, cluster_shape, use_2cta, resolved_group_hint, load_balance_mode, use_bf16_redg) = ( - tactic - ) + ( + mma_tiler, + cluster_shape, + group_hint, + load_balance_mode, + token_back_mode, + use_bulk_fc2_store, + flag_batch, + epi_flag_batch, + ) = _unpack_tactic(tactic) if (not isinstance(mma_tiler, (list, tuple))) or len(mma_tiler) != 3: raise ValueError(f"mma_tiler_mnk must be a 3-tuple/list, got {mma_tiler!r}") @@ -195,6 +354,8 @@ def validate_megamoe_tactic(tactic: Tuple) -> None: f"_validate_mma_*." ) + use_2cta = mma_tiler[0] == 256 + if (not isinstance(cluster_shape, (list, tuple))) or len(cluster_shape) != 3: raise ValueError(f"cluster_shape_mnk must be a 3-tuple/list, got {cluster_shape!r}") if cluster_shape[2] != 1: @@ -215,87 +376,96 @@ def validate_megamoe_tactic(tactic: Tuple) -> None: f"{cluster_shape[0] * cluster_shape[1]}." ) - if not isinstance(use_2cta, bool): - raise ValueError(f"use_2cta_instrs must be bool, got {use_2cta!r}.") - expected_2cta = mma_tiler[0] == 256 - if use_2cta != expected_2cta: - raise ValueError( - f"use_2cta_instrs must be {expected_2cta} for mma_tiler_mnk[0]={mma_tiler[0]}, got {use_2cta}." - ) - if cluster_shape[0] % (2 if use_2cta else 1) != 0: raise ValueError( f"cluster_shape_mnk[0] ({cluster_shape[0]}) must be divisible by " - f"{(2 if use_2cta else 1)} when use_2cta_instrs={use_2cta}." + f"{(2 if use_2cta else 1)} when use_2cta_instrs={use_2cta} " + f"(derived from mma_tiler_mnk[0]={mma_tiler[0]})." ) - if (not isinstance(resolved_group_hint, int)) or resolved_group_hint <= 0: - raise ValueError( - f"resolved_group_hint must be a positive int (resolved before " - f"cache lookup), got {resolved_group_hint!r}." - ) + if (not isinstance(group_hint, int)) or isinstance(group_hint, bool) or group_hint < 512: + raise ValueError(f"group_hint must be an int >= 512, got {group_hint!r}.") if load_balance_mode not in {"static", "atomic_counter"}: raise ValueError( f"load_balance_mode must be 'static' or 'atomic_counter', got {load_balance_mode!r}." ) - if not isinstance(use_bf16_redg, bool): - raise ValueError(f"use_bf16_redg must be bool, got {use_bf16_redg!r}.") - if use_bf16_redg: + if token_back_mode not in {"epi_warps", "standalone_warps", "reuse_dispatch_warps"}: raise ValueError( - "use_bf16_redg=True (form-B in-kernel top-k reduction) is not " - "wired in MegaMoECuteDsl yet. The backend allocates form-A " - "combine_output with shape (T, top_k, hidden) and performs the " - "top-k reduction on the host, so cached or manually supplied " - "form-B tactics are rejected." + f"token_back_mode must be one of 'epi_warps' / 'standalone_warps' / " + f"'reuse_dispatch_warps', got {token_back_mode!r}." ) + if not isinstance(use_bulk_fc2_store, bool): + raise ValueError(f"use_bulk_fc2_store must be bool, got {use_bulk_fc2_store!r}.") + if use_bulk_fc2_store and token_back_mode != "epi_warps": # nosec B105 + raise ValueError( + f"use_bulk_fc2_store=True requires token_back_mode='epi_warps', got " + f"{token_back_mode!r} (non-epi token-back forces non-bulk fc2 store)." + ) + if mma_tiler[1] == 128 and token_back_mode != "epi_warps": # nosec B105 + raise ValueError( + f"mma_tiler_mnk[1]=128 (N128) is only recommended with " + f"token_back_mode='epi_warps'; got {token_back_mode!r}." + ) -def resolve_megamoe_group_hint(cluster_shape_mnk: Tuple[int, int, int]) -> int: - """Resolve ``group_hint=None`` to ``HardwareInfo().get_max_active_clusters``. - - The kernel uses ``group_hint`` as a construction-time constant - (``Sm100MegaMoEKernel.__init__``); caching under ``None`` would - produce a wrong cache key. Falls back to 1 on hosts without - CUDA / Cutlass DSL so the tactic remains JSON-serializable. - """ - cluster_size = cluster_shape_mnk[0] * cluster_shape_mnk[1] * cluster_shape_mnk[2] - if cluster_size <= 0: - cluster_size = 1 - try: - from cutlass.utils import HardwareInfo + if ( + (not isinstance(flag_batch, int)) + or isinstance(flag_batch, bool) + or not (1 <= flag_batch <= _FLAG_BATCH_MAX) + ): + raise ValueError( + f"flag_batch must be an int in [1, {_FLAG_BATCH_MAX}] (kernel " + f"TokenInPullTokenBackPush hard limit), got {flag_batch!r}." + ) + if token_back_mode == "standalone_warps" and flag_batch != 1: # nosec B105 + raise ValueError( + f"token_back_mode='standalone_warps' requires flag_batch == 1, got {flag_batch}." + ) - return max(1, int(HardwareInfo().get_max_active_clusters(cluster_size))) - except Exception: # pragma: no cover - host without CUDA / Cutlass DSL - return 1 + if (not isinstance(epi_flag_batch, (tuple, list))) or len(epi_flag_batch) != 2: + raise ValueError(f"epi_flag_batch must be a 2-tuple (fc1, fc2), got {epi_flag_batch!r}.") + for v in epi_flag_batch: + if (not isinstance(v, int)) or isinstance(v, bool) or not (1 <= v <= _FLAG_BATCH_MAX): + raise ValueError( + f"epi_flag_batch entries must be ints in [1, {_FLAG_BATCH_MAX}] " + f"(epilogue clamp range), got {epi_flag_batch!r}." + ) -def enumerate_megamoe_candidate_tactics() -> List[Tuple]: - """Return the integrated candidate tactic list, fully resolved. +def enumerate_megamoe_candidate_tactics(num_tokens: int) -> List[Tuple]: + """Return the curated candidate tactic list for the current token bucket. - Each candidate has its ``resolved_group_hint`` stamped to the value - returned by ``HardwareInfo.get_max_active_clusters`` for that - cluster shape. Form A is the only supported reduction mode until the - backend wires the form-B output buffer and reduction path. + Cartesian product over the curated dimension ranges, filtered through + ``validate_megamoe_tactic``. ``epi_flag_batch`` is fixed per token + bucket (see ``_epi_flag_batch_for_tokens``). """ + epi_flag_batch = _epi_flag_batch_for_tokens(num_tokens) + candidates: List[Tuple] = [] - for mma_tiler, cluster_shape, use_2cta in _RUN_MEGA_TESTS_CANDIDATE_GEOMETRIES: - for load_balance_mode in _LOAD_BALANCE_MODE_CANDIDATES: - tactic = ( - list(mma_tiler), - list(cluster_shape), - use_2cta, - resolve_megamoe_group_hint(cluster_shape), - load_balance_mode, - False, - ) - try: - validate_megamoe_tactic(tactic) - except ValueError as e: - logger.debug(f"[MegaMoE] dropping candidate tactic {tactic!r}: {e}") - continue - candidates.append(tactic) + for mma_tiler, cluster_shape in _GEOMETRIES: + for token_back_mode in _TOKEN_BACK_MODES: + use_bulk = _TOKEN_BACK_STORE_BINDING[token_back_mode] + for load_balance_mode in _LOAD_BALANCE_MODE_CANDIDATES: + for group_hint in _GROUP_HINTS: + for flag_batch in _FLAG_BATCHES: + tactic = ( + list(mma_tiler), + list(cluster_shape), + int(group_hint), + load_balance_mode, + token_back_mode, + use_bulk, + int(flag_batch), + tuple(epi_flag_batch), + ) + try: + validate_megamoe_tactic(tactic) + except ValueError as e: + logger.debug(f"[MegaMoE] dropping candidate tactic {tactic!r}: {e}") + continue + candidates.append(tactic) return candidates @@ -310,8 +480,6 @@ def enumerate_megamoe_candidate_tactics() -> List[Tuple]: # op is unregistered. try: import cutlass - import cutlass.cute as cute - import cutlass.torch as cutlass_torch import torch.distributed as dist import torch.distributed._symmetric_memory as torch_symm_mem from cutlass.cute.nvgpu import cpasync, tcgen05 # noqa: F401 @@ -328,7 +496,6 @@ def enumerate_megamoe_candidate_tactics() -> List[Tuple]: Nvfp4BlockSize, # noqa: F401 SfPaddingBlock, ) - from ..cute_dsl_kernels.mega_moe_nvfp4.sym_buffer import SymBufferHost IS_MEGAMOE_OP_AVAILABLE = True except Exception as _megamoe_import_err: # pragma: no cover - env-specific @@ -357,51 +524,89 @@ def enumerate_megamoe_candidate_tactics() -> List[Tuple]: # multi-rank and is supplied by the caller (the MegaMoECuteDsl backend's # MegaMoeSymmMemProvider carves it out of the rendezvous'd buffer). _MEGAMOE_LOCAL_WORKSPACE_CACHE: dict = {} + # Cache keys allocated during the autotune sweep. Each candidate has a + # DISTINCT key (workspace SIZE varies with the tactic) and a multi-GiB + # workspace; persisting all would OOM, so each candidate's workspace is + # freed when the next one allocates (only the winner is reused). + _MEGAMOE_TUNING_WORKSPACE_KEYS: set = set() + + @functools.lru_cache(maxsize=1) + def _cute_launch_helpers(): + """Import the cutlass-dsl launch surface once per process and build the + tensor/pointer converters (hoisted out of forward(); lazy so the + module imports without cutlass-dsl / a GPU).""" + import cutlass + import cutlass.cute as cute + import cutlass.torch as cutlass_torch + import cutlass.utils as cutlass_utils + from cutlass.cute.typing import AddressSpace + + from ..cute_dsl_kernels.mega_moe_nvfp4.sym_buffer import SymBufferHost + + def to_cute(t, assumed_align=16): + ct = cutlass_torch.from_dlpack(t, assumed_align=assumed_align) + return ct.mark_layout_dynamic(leading_dim=cutlass_torch.get_leading_dim(t)) + + def to_cute_ptr(t, assumed_align=16): + return cute.runtime.make_ptr( + cutlass.Uint8, t.data_ptr(), AddressSpace.gmem, assumed_align=assumed_align + ) + + return to_cute, to_cute_ptr, SymBufferHost, cute, cutlass_utils + + def _evict_latched_tuning_workspaces() -> None: + """Pop every tuning-latched local workspace and prune its fence keys. + + Fence keys hold the raw local ptr; if the allocator hands the freed + block to the next candidate, a stale key would ABA-skip the + barrier-reset fence and hang the dispatch barrier mid-sweep. + """ + for stale_key in _MEGAMOE_TUNING_WORKSPACE_KEYS: + stale_ws = _MEGAMOE_LOCAL_WORKSPACE_CACHE.pop(stale_key, None) + if stale_ws is not None: + _stale_ptr = int(stale_ws.data_ptr()) + _megamoe_prune_fenced_keys({k for k in _MEGAMOE_FENCED_KEYS if k[0] == _stale_ptr}) + _MEGAMOE_TUNING_WORKSPACE_KEYS.clear() def _get_or_alloc_local_workspace( - kernel, cache_key: Tuple, device: torch.device + kernel, cache_key: Tuple, device: torch.device, latch_in_tuning_mode: bool = True ) -> torch.Tensor: cached = _MEGAMOE_LOCAL_WORKSPACE_CACHE.get(cache_key) if cached is not None: return cached + # Cache MISS = new (tactic, shape): release any workspace latched for + # a previously-profiled candidate (never relaunched; profiling is + # eager, so no captured graph holds the buffer). Within one + # candidate's multi-launch profiling the key HITs, so nothing is + # freed mid-candidate; outside tuning the latch set is empty. + if _MEGAMOE_TUNING_WORKSPACE_KEYS: + _evict_latched_tuning_workspaces() local_bytes, _ = kernel.get_workspace_sizes() - # MUST be zero-initialised: the local workspace embeds Int32 - # atomic counters (l1_arrival_count, fc1_done_counter, - # grid_sync_counter) whose spin_wait expects v >= positive - # threshold; a stray negative byte from ``torch.empty`` makes - # the wait unsatisfiable and hangs the kernel at 100% SM. - local_workspace = torch.zeros(local_bytes, dtype=torch.uint8, device=device) + # Only the atomic counters need zeroing (a stray negative byte hangs + # spin_wait at 100% SM); the multi-GiB bulk is overwritten every launch + # (full ``torch.zeros`` = huge wasted memset). Zero the leading counter + # prefix plus the persisted ``nvlink_barrier_counter``, deliberately + # OUTSIDE that prefix so the device tail-reset never touches it. + local_workspace = torch.empty(local_bytes, dtype=torch.uint8, device=device) + _lead = int(kernel.require_zero_workspace_leading_bytes[0]) + local_workspace[:_lead].zero_() + _bc_off = int(kernel._local_offsets["nvlink_barrier_counter"]) + _bc_n = int(kernel._local_region_by_name["nvlink_barrier_counter"].nbytes) + local_workspace[_bc_off : _bc_off + _bc_n].zero_() _MEGAMOE_LOCAL_WORKSPACE_CACHE[cache_key] = local_workspace + # ``latch_in_tuning_mode=False`` (tactic_autotune opted OUT): launches + # inside a global autotune() are real fallback-tactic runs, not sweep + # candidates, so their workspaces must persist. + if latch_in_tuning_mode and AutoTuner.get().is_tuning_mode: + _MEGAMOE_TUNING_WORKSPACE_KEYS.add(cache_key) return local_workspace - def _zero_local_workspace_preserving_phase(local_workspace, kernel) -> None: - """Per-launch zero of the local workspace that PRESERVES the - self-priming ``nvlink_barrier_counter`` region (multi-rank EP path). - - The kernel's reusable phase-flip NVLink barrier keeps its cross-rank - ``nvlink_barrier_signal`` (in the symmetric shared workspace, which is - NOT re-zeroed per launch) in lockstep with this per-rank - ``nvlink_barrier_counter``. Re-zeroing the counter while the signal is - not reset would decouple the phase and deadlock the barrier. Every - other local counter (l1_arrival_count, fc1_done_counter, - fc2_done_counter, expert_send_count, ...) still needs a per-launch - reset, so we zero the whole buffer except the counter's byte range. - """ - off = int(kernel._local_offsets["nvlink_barrier_counter"]) - nbytes = int(kernel._local_region_by_name["nvlink_barrier_counter"].nbytes) - total = local_workspace.numel() - if off > 0: - local_workspace[:off].zero_() - end = off + nbytes - if end < total: - local_workspace[end:].zero_() - # ----- Symmetric-memory provider (NVSHMEM-equivalent) ------------------- # # PyTorch's ``torch.distributed._symmetric_memory`` is an NVSHMEM-equivalent # symmetric-heap provider built on cuMem APIs. It exposes per-rank buffer # pointers (``handle.buffer_ptrs``) which we use to populate the - # ``SymBufferHost(base_addr, offsets, rank_idx, num_max_ranks)`` payload + # ``SymBufferHost(offsets, rank_idx, num_max_ranks)`` payload # the MegaMoE kernel expects. # # We allocate ONE large symmetric buffer per (group, layout_key) and @@ -416,9 +621,6 @@ def _zero_local_workspace_preserving_phase(local_workspace, kernel) -> None: # ``_MEGA_MOE_SYMM_BUFFER_CACHE`` in ``mega_moe_deepgemm.py``). _MEGAMOE_SYMM_PROVIDER_CACHE: dict = {} - def _round_up_to(value: int, alignment: int) -> int: - return ((value + alignment - 1) // alignment) * alignment - @dataclasses.dataclass class MegaMoeSymmRegions: """User-domain symmetric tensors carved out of a single rendezvous'd @@ -439,7 +641,8 @@ class MegaMoeSymmRegions: # kernel sf_addr formula in dispatch_kernel.py. activation_sf: torch.Tensor # (max_T, sf_bytes_per_row) uint8 (FP8 SF) topk_weights: torch.Tensor # (max_T, num_topk) float32 - combine_output: torch.Tensor # (max_T, num_topk, hidden) output_dtype + # (max_T, 1, hidden): both reduction forms collapse the top-k axis in-op. + combine_output: torch.Tensor shared_workspace: torch.Tensor # (shared_ws_bytes,) uint8 peer_offsets: List[int] # symmetric peer-pointer deltas rank: int @@ -504,6 +707,10 @@ def __init__( self.max_tokens_per_rank = int(max_tokens_per_rank) self.num_topk = int(num_topk) self.output_dtype = output_dtype + # Kernel output is unified (T, hidden), so the symmetric combine + # region is (max_T, 1, hidden). combine_k MUST match the kernel + # ``combine_output`` shape or the region is silently corrupted. + self.combine_k = 1 # Region byte sizes (worst case across launches; staging # writes only the live ``T`` rows). NVFP4 packs 2 elems / byte @@ -515,17 +722,13 @@ def __init__( act_bytes_per_row = hidden_size // 2 sf_bytes_per_row = megamoe_activation_sf_bytes_per_row(hidden_size) topkw_bytes_per_row = num_topk * 4 # float32 - combine_bytes_per_row = num_topk * hidden_size * output_dtype.itemsize + combine_bytes_per_row = self.combine_k * hidden_size * output_dtype.itemsize - act_region = _round_up_to(max_tokens_per_rank * act_bytes_per_row, self._REGION_ALIGN) - sf_region = _round_up_to(max_tokens_per_rank * sf_bytes_per_row, self._REGION_ALIGN) - topkw_region = _round_up_to( - max_tokens_per_rank * topkw_bytes_per_row, self._REGION_ALIGN - ) - combine_region = _round_up_to( - max_tokens_per_rank * combine_bytes_per_row, self._REGION_ALIGN - ) - shared_region = _round_up_to(shared_workspace_bytes, self._REGION_ALIGN) + act_region = pad_up(max_tokens_per_rank * act_bytes_per_row, self._REGION_ALIGN) + sf_region = pad_up(max_tokens_per_rank * sf_bytes_per_row, self._REGION_ALIGN) + topkw_region = pad_up(max_tokens_per_rank * topkw_bytes_per_row, self._REGION_ALIGN) + combine_region = pad_up(max_tokens_per_rank * combine_bytes_per_row, self._REGION_ALIGN) + shared_region = pad_up(shared_workspace_bytes, self._REGION_ALIGN) self._region_offsets: dict = {} self._region_sizes: dict = {} @@ -562,7 +765,7 @@ def __init__( peer_ptr = int(self._handle.buffer_ptrs[r]) self.peer_offsets.append(peer_ptr - local_base) - logger.info( + logger.debug( "[MegaMoeSymmMemProvider] group=%s rank=%d/%d total_bytes=%d " "(activation=%d sf=%d topk_weights=%d combine=%d shared=%d)", self.group_name, @@ -591,6 +794,7 @@ def get_regions(self) -> MegaMoeSymmRegions: hidden = self.hidden_size max_t = self.max_tokens_per_rank top_k = self.num_topk + combine_k = self.combine_k # ``sf_bytes_per_row`` MUST match the byte width used at # allocation time (``__init__`` above) and at the backend's # ``quantize_input`` output: kernel reads @@ -606,7 +810,7 @@ def get_regions(self) -> MegaMoeSymmRegions: ), topk_weights=self._region_view("topk_weights", (max_t, top_k), torch.float32), combine_output=self._region_view( - "combine_output", (max_t, top_k, hidden), self.output_dtype + "combine_output", (max_t, combine_k, hidden), self.output_dtype ), shared_workspace=self._region_view( "shared_workspace", (self._region_sizes["shared_workspace"],), torch.uint8 @@ -625,6 +829,7 @@ def get_megamoe_symm_provider( max_tokens_per_rank: int, num_topk: int, output_dtype: torch.dtype, + combine_format: str, shared_workspace_bytes: int, ) -> MegaMoeSymmMemProvider: """Return a cached provider for (group, layout). The cache is @@ -648,6 +853,7 @@ def get_megamoe_symm_provider( int(max_tokens_per_rank), int(num_topk), str(output_dtype), + str(combine_format), int(shared_workspace_bytes), ) cached = _MEGAMOE_SYMM_PROVIDER_CACHE.get(cache_key) @@ -666,10 +872,137 @@ def get_megamoe_symm_provider( _MEGAMOE_SYMM_PROVIDER_CACHE[cache_key] = provider return provider + # ---- AutoTuner profiling scratch (symmetric, transient) --------------- + # The kernel's SINGLE ``peer_rank_ptr_mapper`` requires ALL cross-rank + # tensors to live in one symmetric allocation at consistent offsets; the + # AutoTuner's regenerated NON-symmetric inputs would peer-map outside any + # symmetric region (multi-rank IMA). So profiling runs on a SEPARATE + # symmetric scratch (never the real staging buffer), shared across + # same-layout layers and freed by ``release_megamoe_profiling_scratch()``. + _MEGAMOE_PROFILING_SCRATCH_CACHE: dict = {} + _ACTIVE_MEGAMOE_PROFILING_SCRATCH = None + # Deferred variant: zero-arg callable returning MegaMoeSymmRegions, set + # when the scratch is not allocated. The profiling pre-hook calls it only + # when REAL profiling launches, so a tuning-mode cache HIT never pays the + # multi-GiB allocation. + _ACTIVE_MEGAMOE_PROFILING_SCRATCH_FACTORY = None + + def get_megamoe_profiling_scratch( + *, + process_group, + world_size: int, + rank: int, + hidden_size: int, + max_tokens_per_rank: int, + num_topk: int, + output_dtype: torch.dtype, + combine_format: str, + shared_workspace_bytes: int, + ): + """Return a cached, transient symmetric scratch provider used ONLY for + AutoTuner profiling; ``None`` for single-rank.""" + if int(world_size) <= 1: + return None + if not hasattr(process_group, "group_name"): + raise RuntimeError( + "get_megamoe_profiling_scratch requires a ProcessGroup with " + ".group_name (mapping.moe_ep_group_pg)." + ) + cache_key = ( + str(process_group.group_name), + int(hidden_size), + int(max_tokens_per_rank), + int(num_topk), + str(output_dtype), + str(combine_format), + int(shared_workspace_bytes), + ) + cached = _MEGAMOE_PROFILING_SCRATCH_CACHE.get(cache_key) + if cached is not None: + return cached + provider = MegaMoeSymmMemProvider( + process_group=process_group, + world_size=world_size, + rank=rank, + hidden_size=hidden_size, + max_tokens_per_rank=max_tokens_per_rank, + num_topk=num_topk, + output_dtype=output_dtype, + shared_workspace_bytes=shared_workspace_bytes, + ) + _MEGAMOE_PROFILING_SCRATCH_CACHE[cache_key] = provider + return provider + + def set_active_megamoe_profiling_scratch(regions) -> None: + """Set (``None`` clears) the :class:`MegaMoeSymmRegions` used for + AutoTuner profiling on the current call (set by the backend around + the op invocation; consumed inside ``choose_one``).""" + global _ACTIVE_MEGAMOE_PROFILING_SCRATCH + _ACTIVE_MEGAMOE_PROFILING_SCRATCH = regions + + def set_active_megamoe_profiling_scratch_factory(factory) -> None: + """Set (``None`` clears) a zero-arg factory returning the profiling + scratch regions; used instead of the eager setter when the scratch is + not allocated. See ``_ACTIVE_MEGAMOE_PROFILING_SCRATCH_FACTORY``.""" + global _ACTIVE_MEGAMOE_PROFILING_SCRATCH_FACTORY + _ACTIVE_MEGAMOE_PROFILING_SCRATCH_FACTORY = factory + + def release_megamoe_profiling_scratch() -> None: + """Free all profiling-scratch symmetric buffers (call after warmup); + profiling re-allocates lazily if it runs again.""" + global _ACTIVE_MEGAMOE_PROFILING_SCRATCH + global _ACTIVE_MEGAMOE_PROFILING_SCRATCH_FACTORY + _ACTIVE_MEGAMOE_PROFILING_SCRATCH = None + _ACTIVE_MEGAMOE_PROFILING_SCRATCH_FACTORY = None + # Prune fence keys pairing with the dying scratch shared buffers: a + # future allocation on the recycled VA must re-fence (ABA). + _dead_shared = { + int(_p.get_regions().shared_workspace.data_ptr()) + for _p in _MEGAMOE_PROFILING_SCRATCH_CACHE.values() + } + if _dead_shared: + _megamoe_prune_fenced_keys({k for k in _MEGAMOE_FENCED_KEYS if k[1] in _dead_shared}) + _MEGAMOE_PROFILING_SCRATCH_CACHE.clear() + # EVICT stale local workspaces -- do not merely clear the latch: the + # latch may hold the LAST LOSER, and a losing default's multi-GiB + # workspace would shrink the KV-cache budget for the process lifetime. + # Before the FIRST capture nothing holds raw pointers, so the whole + # cache is evictable; after a capture only latched entries go. + if not _MEGAMOE_GRAPH_CAPTURE_SEEN: + for _key, _ws in list(_MEGAMOE_LOCAL_WORKSPACE_CACHE.items()): + _ptr = int(_ws.data_ptr()) + _megamoe_prune_fenced_keys({k for k in _MEGAMOE_FENCED_KEYS if k[0] == _ptr}) + del _MEGAMOE_LOCAL_WORKSPACE_CACHE[_key] + _MEGAMOE_TUNING_WORKSPACE_KEYS.clear() + else: + _evict_latched_tuning_workspaces() + + def reset_megamoe_workspace_state() -> None: + """TEST/BENCH ONLY: drop ALL process-global MegaMoE workspace state. + + PRECONDITION: every CUDA graph that replayed a MegaMoE launch is + destroyed and nothing will launch on the dropped workspaces. Fence / + captured-ptr sets are cleared TOGETHER with the buffer caches (stale + entries on recycled addresses would ABA-skip the fence or false-positive + the capture guard). Must run in lockstep on ALL EP ranks. + """ + global _ACTIVE_MEGAMOE_PROFILING_SCRATCH + global _ACTIVE_MEGAMOE_PROFILING_SCRATCH_FACTORY + global _MEGAMOE_GRAPH_CAPTURE_SEEN + _ACTIVE_MEGAMOE_PROFILING_SCRATCH = None + _ACTIVE_MEGAMOE_PROFILING_SCRATCH_FACTORY = None + _MEGAMOE_PROFILING_SCRATCH_CACHE.clear() + _MEGAMOE_SYMM_PROVIDER_CACHE.clear() + _MEGAMOE_LOCAL_WORKSPACE_CACHE.clear() + _MEGAMOE_TUNING_WORKSPACE_KEYS.clear() + _MEGAMOE_FENCED_KEYS.clear() + _MEGAMOE_FENCED_KEY_WS_REFS.clear() + _MEGAMOE_CAPTURED_SHARED_PTRS.clear() + _MEGAMOE_GRAPH_CAPTURE_SEEN = False + def query_megamoe_shared_workspace_bytes( *, world_size: int, - local_rank: int, num_topk: int, num_experts_per_rank: int, hidden_size: int, @@ -679,43 +1012,43 @@ def query_megamoe_shared_workspace_bytes( tactic: Optional[Tuple] = None, apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, + in_kernel_fc2_reduce: bool = False, + combine_format: str = "bf16", ) -> int: """Probe ``Sm100MegaMoEKernel.get_workspace_sizes()`` for the - shared workspace byte count. The shared workspace size is + shared workspace byte count. The SHARED workspace size is invariant across all candidate tactics and across the codegen-time graph/clamp modes (its regions depend only on world_size / num_experts_per_rank / num_topk / max_tokens_per_rank -- see _build_shared_region_specs in megamoe_kernel.py), so we use the - default tactic for the probe. ``apply_topk_in_fc1`` / ``gate_up_clamp`` - are still threaded so the probe kernel ctor signature is satisfied - and matches the real build. + default 8-tuple tactic for the probe; the remaining kwargs are + threaded only to satisfy the kernel ctor and match the real build. """ - from ..cute_dsl_kernels.mega_moe_nvfp4 import import_kernel if tactic is None: - cluster = tuple(DEFAULT_MEGAMOE_TACTIC[1]) - tactic = ( - list(DEFAULT_MEGAMOE_TACTIC[0]), - list(cluster), - DEFAULT_MEGAMOE_TACTIC[2], - resolve_megamoe_group_hint(cluster), - DEFAULT_MEGAMOE_TACTIC[4], - DEFAULT_MEGAMOE_TACTIC[5], - ) + # Sizing is tactic-invariant, but quantized combine (fp8/fp4) + # rejects the bulk fc2 store at kernel construction, so those + # modes must probe with the non-bulk standalone tactic. + if combine_format != "bf16": + tactic = _MEGAMOE_NONBULK_STANDALONE_TACTIC + else: + tactic = default_megamoe_tactic(0) ( mma_tiler, cluster_shape, - use_2cta, - resolved_group_hint, + group_hint, load_balance_mode, - use_bf16_redg, - ) = tactic - kernel_cls = import_kernel() - probe = kernel_cls( - mma_tiler_mnk=tuple(mma_tiler), + token_back_mode, + use_bulk_fc2_store, + flag_batch, + epi_flag_batch, + ) = _unpack_tactic(tactic) + mma_tiler = tuple(mma_tiler) + common = dict( + mma_tiler_mnk=mma_tiler, cluster_shape_mnk=tuple(cluster_shape), - use_2cta_instrs=bool(use_2cta), - group_hint=int(resolved_group_hint), + use_2cta_instrs=bool(mma_tiler[0] == 256), + group_hint=int(group_hint), token_padding_block=64, sf_padding_block=SfPaddingBlock, load_balance_mode=str(load_balance_mode), @@ -725,34 +1058,35 @@ def query_megamoe_shared_workspace_bytes( hidden_size, ), world_size=int(world_size), - local_rank=int(local_rank), num_topk=int(num_topk), max_tokens_per_rank=int(max_tokens_per_rank), hidden=int(hidden_size), fc2_output_dtype=cutlass.BFloat16, - in_kernel_fc2_reduce=bool(use_bf16_redg), + in_kernel_fc2_reduce=bool(in_kernel_fc2_reduce), + token_back_mode=str(token_back_mode), + non_ubulk_fc2_store=(not bool(use_bulk_fc2_store)), + flag_batch=int(flag_batch), + epi_flag_batch=tuple(epi_flag_batch), apply_topk_in_fc1=bool(apply_topk_in_fc1), gate_up_clamp=(None if gate_up_clamp is None else float(gate_up_clamp)), **_LOCKED_KERNEL_KWARGS, ) + # The probe MUST build the SAME kernel that runs (same combine_format): + # otherwise the provider carves an undersized shared region and the + # combine staging writes OOB (single-rank IMA; EP>1 looks like a hang). + kernel_cls, CombineFormat = _import_megamoe_kernel() + probe = kernel_cls(combine_format=CombineFormat.parse(combine_format), **common) _, shared_bytes = probe.get_workspace_sizes() return int(shared_bytes) - def _to_cute( - tensor: torch.Tensor, - assumed_align: int = 16, - force_static_layout: bool = False, - ) -> "cute.Tensor": - cute_tensor = cutlass_torch.from_dlpack(tensor, assumed_align=assumed_align) - # The local workspace's internal region offsets/strides are codegen-time - # static constants (see megamoe_kernel _layout_regions); marking it - # layout-dynamic invalidates those static accesses and corrupts the - # FC1-output / pool / counter regions. The upstream runner passes the - # local workspace with force_static_layout=True for exactly this reason. - if force_static_layout: - return cute_tensor - leading_dim = cutlass_torch.get_leading_dim(tensor) - return cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) + def _megamoe_autotune_num_tokens(shapes: List[torch.Size]) -> int: + """Shape-derivation rule: the activation token count (input[0] dim0) + drives every other tensor's leading axis. + + Module-scope on purpose: ConstraintSpec hashes include the callable + and the AutoTuner lru_caches on it; a per-call closure would defeat + memoization and leak cache entries.""" + return shapes[0][0] class Sm100MegaMoENvfp4Runner(TunableRunner): """TunableRunner for the ported MegaMoE CuteDSL NVFP4 kernel. @@ -767,13 +1101,6 @@ class Sm100MegaMoENvfp4Runner(TunableRunner): # Module-scope compile cache shared by every runner instance. kernel_cache: dict = {} - # Module-scope tuning-config cache keyed on ``unique_id()``. The op - # rebuilds a runner per call, so an instance-level cache would never - # hit; keeping it at class scope amortizes the config build across - # calls (mirrors the ``tuning_config_cache`` of the CuteDSL - # grouped-gemm runners in ``cute_dsl_custom_ops.py``). - tuning_config_cache: dict = {} - def __init__( self, *, @@ -788,8 +1115,9 @@ def __init__( output_dtype: torch.dtype, apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, - token_back_by_dispatch: bool = False, - non_ubulk_fc2_store: bool = True, + in_kernel_fc2_reduce: bool = False, + combine_format: str = "bf16", + tactic_autotune: bool = False, ) -> None: super().__init__() if (sm_version := get_sm_version()) not in (100, 103): @@ -818,18 +1146,34 @@ def __init__( ) self.max_tokens_per_rank = int(max_tokens_per_rank) self.output_dtype = output_dtype - # Codegen-time graph/clamp modes. They change the generated - # kernel, so they are part of ``unique_id`` (and therefore the - # compile-cache key) -- never per-call runtime kwargs. + # Codegen-time modes: they change the generated kernel, so they are + # part of ``unique_id`` (compile + workspace cache key), never + # per-call runtime kwargs. self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) self.gate_up_clamp = None if gate_up_clamp is None else float(gate_up_clamp) - self.token_back_by_dispatch = bool(token_back_by_dispatch) - self.non_ubulk_fc2_store = bool(non_ubulk_fc2_store) + # Symmetric profiling scratch, set by the op around ``choose_one`` + # so the pre-hook routes cross-rank inputs through symmetric + # memory; None outside tuning and for single-rank. + self._profiling_scratch = None + # Deferred scratch factory; the pre-hook materializes it on the + # first REAL profiling launch (tuning-mode cache HITs never allocate). + self._profiling_scratch_factory = None + self.in_kernel_fc2_reduce = bool(in_kernel_fc2_reduce) + # combine wire format: bf16 / 32e4m3xe8m0 (fp8) / 16e2m1xbf16 (fp4); + # quantized changes the shared_workspace size, so part of unique_id. + self.combine_format = str(combine_format) + # Tactic-autotune opt-in (default OFF). Does NOT change the + # generated kernel, so deliberately EXCLUDED from unique_id / + # _tactic_cache_key: opted-in and opted-out runs share caches. + self.tactic_autotune = bool(tactic_autotune) def unique_id(self): + # local_rank is intentionally excluded: every EP rank must run the + # SAME tactic and MERGE merges timings by cache key across ranks, + # so the key MUST be rank-identical. world_size stays so + # single-rank and multi-rank never share entries. return ( self.world_size, - self.local_rank, self.num_topk, self.num_experts_per_rank, self.hidden_size, @@ -839,8 +1183,8 @@ def unique_id(self): str(self.output_dtype), self.apply_topk_in_fc1, self.gate_up_clamp, - self.token_back_by_dispatch, - self.non_ubulk_fc2_store, + self.in_kernel_fc2_reduce, + self.combine_format, ) def get_valid_tactics( @@ -849,44 +1193,48 @@ def get_valid_tactics( profile: OptimizationProfile, **kwargs, ) -> List[Tuple]: - del inputs, profile, kwargs - return enumerate_megamoe_candidate_tactics() + del profile, kwargs + num_tokens = int(inputs[0].shape[0]) + candidates = enumerate_megamoe_candidate_tactics(num_tokens) + # Non-bulk is HARD for form-B (a bulk store collapses the K routes + # UNSUMMED -> silent wrong output) and fp4 combine (UBLK cannot + # scalar-deref sub-byte data -> kernel raises); fp8 bulk is legal + # and stays tunable. ``_unpack_tactic(t)[5]`` is use_bulk_fc2_store. + if self.in_kernel_fc2_reduce or self.combine_format.startswith("16e2m1"): + candidates = [t for t in candidates if not _unpack_tactic(t)[5]] + return candidates def _autotuner_inputs_pre_hook(self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: """Sanitize ONLY the autotuner-regenerated fake inputs. - ``AutoTuner._prepare_input_tensors`` rebuilds fresh fake tensors - for the dynamic / constraint inputs -- activation (0), - activation_sf (1), topk_idx (2), topk_weights (3), - combine_output (11) -- and passes every STATIC input through BY - REFERENCE (``tensor = inputs[i]`` for non-dynamic dims). The static - inputs here are the caller's REAL weight-side tensors: fc1_weight - (4), fc1_weight_sf (5), fc2_weight (6), fc2_weight_sf (7), - fc1_alpha (8), fc2_alpha (9), fc1_norm_const (10). - - Therefore this hook must mirror ``CuteDslFusedMoE.inputs_pre_hook``: - only fix up the regenerated tensors and pass the real weights / - scales through untouched. The fresh ``topk_idx`` is filled with - random ints in ``[-5, 4]`` whose out-of-range values index a - per-CTA SMEM histogram + the peer-rank pointer table and trigger - illegal memory access, so we rewrite it to a valid round-robin; - the fresh ``activation_sf`` / ``topk_weights`` are zeroed to keep - the FP8/FP32 epilogue NaN-free (autotuning measures runtime, not - numerics). - - We intentionally do NOT touch indices 4-10. An in-place - ``zero_()`` / ``fill_()`` on those would permanently clobber the - caller's REAL per-expert weight scale factors / alphas (they are - not regenerated), zeroing the weight SF and forcing every - post-tuning forward to emit an all-zero ``combine_output``. The - real weights are already valid (no NaN, non-zero norm_const), so - they need no sanitization. This keeps the hook copy-free. + AutoTuner rebuilds inputs 0-3 and 11; the static inputs (4-10) are + the caller's REAL weights/scales, passed by reference -- filling + those would clobber them. topk_idx becomes a valid round-robin + (random ints index OOB); SF / weights are filled with 1.0, NOT 0 + (SF==0 degenerates the GEMMs and skews tactic timing). """ + # Runs ONLY on a real profiling MISS (a tuning-mode cache HIT + # early-returns in choose_one), so materializing the DEFERRED + # scratch here keeps cache-hitting forwards allocation-free. The + # factory's collective rendezvous is safe: MERGE lockstep tuning + # gets every EP rank to this pre-hook before any coupled launch. + if self._profiling_scratch is None and self._profiling_scratch_factory is not None: + self._profiling_scratch = self._profiling_scratch_factory() + if self._profiling_scratch is None and self.world_size > 1: + # Multi-rank profiling without symmetric scratch would peer-map + # non-symmetric tensors: cross-rank IMA. Fail loud. + raise RuntimeError( + "MegaMoE-CuteDSL autotune profiling requires the " + "symmetric profiling scratch on multi-rank, but none is " + "active (and no scratch factory produced one)." + ) inputs = list(inputs) total_experts = self.num_experts_per_rank * self.world_size if total_experts <= 0: return inputs + # topk_idx is consumed LOCALLY (not peer-mapped), so the + # regenerated tensor stays; just make the fake ids in-range. topk_idx = inputs[2] if isinstance(topk_idx, torch.Tensor) and topk_idx.dim() == 2: T, K = topk_idx.shape @@ -900,17 +1248,41 @@ def _autotuner_inputs_pre_hook(self, inputs: List[torch.Tensor]) -> List[torch.T ).view(T, K) topk_idx.copy_(valid) - # activation_sf (1) and topk_weights (3) are autotuner-regenerated - # fresh tensors; zero them to keep the FC1/FC2 epilogue paths - # NaN-free against random ``uint8`` -> FP8 reinterpretation. The - # weight SF (5, 7) and per-expert alphas (8, 9, 10) are the real, - # already-valid backend tensors and are deliberately left alone. + scratch = self._profiling_scratch + if scratch is not None: + # Multi-rank: route the CROSS-RANK inputs through the symmetric + # scratch (sliced to the regenerated token count) so the peer + # mapper resolves inside a symmetric region; staging untouched. + m = int(inputs[0].shape[0]) + act = scratch.activation[:m] + act.copy_(inputs[0].view(torch.uint8)) + inputs[0] = act.view(inputs[0].dtype) + + sf = scratch.activation_sf[:m] + sf.fill_(0x38) # raw byte == FP8 1.0 + inputs[1] = sf if inputs[1].dtype == torch.uint8 else sf.view(inputs[1].dtype) + + w = scratch.topk_weights[:m] + w.fill_(1.0) + inputs[3] = w + + comb = scratch.combine_output[:m] + comb.zero_() # form-B accumulates onto live rows; form-A overwrites + inputs[11] = comb + return inputs + + # Single-rank / scratch-disabled: sanitize in place. fill_(1.0) on + # the FP8 view writes byte 0x38; on a uint8 view write the raw + # byte. Do NOT fill_(0x38) on the FP8 view -- that is 56.0. activation_sf = inputs[1] if isinstance(activation_sf, torch.Tensor): - activation_sf.zero_() + if activation_sf.dtype == torch.float8_e4m3fn: + activation_sf.fill_(1.0) # FP8 1.0 (exact, byte 0x38) + else: + activation_sf.view(torch.uint8).fill_(0x38) # raw byte = FP8 1.0 topk_weights = inputs[3] if isinstance(topk_weights, torch.Tensor): - topk_weights.zero_() + topk_weights.fill_(1.0) return inputs @@ -922,20 +1294,11 @@ def get_tuning_config(self) -> TuningConfig: autotuner does not double-enumerate tile sizes for independent token axes. - The config is cached at class scope keyed on ``unique_id()``. - Every field below is a constant except ``inputs_pre_hook``. + Rebuilt on every call, NEVER cached across runners: + ``inputs_pre_hook`` is bound to THIS runner's ``_profiling_scratch``; + a cached hook bound to a dead runner would mix non-symmetric + profiling tensors with the live runner's peer offsets (IMA + hang). """ - key = self.unique_id() - cached = self.__class__.tuning_config_cache.get(key) - if cached is not None: - return cached - - # Constraints reuse the runner's own shape-derivation rules - # (the activation token count drives every other tensor's - # leading axis). We pass shape-derivation lambdas that pull - # the runtime ``num_tokens`` from input[0]. - def _num_tokens(shapes: List[torch.Size]) -> int: - return shapes[0][0] config = TuningConfig( dynamic_tensor_specs=( @@ -944,59 +1307,54 @@ def _num_tokens(shapes: List[torch.Size]) -> int: ), ), constraint_specs=( - ConstraintSpec(1, 0, _num_tokens), # activation_sf - ConstraintSpec(2, 0, _num_tokens), # topk_idx - ConstraintSpec(3, 0, _num_tokens), # topk_weights + ConstraintSpec(1, 0, _megamoe_autotune_num_tokens), # activation_sf + ConstraintSpec(2, 0, _megamoe_autotune_num_tokens), # topk_idx + ConstraintSpec(3, 0, _megamoe_autotune_num_tokens), # topk_weights # combine_output moved from idx 8 -> 11 after inserting # fc1_alpha(8) / fc2_alpha(9) / fc1_norm_const(10). - ConstraintSpec(11, 0, _num_tokens), # combine_output + ConstraintSpec(11, 0, _megamoe_autotune_num_tokens), # combine_output ), - # ``inputs_pre_hook`` is a bound method of THIS runner - # instance, yet caching the whole config across instances is - # safe: the hook only reads ``num_experts_per_rank`` and - # ``world_size`` (see ``_autotuner_inputs_pre_hook``), and both - # are part of ``unique_id()`` -- so every runner that maps to - # the same cache key has a functionally identical hook. The - # first instance for a given key is retained alive by this - # bound method (one runner object per distinct layer config, - # negligible). Mirrors how the CuteDSL grouped-gemm runners - # cache a ``helper.inputs_pre_hook`` keyed on ``unique_id()``. inputs_pre_hook=self._autotuner_inputs_pre_hook, - use_cold_l2_cache=True, + # MUST stay False for multi-rank: cold-L2 rotation clone()s the + # profiling inputs into NON-symmetric buffers, so the kernel's + # peer mapping (clone ptr + scratch peer_offsets) resolves + # outside any symmetric region -> cross-rank IMA while tuning. + use_cold_l2_cache=False, + # Pin the bucket ladder to the per-rank token ceiling instead + # of whatever (KV-clamped, possibly non-pow2) num_tokens the + # warmup forward fed. Mirrors trtllm-gen's tune_max_num_tokens. + tune_max_num_tokens=self.max_tokens_per_rank, # CUDA Graph capture cannot reproduce MegaMoE's runtime # peer-pointer table / dispatch-counter view and would # spin inside the captured barrier when the autotuner's # L2-cache buffers rotate. Plain repeat-loop profiling # is correct and only marginally slower. use_cuda_graph=False, - # FUSED_COMM hard requirement: every EP rank must run - # the same compiled tactic per chunk so the NVLink - # dispatch barrier and peer pointer mapping line up. - # PARALLEL strategy keeps tactic selection lockstep - # across ranks (same as every multi-rank CuteDSL op in - # ``cute_dsl_custom_ops.py``). - distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL, + # Ranks couple inside the fused kernel, so every EP rank must + # profile the SAME tactic in lockstep with identical launch + # counts; MERGE does that and all-gathers timings so all ranks + # converge on one global-best tactic. PARALLEL would split + # tactics across ranks -> contaminated timing + barrier desync. + distributed_tuning_strategy=DistributedTuningStrategy.MERGE, ) - self.__class__.tuning_config_cache[key] = config return config def _build_kernel(self, tactic: Tuple): ( mma_tiler, cluster_shape, - use_2cta, - resolved_group_hint, + group_hint, load_balance_mode, - use_bf16_redg, - ) = tactic - from ..cute_dsl_kernels.mega_moe_nvfp4 import import_kernel - - kernel_cls = import_kernel() - return kernel_cls( - mma_tiler_mnk=tuple[Any, ...](mma_tiler), + token_back_mode, + use_bulk_fc2_store, + flag_batch, + epi_flag_batch, + ) = _unpack_tactic(tactic) + common = dict( + mma_tiler_mnk=tuple(mma_tiler), cluster_shape_mnk=tuple(cluster_shape), - use_2cta_instrs=bool(use_2cta), - group_hint=int(resolved_group_hint), + use_2cta_instrs=bool(mma_tiler[0] == 256), + group_hint=int(group_hint), token_padding_block=64, sf_padding_block=SfPaddingBlock, load_balance_mode=str(load_balance_mode), @@ -1006,81 +1364,52 @@ def _build_kernel(self, tactic: Tuple): self.hidden_size, ), world_size=self.world_size, - local_rank=self.local_rank, num_topk=self.num_topk, max_tokens_per_rank=self.max_tokens_per_rank, hidden=self.hidden_size, fc2_output_dtype=cutlass.BFloat16, - in_kernel_fc2_reduce=bool(use_bf16_redg), + in_kernel_fc2_reduce=self.in_kernel_fc2_reduce, + token_back_mode=str(token_back_mode), + non_ubulk_fc2_store=(not bool(use_bulk_fc2_store)), + flag_batch=int(flag_batch), + epi_flag_batch=tuple(epi_flag_batch), apply_topk_in_fc1=self.apply_topk_in_fc1, gate_up_clamp=self.gate_up_clamp, - token_back_by_dispatch=self.token_back_by_dispatch, - non_ubulk_fc2_store=self.non_ubulk_fc2_store, **_LOCKED_KERNEL_KWARGS, ) + kernel_cls, CombineFormat = _import_megamoe_kernel() + return kernel_cls(combine_format=CombineFormat.parse(self.combine_format), **common) def _tactic_cache_key(self, tactic: Tuple) -> Tuple: - # Hashable cache key shared by the compile cache and the - # local-workspace cache. ``unique_id()`` already carries - # apply_topk_in_fc1 / gate_up_clamp, so the codegen-time - # graph/clamp modes are part of the cache key without listing - # them again here. + """Hashable cache key over ``unique_id()`` + the FULL 8-tuple. + + Both the compile cache and the local-workspace cache MUST key on + the full tactic: the local workspace SIZE varies with it (non-epi + token_back / atomic_counter add large regions), so an under-keyed + cache would reuse a wrong-sized buffer -> OOB / 100% SM hang. + """ ( mma_tiler, cluster_shape, - use_2cta, - resolved_group_hint, + group_hint, load_balance_mode, - use_bf16_redg, - ) = tactic + token_back_mode, + use_bulk_fc2_store, + flag_batch, + epi_flag_batch, + ) = _unpack_tactic(tactic) return ( self.unique_id(), tuple(mma_tiler), tuple(cluster_shape), - bool(use_2cta), - int(resolved_group_hint), + int(group_hint), str(load_balance_mode), - bool(use_bf16_redg), + str(token_back_mode), + bool(use_bulk_fc2_store), + int(flag_batch), + tuple(epi_flag_batch), ) - def _compile_or_get(self, tactic: Tuple, kernel, runtime_kwargs): - ( - mma_tiler, - cluster_shape, - use_2cta, - resolved_group_hint, - load_balance_mode, - use_bf16_redg, - ) = tactic - cache_key = self._tactic_cache_key(tactic) - compiled = self.__class__.kernel_cache.get(cache_key) - if compiled is not None: - return compiled - compile_kwargs = dict(runtime_kwargs) - hardware_info = cutlass.utils.HardwareInfo() - cluster_size = cluster_shape[0] * cluster_shape[1] * cluster_shape[2] - compile_kwargs["max_active_clusters"] = hardware_info.get_max_active_clusters( - max(cluster_size, 1) - ) - # CuTe DSL compile is the dominant first-launch cost; log - # start/end at INFO so the long compile gap is visible through - # the standard TRT-LLM logger (honors TLLM_LOG_LEVEL). - logger.info( - f"[MegaMoECuteDsl] cute.compile START tactic=" - f"(mma_tiler={mma_tiler}, cluster={cluster_shape}, " - f"use_2cta={use_2cta}, group_hint={resolved_group_hint}, " - f"load_balance={load_balance_mode!r}, use_bf16_redg={use_bf16_redg})" - ) - t_compile_start = time.perf_counter() - compiled = cute.compile(kernel, **compile_kwargs) - t_compile_ms = (time.perf_counter() - t_compile_start) * 1000 - logger.info( - f"[MegaMoECuteDsl] cute.compile DONE in {t_compile_ms:.0f} ms " - f"(cache_keys_now={len(self.__class__.kernel_cache) + 1})" - ) - self.__class__.kernel_cache[cache_key] = compiled - return compiled - def forward( self, inputs: List[torch.Tensor], @@ -1091,17 +1420,15 @@ def forward( **kwargs, ) -> None: del kwargs - t_forward_start = time.perf_counter() - # Resolve fallback tactic. if tactic == -1 or tactic is None: - tactic_t = ( - list(DEFAULT_MEGAMOE_TACTIC[0]), - list(DEFAULT_MEGAMOE_TACTIC[1]), - DEFAULT_MEGAMOE_TACTIC[2], - resolve_megamoe_group_hint(tuple(DEFAULT_MEGAMOE_TACTIC[1])), - DEFAULT_MEGAMOE_TACTIC[4], - DEFAULT_MEGAMOE_TACTIC[5], - ) + num_tokens = int(inputs[0].shape[0]) + # Form-B / quantized combine require non-bulk (see + # get_valid_tactics); this deterministic fallback keeps ALL + # quantized combine non-bulk (the tuner may recover fp8 bulk). + if self.in_kernel_fc2_reduce or self.combine_format != "bf16": + tactic_t = _MEGAMOE_NONBULK_STANDALONE_TACTIC + else: + tactic_t = default_megamoe_tactic(num_tokens) elif isinstance(tactic, list): tactic_t = tuple(tactic) else: @@ -1122,6 +1449,13 @@ def forward( fc1_norm_const, combine_output, ) = inputs[:12] + if self._profiling_scratch is not None: + # Profiling launch: with the DEFERRED factory, choose_one's + # kwargs were bound to the STAGING buffers before the pre-hook + # materialized the scratch -- rebind both cross-rank args here. + # No-op for an eager scratch; never fires on the real run. + peer_offsets = list(self._profiling_scratch.peer_offsets) + shared_workspace = self._profiling_scratch.shared_workspace assert peer_offsets is not None, ( "Sm100MegaMoENvfp4Runner.forward requires peer_offsets kwarg " "(length = world_size); single-rank degenerate mode passes " @@ -1133,11 +1467,18 @@ def forward( kernel = self._build_kernel(tactic_t) - # ``local_workspace`` is per-rank private; cached across calls. + # form-B accumulates the top-k reduction into ``combine_output``, + # so its live rows MUST start at 0; the CALLER zeros [:num_tokens] + # (the op has no live token count -- activation is padded to max_T + # -- and a full zero wastes decode time). form-A overwrites. + + # Cached per FULL tactic: local workspace SIZE is tactic-dependent + # (see _tactic_cache_key). local_workspace = _get_or_alloc_local_workspace( kernel, cache_key=self._tactic_cache_key(tactic_t), device=activation.device, + latch_in_tuning_mode=self.tactic_autotune, ) # ``shared_workspace`` is peer-mapped (symmetric heap) for # multi-rank or local CUDA for the single-rank degenerate @@ -1168,92 +1509,149 @@ def forward( # peer rank's in-kernel dispatch barrier write into this rank's # ``nvlink_barrier_signal``: a fast peer ``red_add(+1)``s our slot, # then our late ``zero_()`` wipes it, so the barrier never reaches - # ``world_size`` and the whole grid deadlocks (the EPLB multi-rank - # dispatch-barrier hang). The symmetric workspace's peer-written - # count regions (expert_recv_count[_sum]) are instead reset - # device-side by the kernel's ``tail_reset_shared_counters``, - # ``nvlink_barrier_signal`` self-primes (phase-flip), and - # ``src_token_topk_idx`` is overwritten by dispatch each launch -- - # so the shared workspace needs no per-launch host zero at all. - if self.world_size > 1: - _zero_local_workspace_preserving_phase(local_workspace, kernel) - else: - shared_workspace.zero_() - local_workspace.zero_() - - activation_cute = _to_cute(activation) - activation_sf_cute = _to_cute(activation_sf) - topk_idx_cute = _to_cute(topk_idx) - topk_weights_cute = _to_cute(topk_weights) - # The weights are stored ``(slots, N, K_bytes)`` (K = hidden//2 for - # fc1 / intermediate//2 for fc2, innermost / stride-1). The kernel - # reads them K-major with K innermost; present a ``transpose(1, 2)`` - # VIEW ``(slots, K_bytes, N)`` so K stays stride-1. Do NOT - # ``.contiguous()`` -- materializing would move K off the innermost - # axis (N would become stride-1) and corrupt the GEMM (cosine ~0). - fc1_weight_cute = _to_cute(fc1_weight.transpose(1, 2)) - fc1_weight_sf_cute = _to_cute(fc1_weight_sf) - fc2_weight_cute = _to_cute(fc2_weight.transpose(1, 2)) - fc2_weight_sf_cute = _to_cute(fc2_weight_sf) - # Per-expert fp32 scale tensors are 1-D ``(num_local_slots,)``; - # 4-byte alignment matches the fp32 element size (the kernel - # reads them as a plain fp32 vector, no 16-byte TMA tile). - fc1_alpha_cute = _to_cute(fc1_alpha, assumed_align=4) - fc2_alpha_cute = _to_cute(fc2_alpha, assumed_align=4) - fc1_norm_const_cute = _to_cute(fc1_norm_const, assumed_align=4) - combine_output_cute = _to_cute(combine_output) - local_workspace_cute = _to_cute(local_workspace, force_static_layout=True) - shared_workspace_cute = _to_cute(shared_workspace) - - torch_stream = torch.cuda.current_stream() - stream = cuda.CUstream(torch_stream.cuda_stream) - - # SymBufferHost contract: ``base_addr`` is any local pointer - # inside the symmetric heap; ``offsets[r] = peer_base - - # local_base``. All five regions share the same delta - # because ``MegaMoeSymmMemProvider`` carves them out of one - # symmetric allocation, so peer_rank_ptr_mapper.map(local, - # r, off) maps any region's local pointer to its peer. - sym_buf = SymBufferHost( - base_addr=int(activation.data_ptr()), - offsets=tuple(int(off) for off in peer_offsets), - rank_idx=int(self.local_rank), - num_max_ranks=int(self.world_size), + # ``world_size`` and the whole grid deadlocks. Peer-written count + # regions are reset device-side, the signal self-primes + # (phase-flip) -- no per-launch host zero needed. + # + # Phase coherence: the sense-reversing dispatch barrier pairs a + # PERSISTED local nvlink_barrier_counter with a PERSISTED shared + # nvlink_barrier_signal; on a TACTIC CHANGE (autotune boundary, + # profiling->real) the pair can be phase-DECOUPLED -> the next + # barrier spins forever. So reset both to phase 0 under a + # cross-rank fence ONCE per tactic. The fence (host sync + + # collective barrier) is ILLEGAL under capture (eager warmup + # fences first) and valid only for pure DEP (guard below). + _capturing = torch.cuda.is_current_stream_capturing() + # Fence ONCE per (local_ws, shared_ws, world) key, tracked + # module-globally. Keys are raw pointers, so "already fenced" also + # requires the STORAGES to be alive (ABA; _megamoe_fence_key_live). + _fkey = ( + int(local_workspace.data_ptr()), + int(shared_workspace.data_ptr()), + int(self.world_size), ) - + _fkey_live = _megamoe_fence_key_live(_fkey) + if _capturing: + if not _fkey_live: + # Eager pre-passes fence every pairing a captured forward + # uses; reaching capture unfenced means two runners share a + # workspace side. Fail loud (replay would hang silently). + raise RuntimeError( + "MegaMoE-CuteDSL: CUDA-graph capture reached an " + "unfenced (local, shared) workspace pairing; the " + "barrier-reset fence cannot run during capture." + ) + # A later EAGER fence must never zero this captured shared side. + _MEGAMOE_CAPTURED_SHARED_PTRS.add(_fkey[1]) + global _MEGAMOE_GRAPH_CAPTURE_SEEN + _MEGAMOE_GRAPH_CAPTURE_SEEN = True + if (not _fkey_live) and not _capturing: + import torch.distributed as _dist + + _sig_off = int(kernel._shared_offsets["nvlink_barrier_signal"]) + _sig_n = int(kernel._shared_region_by_name["nvlink_barrier_signal"].nbytes) + _bc_off = int(kernel._local_offsets["nvlink_barrier_counter"]) + _bc_n = int(kernel._local_region_by_name["nvlink_barrier_counter"].nbytes) + _have_dist = _dist.is_available() and _dist.is_initialized() + # The fence barriers the default (WORLD) group -- correct ONLY + # for pure DEP. Under TP x EP / PP the WORLD barrier spans ranks + # that never reach this fence -> deadlock; fail loud instead. + # (Threading the EP subgroup is the general fix.) + if _have_dist and _dist.get_world_size() != int(self.world_size): + # Not an assert: the guard must survive ``python -O``. + raise RuntimeError( + "MegaMoE-CuteDSL barrier-reset fence uses the WORLD process " + "group, valid only for pure DEP (world == EP). Detected WORLD=" + f"{_dist.get_world_size()} != EP={self.world_size} (TP x EP / " + "PP) -- fence must run on the EP subgroup." + ) + # Never re-zero a shared side baked into a captured graph + # (captured pairings cannot re-fence -> silent hang). + if _fkey[1] in _MEGAMOE_CAPTURED_SHARED_PTRS: + raise RuntimeError( + "MegaMoE-CuteDSL: eager barrier-reset fence would " + "zero a shared workspace already used by a captured " + "CUDA graph; the captured pairings cannot re-fence." + ) + torch.cuda.current_stream().synchronize() + if _have_dist: + _dist.barrier() + # Reset BOTH sides to phase 0. The local counter can be + # ALREADY-ADVANCED here (at the profiling->real transition the + # winning tactic's cached local ran under the scratch's shared + # ptr; the real launch's new fkey re-fences it). Zero only the + # tiny counter slice; the multi-GiB bulk stays uninitialized. + local_workspace[_bc_off : _bc_off + _bc_n].zero_() + shared_workspace[_sig_off : _sig_off + _sig_n].zero_() + # The zeros are stream-async and dist.barrier() is NOT + # device-ordered after them: a released peer's in-kernel signal + # WRITE could be CLOBBERED by our still-pending zero (barrier + # spins forever). Sync so the zeros LAND before peers launch. + torch.cuda.current_stream().synchronize() + if _have_dist: + _dist.barrier() + # Both sides were just re-zeroed, so every OTHER fenced pairing + # sharing either side is now phase-stale (the autotune + # real->scratch->real ABA); drop them so they re-fence. + # Invariant: fkey in set => neither side touched under any + # other pairing since its fence. + _megamoe_prune_fenced_keys( + {k for k in _MEGAMOE_FENCED_KEYS if k[0] == _fkey[0] or k[1] == _fkey[1]} + ) + _MEGAMOE_FENCED_KEYS.add(_fkey) + # Record the fenced ALLOCATIONS for recycled-VA (ABA) detection. + _MEGAMOE_FENCED_KEY_WS_REFS[_fkey] = ( + weakref.ref(local_workspace.untyped_storage()), + weakref.ref(shared_workspace.untyped_storage()), + ) + # else: SAME tactic -- the kernel's device-side tail self-reset suffices. + + # cute.compile (JIT) launch. The + # uint8 workspaces MUST be cute.Pointer, NOT cute.Tensor: the + # 32-bit memref shape field overflows once shared_workspace passes + # 2 GiB (the kernel addresses by raw base + Int64 byte offset). + _to_cute, _to_cute_ptr, SymBufferHost, cute, cutlass_utils = _cute_launch_helpers() + + # combine_output (max_T, 1, hidden) reshapes freely to 2D. Weights + # present K stride-1 via a transpose VIEW (DLPack carries the + # strides); do NOT ``.contiguous()``. + output_activation = combine_output.reshape(combine_output.shape[0], self.hidden_size) runtime_kwargs = dict( - activation=activation_cute, - activation_sf=activation_sf_cute, - topk_idx=topk_idx_cute, - topk_weights=topk_weights_cute, - fc1_weight=fc1_weight_cute, - fc1_weight_sf=fc1_weight_sf_cute, - fc2_weight=fc2_weight_cute, - fc2_weight_sf=fc2_weight_sf_cute, - fc1_alpha=fc1_alpha_cute, - fc2_alpha=fc2_alpha_cute, - fc1_norm_const=fc1_norm_const_cute, - combine_output=combine_output_cute, - local_workspace=local_workspace_cute, - shared_workspace=shared_workspace_cute, - peer_rank_ptr_mapper_host=sym_buf, - stream=stream, + activation=_to_cute(activation), + activation_sf=_to_cute(activation_sf), + topk_idx=_to_cute(topk_idx), + topk_weights=_to_cute(topk_weights), + fc1_weight=_to_cute(fc1_weight.transpose(1, 2)), + fc1_weight_sf=_to_cute(fc1_weight_sf), + fc2_weight=_to_cute(fc2_weight.transpose(1, 2)), + fc2_weight_sf=_to_cute(fc2_weight_sf), + fc1_alpha=_to_cute(fc1_alpha, assumed_align=4), + fc2_alpha=_to_cute(fc2_alpha, assumed_align=4), + fc1_norm_const=_to_cute(fc1_norm_const, assumed_align=4), + output_activation=_to_cute(output_activation), + local_workspace=_to_cute_ptr(local_workspace), + shared_workspace=_to_cute_ptr(shared_workspace), + peer_rank_ptr_mapper_host=SymBufferHost( + offsets=tuple(int(off) for off in peer_offsets), + rank_idx=int(self.local_rank), + num_max_ranks=int(self.world_size), + ), + stream=cuda.CUstream(torch.cuda.current_stream().cuda_stream), ) - compiled = self._compile_or_get(tactic_t, kernel, runtime_kwargs) - t_launch_start = time.perf_counter() + cache_key = self._tactic_cache_key(tactic_t) + compiled = self.__class__.kernel_cache.get(cache_key) + if compiled is None: + # ``max_active_clusters`` is a compile-time Constexpr (baked in, + # omitted at launch); cluster_size = cluster_shape M*N. + _cluster = _unpack_tactic(tactic_t)[1] + _max_active_clusters = cutlass_utils.HardwareInfo().get_max_active_clusters( + int(_cluster[0]) * int(_cluster[1]) + ) + compiled = cute.compile( + kernel, max_active_clusters=_max_active_clusters, **runtime_kwargs + ) + self.__class__.kernel_cache[cache_key] = compiled compiled(**runtime_kwargs) - t_launch_ms = (time.perf_counter() - t_launch_start) * 1000 - t_forward_ms = (time.perf_counter() - t_forward_start) * 1000 - logger.debug( - "[MegaMoECuteDsl] forward DONE tactic=" - "(mma_tiler=%s, cluster=%s, load_balance=%r) " - "launch+sync=%.0fms total=%.0fms", - tactic_t[0], - tactic_t[1], - tactic_t[4], - t_launch_ms, - t_forward_ms, - ) return combine_output # ----- torch op --------------------------------------------------------- @@ -1288,24 +1686,32 @@ def cute_dsl_megamoe_nvfp4_blackwell( peer_offsets: List[int], apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, - token_back_by_dispatch: bool = False, - non_ubulk_fc2_store: bool = True, + in_kernel_fc2_reduce: bool = False, + combine_format: str = "bf16", + tactic_autotune: bool = False, + num_tokens: int = -1, ) -> None: """Run the fused MegaMoE CuteDSL NVFP4 kernel. Inputs are pre-staged by the caller (the ``MegaMoECuteDsl`` - backend in ``mega_moe_cute_dsl.py``). The op runs AutoTuner once - per call to pick the best tactic for the current shape and - invokes the runner. + backend in ``mega_moe_cute_dsl.py``). ``tactic_autotune=True`` + (``MEGAMOE_TACTIC_AUTOTUNE=1`` bench opt-in) picks the tactic via + AutoTuner per call; the default runs the deterministic token-bucket + heuristic (``default_megamoe_tactic``). ``shared_workspace`` MUST be a symmetric-heap tensor for ``world_size > 1`` (use :class:`MegaMoeSymmMemProvider`); a local CUDA tensor is acceptable for the single-rank degenerate path. ``combine_output`` is mutated in place; the op does not return it because torch custom_op forbids the return value from - aliasing any mutated input. Form A semantics: ``combine_output`` - keeps its ``(T, num_topk, hidden)`` layout, and the caller is - responsible for the host-side ``.sum(dim=1)``. + aliasing any mutated input. + + ``in_kernel_fc2_reduce`` selects the reduction form: ``False`` + (form-A) runs the deterministic standalone TopkReduce; ``True`` + (form-B) folds top-k into the kernel and is NON-deterministic (float + accumulation order). Both write ``combine_output`` with shape + ``(T, 1, hidden)``. Perf knobs come from the tactic, not op + arguments. """ sm_version = get_sm_version() if sm_version not in (100, 103): @@ -1314,6 +1720,14 @@ def cute_dsl_megamoe_nvfp4_blackwell( f"SM 103 (B300); got SM {sm_version}." ) + # Live-token trim: TopkReduce sizes its grid from THIS tensor's dim0, + # so slicing to the live count skips dead reduce rows (~5-8us/layer at + # decode). num_tokens < 0 or oversized keeps the full bucket; so does + # num_tokens == 0 (a zero-token rank still launches so peers can cross + # the NVLink barrier -- a 0-row slice would build a zero grid). + if num_tokens > 0: + combine_output = combine_output[: min(num_tokens, combine_output.shape[0])] + runner = Sm100MegaMoENvfp4Runner( world_size=world_size, local_rank=local_rank, @@ -1326,8 +1740,9 @@ def cute_dsl_megamoe_nvfp4_blackwell( output_dtype=combine_output.dtype, apply_topk_in_fc1=apply_topk_in_fc1, gate_up_clamp=gate_up_clamp, - token_back_by_dispatch=token_back_by_dispatch, - non_ubulk_fc2_store=non_ubulk_fc2_store, + in_kernel_fc2_reduce=in_kernel_fc2_reduce, + combine_format=combine_format, + tactic_autotune=tactic_autotune, ) inputs = [ activation, @@ -1343,15 +1758,67 @@ def cute_dsl_megamoe_nvfp4_blackwell( fc1_norm_const, combine_output, ] + if not tactic_autotune: + # Opt-OUT (default; serving never opts in): skip the AutoTuner. + # ``choose_one`` is NOT a safe no-op -- in tuning mode a cache miss + # materializes the multi-GiB scratch and MERGE-sweeps ~36 + # candidates. tactic=-1 guarantees the deterministic heuristic + # with no tuning collectives, even inside a global autotune(). + runner( + inputs, + tactic=-1, + peer_offsets=peer_offsets, + shared_workspace=shared_workspace, + ) + return tuner = AutoTuner.get() - _, best_tactic = tuner.choose_one( - "trtllm::cute_dsl_megamoe_nvfp4_blackwell", - [runner], - runner.get_tuning_config(), - inputs, - peer_offsets=peer_offsets, - shared_workspace=shared_workspace, - ) + # Opt-IN: in tuning mode the MERGE lockstep sweep runs (made safe by + # the tactic-change fence in forward); outside it choose_one is a pure + # cache lookup (tuned tactic, or -1 -> deterministic heuristic). + # Profiling must use SYMMETRIC cross-rank buffers (the transient + # scratch); only single-rank may fall back to the staging buffer. + prof_scratch = _ACTIVE_MEGAMOE_PROFILING_SCRATCH if world_size > 1 else None + prof_factory = _ACTIVE_MEGAMOE_PROFILING_SCRATCH_FACTORY if world_size > 1 else None + if ( + prof_scratch is None + and prof_factory is None + and world_size > 1 + and tuner.is_tuning_mode + ): + # Not a fallback: profiling would peer-map non-symmetric tensors + # (cross-rank IMA). The backend hands either a live scratch or a + # DEFERRED factory; reaching here means both were bypassed. + raise RuntimeError( + "MegaMoE-CuteDSL autotune profiling requires the " + "symmetric profiling scratch (or a scratch factory) on " + "multi-rank, but neither is active." + ) + if prof_scratch is not None: + runner._profiling_scratch = prof_scratch + prof_peer_offsets = list(prof_scratch.peer_offsets) + prof_shared_workspace = prof_scratch.shared_workspace + else: + # DEFERRED path: the pre-hook materializes the scratch only if + # real profiling launches; ``forward`` then rebinds these staging + # values from it, and on a cache hit choose_one never launches -- + # so they are never consumed for profiling. + runner._profiling_scratch = None + runner._profiling_scratch_factory = prof_factory + prof_peer_offsets = peer_offsets + prof_shared_workspace = shared_workspace + try: + _, best_tactic = tuner.choose_one( + "trtllm::cute_dsl_megamoe_nvfp4_blackwell", + [runner], + runner.get_tuning_config(), + inputs, + peer_offsets=prof_peer_offsets, + shared_workspace=prof_shared_workspace, + ) + finally: + runner._profiling_scratch = None + runner._profiling_scratch_factory = None + # Real run always uses the caller's staging buffer + peer_offsets. runner( inputs, tactic=best_tactic, @@ -1385,7 +1852,9 @@ def _( peer_offsets: List[int], apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, - token_back_by_dispatch: bool = False, - non_ubulk_fc2_store: bool = True, + in_kernel_fc2_reduce: bool = False, + combine_format: str = "bf16", + tactic_autotune: bool = False, + num_tokens: int = -1, ) -> None: return None diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py index a0eef9b7b9ac..30c56a896cf5 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/__init__.py @@ -2,9 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 """CuteDSL MegaMoE NVFP4 kernel package. -Hosts the ported MegaMoE fused dispatch + FC1 + activation + FC2 + combine -CuteDSL kernel (flattened from the upstream ``moe_nvfp4_swapab/`` + ``src/`` -split). The package is loaded +Hosts the MegaMoE fused dispatch + FC1 + activation + FC2 + combine CuteDSL +kernel. The package is loaded lazily by :mod:`tensorrt_llm._torch.modules.fused_moe.mega_moe.mega_moe_cute_dsl` through :func:`import_kernel` so environments without a CUDA 13 Cutlass DSL runtime can still import the backend file for capability probing. @@ -39,7 +38,6 @@ "from_blocked", "import_kernel", "import_sym_buffer_host", - "import_topk_reduce", "stack_byte_reinterpretable_tensors", "to_blocked", ] @@ -72,20 +70,3 @@ def import_sym_buffer_host(): # SymBufferHost lives at module scope as a factory; the upstream API # constructs the per-world-size variant inside sym_buffer.py. return sym_buffer - - -def import_topk_reduce(): - """Lazily import the standalone CuteDSL top-k reduce kernel API. - - Returns ``(compile_topk_reduce, launch_compiled_topk_reduce)`` from - :mod:`.topk_reduce` (mirrors :func:`import_kernel`). The reduce kernel - is only needed by the opt-in transformers graph - (``apply_topk_in_fc1=False``); the deepgemm-default route reduces on - the host via ``combine_output.sum(dim=1)`` and never imports it. Like - ``import_kernel`` this stays lazy so non-SM100 / no-cutlass-dsl - environments can import the backend for capability probing without - pulling the heavyweight CuteDSL symbols. - """ - from .topk_reduce import compile_topk_reduce, launch_compiled_topk_reduce - - return compile_topk_reduce, launch_compiled_topk_reduce diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py index 579231b1a383..c3965a5cc963 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/contract.py @@ -1,5 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause """Codegen-time finite mapping contracts for RMEM tensor handoff.""" from __future__ import annotations @@ -78,9 +80,8 @@ def linearize(self, coord: Sequence[int]) -> int: """Convert a coordinate tuple into a CuTe-style linear index.""" coord_tuple = tuple(coord) if len(coord_tuple) != self.rank: - raise ContractError( - f"Coordinate rank mismatch for {self.names!r}: {len(coord_tuple)} != {self.rank}" - ) + raise ContractError(f"Coordinate rank mismatch for {self.names!r}: " + f"{len(coord_tuple)} != {self.rank}") linear = 0 stride = 1 @@ -162,8 +163,8 @@ def normalize(self, *, domain: Space, codomain: Space) -> tuple[int, ...]: """Validate and return the canonical table for the given spaces.""" if len(self.table) != domain.size: raise ContractError( - f"TableMapping length must equal domain size {domain.size}, got {len(self.table)}" - ) + f"TableMapping length must equal domain size {domain.size}, " + f"got {len(self.table)}") for idx, value in enumerate(self.table): if value < 0 or value >= codomain.size: raise ContractError( @@ -395,3 +396,27 @@ class TensorWithContract: tensor: Any contract: Contract + + +def eval_function_mapping(contract: Contract, **domain_coord) -> dict: + """Evaluate a FunctionMapping contract at runtime.""" + if not isinstance(contract.mapping, FunctionMapping): + raise TypeError("runtime contract eval requires a FunctionMapping") + + result = contract.mapping.function(**domain_coord) + if isinstance(result, dict): + return result + if isinstance(result, (tuple, list)): + if len(result) != contract.codomain.rank: + raise ValueError( + "FunctionMapping result rank does not match codomain rank: " + f"{len(result)} vs {contract.codomain.rank}") + return { + name: result[i] + for i, name in enumerate(contract.codomain.names) + } + if contract.codomain.rank == 1: + return {contract.codomain.names[0]: result} + raise TypeError( + "FunctionMapping runtime eval must return dict/tuple/list, or scalar " + "for rank-1 codomain") diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py index 0a7acf16ee06..cdf832245946 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/custom_ext.py @@ -1,5 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause """Sched extension for fused fc1+fc2 work-tile enrichment and GMEM slicing.""" from typing import List, Optional, Tuple, Union @@ -37,10 +39,9 @@ def __init__( cumulative_data_physical_row: Int32, cumulative_sf_physical_row: Int32, cumulative_token_block_count: Int32, - valid_tokens_in_tile: Int32, + valid_tokens_in_cta_tile: Int32, phase_and_peek: Int32, ): - # Slot 3 reuses base k_tile_cnt storage. super().__init__( expert_idx, tile_m_idx, @@ -50,7 +51,7 @@ def __init__( self.cumulative_data_physical_row = self.k_tile_cnt self.cumulative_sf_physical_row = cumulative_sf_physical_row self.cumulative_token_block_count = cumulative_token_block_count - self.valid_tokens_in_tile = valid_tokens_in_tile + self.valid_tokens_in_cta_tile = valid_tokens_in_cta_tile # Slot 7 is the packed (BlockPhase | (peek_ready << 16)) field. # The ``.phase`` and ``.peek_ready`` properties below unpack it; # consumers call them directly so the codebase reads as if the @@ -81,14 +82,14 @@ def __extract_mlir_values__(self) -> List[ir.Value]: values = super().__extract_mlir_values__() values.extend(extract_mlir_values(self.cumulative_sf_physical_row)) values.extend(extract_mlir_values(self.cumulative_token_block_count)) - values.extend(extract_mlir_values(self.valid_tokens_in_tile)) + values.extend(extract_mlir_values(self.valid_tokens_in_cta_tile)) values.extend(extract_mlir_values(self.phase_and_peek)) return values def __new_from_mlir_values__( self, values: List[ir.Value]) -> "SwapABSwigluFp4Fc12WorkTileInfo": assert len(values) == 8 - return SwapABSwigluFp4Fc12WorkTileInfo( + return type(self)( expert_idx=new_from_mlir_values(self.expert_idx, [values[0]]), tile_m_idx=new_from_mlir_values(self.tile_m_idx, [values[1]]), tile_n_idx=new_from_mlir_values(self.tile_n_idx, [values[2]]), @@ -98,8 +99,8 @@ def __new_from_mlir_values__( self.cumulative_sf_physical_row, [values[4]]), cumulative_token_block_count=new_from_mlir_values( self.cumulative_token_block_count, [values[5]]), - valid_tokens_in_tile=new_from_mlir_values(self.valid_tokens_in_tile, - [values[6]]), + valid_tokens_in_cta_tile=new_from_mlir_values( + self.valid_tokens_in_cta_tile, [values[6]]), phase_and_peek=new_from_mlir_values(self.phase_and_peek, [values[7]]), ) @@ -112,7 +113,7 @@ def to_rmem(self) -> cute.Tensor: rmem[3] = self.k_tile_cnt # = cumulative_data_physical_row rmem[4] = self.cumulative_sf_physical_row rmem[5] = self.cumulative_token_block_count - rmem[6] = self.valid_tokens_in_tile + rmem[6] = self.valid_tokens_in_cta_tile rmem[7] = self.phase_and_peek return rmem @@ -125,7 +126,7 @@ def from_rmem(cls, rmem: cute.Tensor) -> "SwapABSwigluFp4Fc12WorkTileInfo": cumulative_data_physical_row=rmem[3], # type: ignore[arg-type] cumulative_sf_physical_row=rmem[4], # type: ignore[arg-type] cumulative_token_block_count=rmem[5], # type: ignore[arg-type] - valid_tokens_in_tile=rmem[6], # type: ignore[arg-type] + valid_tokens_in_cta_tile=rmem[6], # type: ignore[arg-type] phase_and_peek=rmem[7], # type: ignore[arg-type] ) @@ -160,7 +161,7 @@ def __init__( # shows enough arrivals. Mirrors ``fc1_done_counter_ptr`` for the # fc1->fc2 link: this side is "fc1 input ready", that side is # "fc1 output done". The threshold per-tile is the tile's - # ``valid_tokens_in_tile`` (dispatch does not pull padding + # ``valid_tokens_in_cta_tile`` (dispatch does not pull padding # tokens), read straight off the base work tile -- no separate # threshold field needed. ``None`` in the lean fc1+fc2 path keeps # ``enrich_work_tile_info`` to its existing fc2-only peek shape and @@ -236,7 +237,7 @@ def enrich_work_tile_info( ``cumulative_token_block_count + tile_n_idx`` against ``self.fc2_spin_threshold`` (work-tile-invariant const). - fc1 tiles peek the dispatch->fc1 ``fc1_ready_counter`` at the - same slot index but with ``valid_tokens_in_tile`` as threshold + same slot index (``tile_n_idx``) but with ``valid_tokens_in_cta_tile`` as threshold (per-tile dynamic). This branch only emits when ``self.fc1_ready_counter_ptr is not None`` (MegaMoE mode). """ @@ -247,8 +248,7 @@ def enrich_work_tile_info( if is_valid: # Same slot index for both phases -- fc1 release-add (dispatch # pull) and fc2 release-add (fc1 epi) target the per-task-tile - # counter slot indexed by ``cumulative_token_block_count + - # tile_n_idx``. + # counter slot indexed by ``cumulative_token_block_count + tile_n_idx``. counter_slot = (base_work.cumulative_token_block_count + base_work.tile_n_idx) is_fc1 = base_work.phase == Int32(int(BlockPhase.Linear1)) @@ -257,14 +257,14 @@ def enrich_work_tile_info( # MegaMoE-only: fc1 phase peek on fc1_ready_counter. Threshold # is dynamic (per-tile valid count) because dispatch does not # pull padding tokens, so the counter's terminal value matches - # the tile's valid_tokens_in_tile (cluster_tile_m for full + # the tile's valid_tokens_in_cta_tile (cluster_tile_m for full # tiles, less for an expert's last partial tile). if cutlass.const_expr(self.fc1_ready_counter_ptr is not None): if is_fc1: counter_ptr = self.fc1_ready_counter_ptr + counter_slot peek_ready = spin_wait( counter_ptr, - lambda v: v >= base_work.valid_tokens_in_tile, + lambda v: v >= base_work.valid_tokens_in_cta_tile, peek_only=True, ) peek_bit = Int32(0) @@ -299,7 +299,7 @@ def enrich_work_tile_info( cumulative_data_physical_row=base_work.cumulative_data_physical_row, cumulative_sf_physical_row=base_work.cumulative_sf_physical_row, cumulative_token_block_count=base_work.cumulative_token_block_count, - valid_tokens_in_tile=base_work.valid_tokens_in_tile, + valid_tokens_in_cta_tile=base_work.valid_tokens_in_cta_tile, phase_and_peek=new_phase_and_peek, ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py index 21d41dd488ca..290d7b861f43 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py @@ -1,14 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause """Autonomous epilogue for the fused fc1+fc2 swap-AB MegaMoE kernel. -Component boundaries use ``TensorWithContract`` to keep per-thread RMEM layout -semantics explicit at the handoff between transpose, SwiGLU, quantize, and fc2 -store components. +Per-thread RMEM tensors flow between the transpose / SwiGLU / quantize / fc2 +store steps as bare ``cute.Tensor`` fragments; their thread distribution is a +fixed physical property of the surrounding atom sequence and is documented in +local comments. A ``Contract`` is only kept where it earns its keep: the fc2 +store-out mapping (``Fc2ProcessPipeline.store_out_mapping``) is evaluated at +runtime to drive which token/hidden cell each issue targets and therefore how +the per-token metadata is fetched. """ import dataclasses -from typing import Callable, List, Optional, Tuple, Type, Union +from typing import Any, Callable, List, Literal, Optional, Tuple, Type, Union import cutlass import cutlass.cute as cute @@ -16,162 +22,508 @@ import cutlass.utils as utils import cutlass.utils.blackwell_helpers as sm100_utils from cutlass._mlir import ir -from cutlass._mlir.dialects import llvm +from cutlass._mlir.dialects import arith, llvm, vector from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass.cute.typing import AddressSpace -from cutlass.cutlass_dsl import dsl_user_op +from cutlass.cutlass_dsl import Float32, Int64, T -from .contract import (Contract, FunctionMapping, Space, TensorWithContract, - assert_contract_equivalent) +from .contract import Contract, FunctionMapping, Space, eval_function_mapping from .fc1_fc2_fuse_sched import BlockPhase +from .flag_batch import GpuReleaseFlagBatchTracker from .iket_compat import iket -from .megamoe_constants import Nvfp4BlockSize +from .megamoe_constants import (Fp8E4M3RcpLimit, Fp32Max, Nvfp4BlockSize, + Nvfp4E2M1RcpLimit) from .moe_persistent_scheduler import (MoESchedConsumer, MoESchedExtension, MoEWorkTileInfo) +from .ptx_helpers import cp_async_bulk_s2g as _cp_async_bulk_s2g +from .ptx_helpers import \ + cp_reduce_async_bulk_add_noftz_bf16_s2g as \ + _cp_reduce_async_bulk_add_noftz_bf16_s2g +from .ptx_helpers import \ + red_add_relaxed_sys_v2_bf16x2 as _red_add_relaxed_sys_v2_bf16x2 from .sym_buffer import SymBufferDeviceBase -from .token_comm import TokenCommArgs +from .token_comm import CombineFormat, TokenCommArgs, TokenSrcMetadata -# ============================================================================= -# Module-local helpers -# ============================================================================= - -@dsl_user_op -def _red_add_relaxed_sys_v2_bf16x2( - addr, - val0_packed_bf16x2, - val1_packed_bf16x2, - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, -) -> None: - """Issue ``red.relaxed.sys.global.add.v2.bf16x2 [addr], {v0, v1};``. - - Used by the fc2 REDG path to atomic-add 4 bf16 cells. Inline asm is - used because cuTeDSL has no vector-form ``red.v2.bf16x2`` surface; the - operands are packed bf16x2 bit patterns carried in 32-bit registers. - """ - llvm.inline_asm( - None, - [ - addr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip), - val0_packed_bf16x2.ir_value(loc=loc, ip=ip), - val1_packed_bf16x2.ir_value(loc=loc, ip=ip), - ], - "red.relaxed.sys.global.add.noftz.v2.bf16x2 [$0], {$1, $2};", - "l,r,r", +@cute.jit +def cvt_f32_to_ue8m0_to_f32(fp32x1, loc=None, ip=None): + """Round-trip fp32 -> ue8m0 -> fp32: the E8M0 decode-scale quantize step + (rp/satfinite downcast, then upcast through bf16).""" + src_fp32 = Float32(fp32x1).ir_value(loc=loc, ip=ip) + + asm_tmpl = ("{\n" + " .reg .b16 bf_lo;\n" + " cvt.rp.satfinite.ue8m0x2.f32 bf_lo, 0f00000000, $1;\n" + " cvt.rn.bf16x2.ue8m0x2 $0, bf_lo;\n" + "}") + packed_i32 = llvm.inline_asm( + T.i32(), + [src_fp32], + asm_tmpl, + "=r,f", has_side_effects=True, is_align_stack=False, asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, ) + vec_bf16_ty = ir.Type.parse("vector<2xbf16>") + bf2_lo = llvm.bitcast(vec_bf16_ty, packed_i32, loc=loc, ip=ip) + h0 = vector.extract(bf2_lo, [], [0], loc=loc, ip=ip) + dst_f32 = arith.extf(Float32.mlir_type, h0, loc=loc, ip=ip) + + return dst_f32 -@dsl_user_op -def _red_add_release_gpu_s32( - counter_ptr, - value, - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, -) -> None: - """Issue ``red.release.gpu.add.s32`` to a GMEM int32 location. - Publishes fc1 task-tile completion after the caller has flushed the fc1 - output stores. Single-thread helper; caller guards the thread predicate. +@dataclasses.dataclass(frozen=True) +class QuantImpl: + """Register-level block quantizer shared by the fc1 / fc2 epilogues. + + Returns ``(data_regs, sf_regs)`` ONLY: the caller pre-multiplies the topk + weight / global scale into ``prequant_reg`` beforehand and owns the data / + sf plane stores afterwards. ``sf_vec_direction`` selects the per-block amax + reduction: + + * ``regs_in_thread`` -- a block is ``sf_vec`` contiguous regs of + one thread; amax is thread-local (packed bf16x2 abs-max for combine, + fp32 ``fmax`` for orthodox). + * ``threads_with_the_same_reg`` -- a block is one reg across ``sf_vec`` warp + lanes; amax is a warp CREDUX (full warp for vec=32, lane-predicated + halves for vec=16) and is fp32-only, so bf16 is upconverted first. + + ``prequant_reg`` must be 1D with size divisible by ``sf_vec``. Combine inputs + are bf16 (fc2's bf16 reorder regs); orthodox nvfp4 input is fp32 (swiglu). """ - llvm.inline_asm( - None, - [ - counter_ptr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip), - value.ir_value(loc=loc, ip=ip), - ], - "red.release.gpu.global.add.s32 [$0], $1;", - "l,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) + quant_kind: Union[str, CombineFormat] + sf_vec_direction: Literal["regs_in_thread", "threads_with_the_same_reg"] + lane_idx: Optional[Any] = None # Int32; only the across-lane path needs it -@dsl_user_op -def _cp_async_bulk_s2g( - dst_gmem, - src_smem, - size_bytes, - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, -) -> None: - """Issue non-tensor descriptor-free ``cp.async.bulk`` SMEM->GMEM. - - cuTeDSL does expose ``cpasync.CopyBulkS2GOp`` / ``cute.copy`` for this - instruction family, but that abstraction bakes the transfer size into - the copy atom / static tensor layout: CuteNvGPU lowers it as an - ``arch.copy.SM90.bulk_copy_s2g`` op whose ``size`` is an ``I32Attr``. - The fc2 UBLK epilogue needs a runtime byte count for the hidden-tail - row (still 16B-aligned, but not necessarily the full 128-hidden row). - Using the cute copy atom would silently encode the wrong semantic - contract, so keep the raw PTX here until the dialect grows a dynamic-size - descriptor-free bulk-copy op. - - This helper only issues the instruction. The caller owns - ``cp_async_bulk_commit_group`` so copy and reduce bulk paths share the - same group boundary. - """ - # with cute.arch.elect_one(): - llvm.inline_asm( - None, - [ - dst_gmem.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip), - src_smem.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip), - size_bytes.ir_value(loc=loc, ip=ip), - ], - "cp.async.bulk.global.shared::cta.bulk_group [$0], [$1], $2;", - "l,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) + # -- config / validation -------------------------------------------------- + + def __post_init__(self): + if isinstance(self.quant_kind, CombineFormat): + if not self.quant_kind.is_quantized: + raise ValueError( + f"QuantImpl combine path needs a quantized CombineFormat, " + f"got {self.quant_kind}.") + elif self.quant_kind != "nvfp4": + raise ValueError(f"quant_kind must be a CombineFormat or 'nvfp4', " + f"got {self.quant_kind!r}.") + if self.sf_vec_direction not in ("regs_in_thread", + "threads_with_the_same_reg"): + raise ValueError( + f"sf_vec_direction must be 'regs_in_thread' or " + f"'threads_with_the_same_reg', got {self.sf_vec_direction!r}.") + # Orthodox nvfp4 only sees the transposed (regs-in-thread) layout fc1 + # produces via its TMEM transpose. + if not isinstance(self.quant_kind, CombineFormat) and ( + self.sf_vec_direction != "regs_in_thread"): + raise NotImplementedError( + "orthodox nvfp4 quant is regs_in_thread only.") + if self.sf_vec_direction == "threads_with_the_same_reg" and ( + self.lane_idx is None): + raise ValueError( + "across-lane quant needs lane_idx for the CREDUX half-warp predicate." + ) + @property + def data_dtype(self): + if isinstance(self.quant_kind, CombineFormat): + return self.quant_kind.act_dtype + return cutlass.Float4E2M1FN # orthodox nvfp4 + + @property + def scale_dtype(self): + if isinstance(self.quant_kind, CombineFormat): + return self.quant_kind.scale_dtype + return cutlass.Float8E4M3FN # orthodox nvfp4 e4m3 sfc + + @property + def sf_vec_size(self) -> int: + if isinstance(self.quant_kind, CombineFormat): + return self.quant_kind.scale_block + return Nvfp4BlockSize # 16 + + @property + def _is_combine(self) -> bool: + return isinstance(self.quant_kind, CombineFormat) + + @property + def _data_rcp_limit(self) -> float: + # 1 / max representable magnitude of the data element type. Only e2m1 + # and e4m3 data planes exist (orthodox nvfp4 + CombineFormat acts). + dt = self.data_dtype + if dt is cutlass.Float4E2M1FN: + return Nvfp4E2M1RcpLimit # 1/6 + assert dt is cutlass.Float8E4M3FN, f"unexpected data dtype {dt}" + return Fp8E4M3RcpLimit # 1/448 + + # -- dispatch ------------------------------------------------------------- -@dsl_user_op -def _cp_reduce_async_bulk_add_noftz_bf16_s2g( - dst_gmem, - src_smem, - size_bytes, - *, - loc: Optional[ir.Location] = None, - ip: Optional[ir.InsertionPoint] = None, -) -> None: - """Issue non-tensor ``cp.reduce.async.bulk`` for BF16 add. - - cuTeDSL currently exposes descriptor-free ``CopyBulkS2GOp`` but not the - matching descriptor-free reduce atom. Keep the fallback local to this - epilogue path so the rest of the bulk pipeline can still share the same - tensor/layout front-end. - """ - # with cute.arch.elect_one(): - llvm.inline_asm( - None, - [ - dst_gmem.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip), - src_smem.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip), - size_bytes.ir_value(loc=loc, ip=ip), - ], - "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.noftz.bf16 [$0], [$1], $2;", - "l,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - loc=loc, - ip=ip, - ) + @cute.jit + def __call__(self, prequant_reg: cute.Tensor, *, norm_const=None): + if cutlass.const_expr(cute.size(prequant_reg) % self.sf_vec_size != 0): + raise ValueError("prequant_reg size must be divisible by sf_vec.") + # Combine quantizes fc2's bf16 reorder regs; orthodox nvfp4 the fp32 swiglu. + expected_in = cutlass.BFloat16 if self._is_combine else cutlass.Float32 + if cutlass.const_expr(prequant_reg.element_type is not expected_in): + raise TypeError( + f"QuantImpl({self.quant_kind}) expects {expected_in} prequant " + f"input, got {prequant_reg.element_type}.") + if cutlass.const_expr(not self._is_combine): + return self.nvfp4_quant_impl(prequant_reg, norm_const=norm_const) + if cutlass.const_expr(self.data_dtype is cutlass.Float4E2M1FN): + if cutlass.const_expr(self.sf_vec_direction == "regs_in_thread"): + return self.nvfp4_combine_quant_regs_in_thread_impl( + prequant_reg) + return self.nvfp4_combine_quant_threads_with_the_same_reg_impl( + prequant_reg) + if cutlass.const_expr(self.sf_vec_direction == "regs_in_thread"): + return self.mxfp8_combine_quant_regs_in_thread_impl(prequant_reg) + return self.mxfp8_combine_quant_threads_with_the_same_reg_impl( + prequant_reg) + + # -- impls ---------------------------------------------------------------- + + # regs_in_thread only; fc1 promises the vec direction via its TMEM transpose. + @cute.jit + def nvfp4_quant_impl( + self, + prequant_reg: cute.Tensor, + *, + norm_const: Optional[cutlass.Float32] = None, + ) -> Tuple[cute.Tensor, cute.Tensor]: + # fp32 in -> e2m1 data + e4m3 sfc. Mirrors the prior nvfp4_quant scale + # math (sfc -> capped/masked acc_scale); topk pre-mult + sf store are the + # caller's job now. + vec = self.sf_vec_size + n_blocks = cute.size(prequant_reg) // vec + data = cute.make_rmem_tensor((cute.size(prequant_reg), ), + cutlass.Float4E2M1FN) + sf = cute.make_rmem_tensor((n_blocks, ), cutlass.Float8E4M3FN) + in_blocks = cute.zipped_divide(prequant_reg, + (vec, )) # ((vec,), (n_blocks,)) + data_blocks = [] + sf_values = [] + rcp_limit = cutlass.Float32(self._data_rcp_limit) + for vec_block_idx in cutlass.range_constexpr(n_blocks): + block = in_blocks[None, vec_block_idx] + amax = self._amax_thread_fp32(block) + if cutlass.const_expr(norm_const is not None): + sfc_fp32 = amax * rcp_limit * norm_const + else: + sfc_fp32 = amax * rcp_limit + sfc_e4m3 = sfc_fp32.to(cutlass.Float8E4M3FN) + sfc_rt = cutlass.Float32(sfc_e4m3) + if cutlass.const_expr(norm_const is not None): + acc_scale = norm_const * cute.arch.rcp_approx(sfc_rt) + else: + acc_scale = cute.arch.rcp_approx(sfc_rt) + acc_scale = cute.arch.fmin(acc_scale, Fp32Max) + mask = cute.arch.fmin(sfc_rt * cutlass.Float32(1e30), + cutlass.Float32(1.0)) + acc_scale = acc_scale * mask + sf_values.append(sfc_e4m3) + data_blocks.append(self._scale_to_e2m1_ssa(block, acc_scale)) + self._store_packed_blocks(data, data_blocks) + sf.store(self._values_to_ssa(sf_values, cutlass.Float8E4M3FN)) + return data, sf + + @cute.jit + def nvfp4_combine_quant_regs_in_thread_impl( + self, + prequant_reg: cute.Tensor, + ) -> Tuple[cute.Tensor, cute.Tensor]: + # bf16 in -> e2m1 data + per-16 bf16 amax. amax found on bf16 (packed). + vec = self.sf_vec_size + n_blocks = cute.size(prequant_reg) // vec + data = cute.make_rmem_tensor((cute.size(prequant_reg), ), + cutlass.Float4E2M1FN) + sf = cute.make_rmem_tensor((n_blocks, ), cutlass.BFloat16) + in_blocks = cute.zipped_divide(prequant_reg, + (vec, )) # ((vec,), (n_blocks,)) + data_blocks = [] + sf_values = [] + for vec_block_idx in cutlass.range_constexpr(n_blocks): + block = in_blocks[None, vec_block_idx] + amax = self._amax_thread_bf16(block) + sf_values.append(amax) + decode_scale = cutlass.Float32(amax) * cutlass.Float32( + self._data_rcp_limit) + data_blocks.append( + self._scale_to_e2m1_ssa(block, self._enc_nvfp4(decode_scale))) + self._store_packed_blocks(data, data_blocks) + sf.store(self._values_to_ssa(sf_values, cutlass.BFloat16)) + return data, sf + + # Mapping: (lane_idx, selected_sf_idx) -> (token_64, hidden_32) + # token_idx = lane_idx % 16 + selected_sf_idx * 16 + # hidden_idx = lane_idx // 16 * 16 + @cute.jit + def nvfp4_combine_quant_threads_with_the_same_reg_impl( + self, + prequant_reg: cute.Tensor, + ) -> Tuple[cute.Tensor, cute.Tensor]: + # UBLK has lane == hidden, so the warp's 32 lanes are 32 consecutive + # hidden. A scale block = sf_vec hidden, so the lanes split along hidden + # into 32 // sf_vec blocks of sf_vec lanes each (warp = blocks_per_warp * + # lanes_per_block, the EP x TP split). Only the sf_vec lanes inside a + # block share its CREDUX scale, so they pool the subtile tokens: sf_vec + # == 32 pools the whole warp, sf_vec < 32 pools fewer (more per lane). + lanes_per_block = self.sf_vec_size + n_tokens = cute.size(prequant_reg) + lane_in_block = self.lane_idx % cutlass.Int32(lanes_per_block) + data = cute.make_rmem_tensor((n_tokens, ), cutlass.Float4E2M1FN) + selected_sf = cute.make_rmem_tensor((n_tokens // lanes_per_block, ), + cutlass.BFloat16) + scaled_vec = cute.full((n_tokens, ), cutlass.Float32(0.0), + cutlass.Float32) + for token_idx in cutlass.range_constexpr(n_tokens): + value = cutlass.Float32(prequant_reg[token_idx]) + amax_bf16 = self._amax_lane(value).to(cutlass.BFloat16) + slot = token_idx // lanes_per_block + if (token_idx % lanes_per_block) == lane_in_block: + selected_sf[slot] = amax_bf16 + else: + selected_sf[slot] = selected_sf[slot] + decode_scale = cutlass.Float32(amax_bf16) * cutlass.Float32( + self._data_rcp_limit) + scaled_value = value * self._enc_nvfp4(decode_scale) + scaled_vec = cute.TensorSSA( + vector.insert( + scaled_value.ir_value(), + scaled_vec.ir_value(), + [], + [token_idx], + ), + (n_tokens, ), + cutlass.Float32, + ) + self._store_packed_data(data, scaled_vec.to(cutlass.Float4E2M1FN)) + return data, selected_sf + + @cute.jit + def mxfp8_combine_quant_regs_in_thread_impl( + self, + prequant_reg: cute.Tensor, + ) -> Tuple[cute.Tensor, cute.Tensor]: + # bf16 in -> e4m3 data + per-32 e8m0. amax found on bf16 (packed). + vec = self.sf_vec_size + n_blocks = cute.size(prequant_reg) // vec + data = cute.make_rmem_tensor((cute.size(prequant_reg), ), + cutlass.Float8E4M3FN) + sf = cute.make_rmem_tensor((n_blocks, ), cutlass.Float8E8M0FNU) + in_blocks = cute.zipped_divide(prequant_reg, + (vec, )) # ((vec,), (n_blocks,)) + data_blocks = [] + sf_values = [] + for vec_block_idx in cutlass.range_constexpr(n_blocks): + block = in_blocks[None, vec_block_idx] + # widen the native-bf16 amax to fp32 for the e8m0 round-up math. + scale_e8m0, scale_f32 = self._e8m0( + cutlass.Float32(self._amax_thread_bf16(block))) + sf_values.append(scale_e8m0) + data_blocks.append( + self._scale_to_e4m3_ssa(block, self._enc_mxfp8(scale_f32))) + self._store_packed_blocks(data, data_blocks) + sf.store(self._values_to_ssa(sf_values, cutlass.Float8E8M0FNU)) + return data, sf + + # Mapping: (lane_idx, selected_sf_idx) -> (token_64, hidden_32) + # token_idx = lane_idx + selected_sf_idx * 32 + # hidden_idx = 0 + @cute.jit + def mxfp8_combine_quant_threads_with_the_same_reg_impl( + self, + prequant_reg: cute.Tensor, + ) -> Tuple[cute.Tensor, cute.Tensor]: + # UBLK has lane == hidden, so the warp's 32 lanes are 32 consecutive + # hidden. A scale block = sf_vec hidden, so the lanes split along hidden + # into 32 // sf_vec blocks of sf_vec lanes each (warp = blocks_per_warp * + # lanes_per_block, the EP x TP split). Only the sf_vec lanes inside a + # block share its CREDUX scale, so they pool the subtile tokens. mxfp8 + # sf_vec == 32 -> the whole warp is one block, all 32 lanes pool. + lanes_per_block = self.sf_vec_size + n_tokens = cute.size(prequant_reg) + lane_in_block = self.lane_idx % cutlass.Int32(lanes_per_block) + data = cute.make_rmem_tensor((n_tokens, ), cutlass.Float8E4M3FN) + selected_sf = cute.make_rmem_tensor((n_tokens // lanes_per_block, ), + cutlass.Float8E8M0FNU) + scaled_vec = cute.full((n_tokens, ), cutlass.Float32(0.0), + cutlass.Float32) + for token_idx in cutlass.range_constexpr(n_tokens): + value = cutlass.Float32(prequant_reg[token_idx]) + scale_e8m0, scale_f32 = self._e8m0(self._amax_lane(value)) + slot = token_idx // lanes_per_block + if (token_idx % lanes_per_block) == lane_in_block: + selected_sf[slot] = scale_e8m0 + else: + selected_sf[slot] = selected_sf[slot] + scaled_value = value * self._enc_mxfp8(scale_f32) + scaled_vec = cute.TensorSSA( + vector.insert( + scaled_value.ir_value(), + scaled_vec.ir_value(), + [], + [token_idx], + ), + (n_tokens, ), + cutlass.Float32, + ) + self._store_packed_data(data, scaled_vec.to(cutlass.Float8E4M3FN)) + return data, selected_sf + + # -- shared sub-steps ----------------------------------------------------- + + @cute.jit + def _scale_to_e2m1_ssa( + self, + block: cute.Tensor, + enc: cutlass.Float32, + ) -> cute.TensorSSA: + block_f32 = block.load().to(cutlass.Float32) + enc_vec = cute.full_like(block_f32, enc, cutlass.Float32) + return (block_f32 * enc_vec).to(cutlass.Float4E2M1FN) + + @cute.jit + def _scale_to_e4m3_ssa( + self, + block: cute.Tensor, + enc: cutlass.Float32, + ) -> cute.TensorSSA: + block_f32 = block.load().to(cutlass.Float32) + enc_vec = cute.full_like(block_f32, enc, cutlass.Float32) + return (block_f32 * enc_vec).to(cutlass.Float8E4M3FN) + + @cute.jit + def _values_to_ssa(self, values, + dtype: Type[cutlass.Numeric]) -> cute.TensorSSA: + vec = vector.from_elements( + T.vector(len(values), dtype.mlir_type), + [values[i].ir_value() for i in range(len(values))], + ) + return cute.TensorSSA(vec, (len(values), ), dtype) + + @cute.jit + def _concat_i32_blocks_ssa(self, blocks) -> cute.TensorSSA: + values = [] + for block_idx in cutlass.range_constexpr(len(blocks)): + packed_block = blocks[block_idx].bitcast(cutlass.Int32) + for elem_idx in cutlass.range_constexpr( + cute.size(packed_block.shape)): + values.append(packed_block[elem_idx].ir_value()) + vec = vector.from_elements( + T.vector(len(values), cutlass.Int32.mlir_type), + values, + ) + return cute.TensorSSA(vec, (len(values), ), cutlass.Int32) + + @cute.jit + def _store_packed_blocks(self, data: cute.Tensor, blocks) -> None: + packed_data = cute.recast_tensor(data, cutlass.Int32) + packed_data.store(self._concat_i32_blocks_ssa(blocks)) + + @cute.jit + def _store_packed_data(self, data: cute.Tensor, + data_ssa: cute.TensorSSA) -> None: + packed_data = cute.recast_tensor(data, cutlass.Int32) + packed_data.store(data_ssa.bitcast(cutlass.Int32)) + + @cute.jit + def _amax_thread_fp32(self, block: cute.Tensor) -> cutlass.Float32: + # max.xorsign.abs reduces |.| in one op per element; the result sign is + # the xor of the inputs (junk for an amax), so clear it at the end. + def max_abs(lhs: cutlass.Float32, + rhs: cutlass.Float32) -> cutlass.Float32: + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [ + cutlass.Float32(lhs).ir_value(), + cutlass.Float32(rhs).ir_value() + ], + "max.xorsign.abs.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + )) + + acc = block[0] + for elem_idx in cutlass.range_constexpr(1, cute.size(block)): + acc = max_abs(acc, block[elem_idx]) + mag_bits = cutlass.Int32( + llvm.bitcast( + T.i32(), + cutlass.Float32(acc).ir_value())) & cutlass.Int32(0x7FFFFFFF) + return cutlass.Float32(llvm.bitcast(T.f32(), mag_bits.ir_value())) + + @cute.jit + def _amax_thread_bf16(self, block: cute.Tensor) -> cutlass.BFloat16: + # Packed bf16x2 abs-max: tree-reduce the pairs, then fold the survivor's + # two halves (high shifted into low). max.xorsign.abs leaves a junk sign, + # so the low bf16 is masked before being read back. The amax is natively + # bf16 -- exactly what the wire format stores. + def max_abs(lhs: cutlass.Int32, rhs: cutlass.Int32) -> cutlass.Int32: + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Int32(lhs).ir_value(), + cutlass.Int32(rhs).ir_value() + ], + "max.xorsign.abs.bf16x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + )) + + pairs = cute.recast_tensor(block, cutlass.Int32) # (vec/2,) bf16x2 + acc = cutlass.Int32(pairs[0]) + for pair_idx in cutlass.range_constexpr(1, cute.size(pairs)): + acc = max_abs(acc, pairs[pair_idx]) + acc = max_abs(acc, acc >> cutlass.Int32(16)) + amax_bits = cute.make_rmem_tensor((1, ), cutlass.Int32) + amax_bits[0] = acc & cutlass.Int32(0x7FFF) + return cute.recast_tensor(amax_bits, cutlass.BFloat16)[0] + + @cute.jit + def _amax_lane(self, v: cutlass.Float32) -> cutlass.Float32: + if cutlass.const_expr(self.sf_vec_size == 32): + return cute.arch.warp_redux_sync(v, "fmax", abs=True) + first_half = (self.lane_idx % cutlass.Int32(32)) < cutlass.Int32(16) + vsel = cutlass.Float32(0.0) + if first_half: + vsel = v + amax = cute.arch.warp_redux_sync(vsel, "fmax", abs=True) + if not first_half: + amax = cute.arch.warp_redux_sync(v, "fmax", abs=True) + return amax + + @cute.jit + def _e8m0( + self, amax: cutlass.Float32 + ) -> Tuple[cutlass.Float8E8M0FNU, cutlass.Float32]: + candidate = amax * cutlass.Float32(self._data_rcp_limit) + scale_f32 = cutlass.Float32(cvt_f32_to_ue8m0_to_f32(candidate)) + return scale_f32.to(cutlass.Float8E8M0FNU), scale_f32 + + @cute.jit + def _enc_nvfp4(self, decode_scale: cutlass.Float32) -> cutlass.Float32: + # rcp.approx.ftz with the fc1 cap+mask idiom (amax==0 -> 0, no inf*0 NaN). + enc = cute.arch.fmin(cute.arch.rcp_approx(decode_scale), Fp32Max) + mask = cute.arch.fmin(decode_scale * cutlass.Float32(1e30), + cutlass.Float32(1.0)) + return enc * mask + + @cute.jit + def _enc_mxfp8(self, scale_f32: cutlass.Float32) -> cutlass.Float32: + # Skip nan + enc = cute.arch.fmin(cute.arch.rcp_approx(scale_f32), Fp32Max) + mask = cute.arch.fmin(scale_f32 * cutlass.Float32(1e30), + cutlass.Float32(1.0)) + return enc * mask # ============================================================================= @@ -192,43 +544,21 @@ class Region: class _TmemTranspose16x32Core: - """Contract-naive physical implementation of the 16x32 -> 32x16 TMEM - in-place transpose. Shared by: - - - ``TmemTranspose16x32`` : fc1 epi codomain naming - (``intermediate_output_idx``); - elements are fp32 (swiglu fold output). - - ``TmemTranspose16x32Packed`` : fc2 epi codomain naming - (``hidden_pair_idx``); elements are - 32-bit packed ``(bf16, bf16)`` pairs. + """Physical implementation of the 16x32 -> 32x16 TMEM in-place transpose. - The (lane_idx, elem_idx) physical distribution is identical for both - subclasses -- the underlying tcgen05 atoms are 32-bit element atoms, - agnostic to whether each 32-bit slot holds an fp32 or a packed bf16x2. - Only the codomain semantic names differ, expressed via the subclass's - ``InputContract`` / ``OutputContract`` class attributes. + The transpose is a fixed sequence of tcgen05 32-bit element atoms; each + 32-bit slot is opaque to it (an fp32 swiglu-fold value for fc1, or a packed + ``(bf16, bf16)`` pair for fc2 -- the physical (lane_idx, elem_idx) + distribution is identical either way). The (thread, reg) -> (tmem_dp, + tmem_col) input / output mapping is documented on the ``TmemTranspose16x32`` + subclass, which is the public entry point. - Per-thread RMEM coordinate convention (used by both subclasses' contracts): + Per-thread RMEM coordinate convention: - ``lane_idx`` -- warp lane id (= thread index within warp), in [0, 32). - ``elem_idx`` -- per-thread reg index, in [0, 16). - - Subclasses MUST override these two class attributes: - ``InputContract`` -- (lane_idx, elem_idx) -> codomain mapping after - R1.Load (or after ``reg_tensor`` is fed in for - skip-R1.Load mode). - ``OutputContract`` -- (lane_idx, elem_idx) -> codomain mapping after - ``r4_perm`` has run all four rounds. - - The Core's ``__init__`` reads ``self.InputContract`` / ``self.OutputContract`` - via Python's normal MRO attribute lookup; the subclass's overrides take - precedence at construction time. """ - # Subclasses MUST override these. - InputContract: Contract - OutputContract: Contract - _PermR1 = (0, 8, 2, 10, 4, 12, 6, 14, 1, 9, 3, 11, 5, 13, 7, 15) _PermR3 = (0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15) _PermR4 = (0, 8, 2, 10, 4, 12, 6, 14, 1, 9, 3, 11, 5, 13, 7, 15) @@ -270,24 +600,19 @@ def load_subtile_raw_acc( ``warp_lane_offset + acc_stage_col_offset + subtile_col_offset``; see ``SwapABSwigluFp4Epilogue._subtile_local_tmem_tensor``). - Returns a 4-tuple of (16,) fp32 RMEM tensors, each carrying - the (lane_idx, elem_idx) -> codomain distribution described by - ``TmemTranspose16x32.InputContract`` / - ``TmemTranspose16x32Packed.InputContract`` (physically identical - for fc1 and fc2, only codomain semantic names differ): + Returns a 4-tuple of (16,) fp32 RMEM tensors, each carrying the + (lane_idx, elem_idx) input distribution documented on + ``TmemTranspose16x32`` (physically identical for fc1 and fc2): [0] gate_lo / first-half top -- subtile cols 0..31, lanes 0..15 [1] up_lo / first-half bot -- subtile cols 0..31, lanes 16..31 [2] raw_top / second-half top -- subtile cols 32..63, lanes 0..15 [3] raw_bot / second-half bot -- subtile cols 32..63, lanes 16..31 - 4 atom calls of ``Ld16x64bOp(Repetition.x16) Float32`` -- the - same atom currently used by the per-subtile entry LDTM in - ``_run_fc1_subtile`` and by ``second_t.r1_load`` / - ``Fc2AccLoadAndPack`` per-half LDTMs. Caller is expected to - wrap each output in ``TensorWithContract`` with - ``TmemTranspose16x32{,Packed}.InputContract`` before handing - them downstream. + 4 atom calls of ``Ld16x64bOp(Repetition.x16) Float32`` -- the same + atom used by the per-subtile entry LDTM. Each output is in the + ``TmemTranspose16x32`` input distribution and can be fed straight + into a transpose as ``reg_tensor`` (skip-R1.Load mode). """ atom_ld16x64 = cute.make_copy_atom( tcgen05.Ld16x64bOp(tcgen05.Repetition.x16), @@ -358,8 +683,18 @@ def __init__( self, tmem_ptr, region: int, - reg_tensor: Optional[TensorWithContract] = None, + reg_tensor: Optional[cute.Tensor] = None, ) -> None: + # The whole transpose is built from 32-bit element atoms; _io_dtype + # drives _src_regs / output / every LDTM/STTM atom below, so guard the + # invariant once here (tautological today, defensive against future + # dtype edits). + if cutlass.const_expr(self._io_dtype.width != 32): + raise TypeError( + f"{type(self).__name__} requires a 32-bit _io_dtype (the " + f"transpose uses 32-bit element atoms), got {self._io_dtype} " + f"(width {self._io_dtype.width}).") + half_lane_off = 16 * self._TmemRowStride if region == Region.Top: src_ptr = tmem_ptr @@ -403,21 +738,31 @@ def __init__( ) self._src_regs = cute.make_rmem_tensor((16, ), self._io_dtype) - output_tensor = cute.make_rmem_tensor((16, ), self._io_dtype) - self.output = TensorWithContract( - tensor=output_tensor, - contract=self.OutputContract, - ) - + # ``output`` is a bare (16,) RMEM fragment; its (lane_idx, elem_idx) + # distribution after all four rounds is the transpose output mapping + # documented on ``TmemTranspose16x32``. + self.output = cute.make_rmem_tensor((16, ), self._io_dtype) + + # skip-R1.Load mode: ``reg_tensor`` must already be in the transpose + # input distribution (see ``TmemTranspose16x32`` / produced by + # ``load_subtile_raw_acc``); we copy it in lieu of the R1 LDTM. + # Weak entry guard (replaces the removed input contract): the transpose + # atoms are 32-bit element atoms over exactly 16 regs/lane, so the fed + # tensor must be a 32-bit element type (fp32 or packed bf16x2) of size 16. self._reg_tensor = reg_tensor if reg_tensor is not None: - assert_contract_equivalent( - reg_tensor.contract, - self.InputContract, - context=f"{type(self).__name__} skip-R1.Load reg_tensor", - ) + if cutlass.const_expr(reg_tensor.element_type.width != 32): + raise TypeError( + f"{type(self).__name__} reg_tensor must be a 32-bit element " + f"type (fp32 or packed bf16x2), got element type " + f"{reg_tensor.element_type} (width {reg_tensor.element_type.width})." + ) + if cutlass.const_expr(cute.size(reg_tensor) != 16): + raise ValueError( + f"{type(self).__name__} reg_tensor must hold exactly 16 " + f"elements, got {cute.size(reg_tensor)}.") for r in range(16): - self._src_regs[r] = reg_tensor.tensor[r] + self._src_regs[r] = reg_tensor[r] # -- R1 ------------------------------------------------------------------ @@ -433,12 +778,12 @@ def r1_load(self) -> None: def r1_perm(self) -> None: for r in range(16): - self.output.tensor[r] = self._src_regs[self._PermR1[r]] + self.output[r] = self._src_regs[self._PermR1[r]] def r1_store(self) -> None: cute.copy( self._atom_st16x128, - self._rmem_copy_view(self.output.tensor, 16), + self._rmem_copy_view(self.output, 16), self._tmem_src_full, ) @@ -476,12 +821,12 @@ def r3_load_bot(self) -> None: def r3_perm(self) -> None: for r in range(16): - self.output.tensor[r] = self._src_regs[self._PermR3[r]] + self.output[r] = self._src_regs[self._PermR3[r]] def r3_store(self) -> None: cute.copy( self._atom_st32x32, - self._rmem_copy_view(self.output.tensor, 16), + self._rmem_copy_view(self.output, 16), self._tmem_dst_full, ) @@ -503,12 +848,12 @@ def r4_load_bot(self) -> None: def r4_perm(self) -> None: for r in range(16): - self.output.tensor[r] = self._src_regs[self._PermR4[r]] + self.output[r] = self._src_regs[self._PermR4[r]] def r4_store(self) -> None: cute.copy( self._atom_st32x32, - self._rmem_copy_view(self.output.tensor, 16), + self._rmem_copy_view(self.output, 16), self._tmem_dst_full, ) @@ -528,65 +873,29 @@ def from_r1_perm_until_last_store(self) -> cute.Tensor: class TmemTranspose16x32(_TmemTranspose16x32Core): - """fc1 epi 16x32 -> 32x16 TMEM in-place transpose. - - Contract summary: - - input : ``token_idx = elem_idx * 2 + ((lane_idx // 2) % 2)`` - - output: ``token_idx = lane_idx`` - The second codomain axis is ``intermediate_output_idx``. - """ + """Public 16x32 -> 32x16 TMEM in-place transpose. - _domain = Space(("lane_idx", "elem_idx"), (32, 16)) - _codomain = Space(("token_idx", "intermediate_output_idx"), (32, 16)) + The per-thread RMEM ``(lane_idx, elem_idx) -> (tmem_dp, tmem_col)`` mapping + is fixed by the underlying atom sequence and is identical for fc1 (each slot + is an fp32 swiglu-fold value, ``tmem_col`` = intermediate-output index) and + fc2 (each slot is a packed bf16x2, ``tmem_col`` = hidden-pair index). Only + the ``tmem_col`` semantic name differs between the two uses; the physical + distribution below is the single source of truth. - InputContract = Contract( - domain=_domain, - codomain=_codomain, - mapping=FunctionMapping( - lambda lane_idx, elem_idx: { - "token_idx": elem_idx * 2 + ((lane_idx // 2) % 2), - "intermediate_output_idx": (lane_idx % 2) * 8 + lane_idx // 4, - }), - ) - OutputContract = Contract( - domain=_domain, - codomain=_codomain, - mapping=FunctionMapping(lambda lane_idx, elem_idx: { - "token_idx": lane_idx, - "intermediate_output_idx": elem_idx, - }), - ) + Input distribution -- what each (lane_idx, elem_idx) reg holds on entry + (i.e. straight after the 16-dp x 32-col source LDTM, or as fed in via + ``reg_tensor`` / ``load_subtile_raw_acc`` for skip-R1.Load mode): + tmem_dp = elem_idx * 2 + (lane_idx // 2) % 2 # in [0, 32) + tmem_col = (lane_idx % 2) * 8 + lane_idx // 4 # in [0, 16) -class TmemTranspose16x32Packed(_TmemTranspose16x32Core): - """fc2 epi 16x32 -> 32x16 TMEM in-place transpose, 32-bit packed - bf16x2 elements. + Output distribution -- after all four rounds, the 32-dp x 16-col result has + each lane owning one full dp-row of 16 cols: - Same physical atom sequence as ``TmemTranspose16x32``; codomain is - ``(token_idx, hidden_pair_idx)`` and each slot holds one packed bf16x2. + tmem_dp = lane_idx # in [0, 32) + tmem_col = elem_idx # in [0, 16) """ - _domain = Space(("lane_idx", "elem_idx"), (32, 16)) - _codomain = Space(("token_idx", "hidden_pair_idx"), (32, 16)) - - InputContract = Contract( - domain=_domain, - codomain=_codomain, - mapping=FunctionMapping( - lambda lane_idx, elem_idx: { - "token_idx": elem_idx * 2 + ((lane_idx // 2) % 2), - "hidden_pair_idx": (lane_idx % 2) * 8 + lane_idx // 4, - }), - ) - OutputContract = Contract( - domain=_domain, - codomain=_codomain, - mapping=FunctionMapping(lambda lane_idx, elem_idx: { - "token_idx": lane_idx, - "hidden_pair_idx": elem_idx, - }), - ) - # ============================================================================= # TmemTranspose32x32Inplace @@ -604,8 +913,8 @@ class TmemTranspose32x32Inplace: def __init__( self, tmem_ptr, - reg_tensor_top: Optional[TensorWithContract] = None, - reg_tensor_bot: Optional[TensorWithContract] = None, + reg_tensor_top: Optional[cute.Tensor] = None, + reg_tensor_bot: Optional[cute.Tensor] = None, ) -> None: if (reg_tensor_top is None) != (reg_tensor_bot is None): raise ValueError( @@ -693,7 +1002,8 @@ def __init__( use_2cta_instrs: bool, sf_vec_size: int, fc1_output_dtype: Type[cutlass.Numeric], - fc2_output_dtype: Type[cutlass.Numeric], + combine_format: + CombineFormat, # fc2 combine wire: act_dtype (data) + scale_dtype (sf) non_ubulk_fc2_store: bool, # Whether epilogue warps use STG or UBLK in fc2 in_kernel_fc2_reduce: @@ -702,13 +1012,12 @@ def __init__( bool = False, # Whether epilogue warps store fc2 to local or peer acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, fc1_output_sf_dtype: Type[cutlass.Numeric] = cutlass.Float8E4M3FN, - fc2_output_sf_dtype: Optional[Type[ - cutlass. - Numeric]] = None, # Reserve for later low precision combine allow_overlap_acc: bool = True, static_expert_shape: Optional[Tuple[ int, int, int]] = None, # [expert, intermediate, hidden] gate_up_clamp: Optional[float] = None, # Swiglu style only + epi_flag_batch: Optional[Tuple[int, int]] = ( + 1, 1), # (fc1, fc2) done-counter publish batch ) -> None: if fc1_output_dtype is not cutlass.Float4E2M1FN: raise NotImplementedError( @@ -726,30 +1035,34 @@ def __init__( self.fc2_use_bulk = not non_ubulk_fc2_store self.reduce_topk_in_kernel = in_kernel_fc2_reduce self.token_back_by_dispatch = token_back_by_dispatch - self.fc2_output_dtype = fc2_output_dtype + self.combine_format = combine_format self.fc1_output_dtype = fc1_output_dtype self.acc_dtype = acc_dtype self.fc1_output_sf_dtype = fc1_output_sf_dtype self.sf_vec_size = sf_vec_size # Swiglu gate/up clamp limit; None disables clamping. self.gate_up_clamp = gate_up_clamp - self.cluster_tile_intermediate_downproj = ( - self._EpilogueFc1IntermediateDownTileSize * cluster_shape_mn[0]) + # Done-counter publish batch granularity + _fc1_eb, _fc2_eb = (1, 1) if epi_flag_batch is None else epi_flag_batch + self.fc1_epi_flag_batch = max(1, min(32, int(_fc1_eb))) + self.fc2_epi_flag_batch = max(1, min(32, int(_fc2_eb))) + self.cluster_tile_intermediate_downproj = self._EpilogueFc1IntermediateDownTileSize * cluster_shape_mn[ + 0] atom_thr_size = 2 if use_2cta_instrs else 1 self.cta_tile_m = self._EpilogueFc2HiddenTileSize self.cta_tile_n = mma_tiler_mnk[1] self.cta_tile_k = mma_tiler_mnk[2] - assert mma_tiler_mnk[0] // atom_thr_size == self.cta_tile_m - assert self.cta_tile_n % self._EpilogueTokenTileSize == 0 + assert (mma_tiler_mnk[0] // atom_thr_size == self.cta_tile_m) + assert (self.cta_tile_n % self._EpilogueTokenTileSize == 0) self.static_expert_shape = static_expert_shape self.acc_tmem_cols = self.cta_tile_n self.acc_sf_cols = (max(self.cta_tile_n // 128, 1) * self.cta_tile_k + max(self.cta_tile_m // 128, 1) * self.cta_tile_k) // self.sf_vec_size - if (static_expert_shape is not None and static_expert_shape[2] % - (self.cta_tile_m * cluster_shape_mn[0]) == 0): + if static_expert_shape is not None and static_expert_shape[2] % ( + self.cta_tile_m * cluster_shape_mn[0]) == 0: self.fc2_hidden_needs_predicate: bool = False else: self.fc2_hidden_needs_predicate: bool = True @@ -766,19 +1079,17 @@ def __init__( self.num_acc_stage = 2 self.num_acc_pipeline_stages = 1 if self.overlapping_accum else self.num_acc_stage self.overlapped_tmem_cols = self._EpilogueTokenTileSize if self.overlapping_accum else 0 - assert not self.overlapping_accum or self.overlapped_tmem_cols >= self.acc_sf_cols + assert (not self.overlapping_accum + or self.overlapped_tmem_cols >= self.acc_sf_cols) self.epi_smem_bytes = 8 * 1024 if self.fc1_output_dtype.width > 4: raise NotImplementedError( "Remember to adjust the smem size when switch to mxfp8 support") - self.tmem_acc_layout_py_obj = ( - (self.cta_tile_m, self.cta_tile_n, self.num_acc_stage), - ( - _TmemTranspose16x32Core._TmemRowStride, - 1, - self.cta_tile_n - self.overlapped_tmem_cols, - ), - ) + self.tmem_acc_layout_py_obj = ((self.cta_tile_m, self.cta_tile_n, + self.num_acc_stage), + (_TmemTranspose16x32Core._TmemRowStride, + 1, self.cta_tile_n - + self.overlapped_tmem_cols)) def get_epi_storage_type(self) -> Type: # This could be extended to take atoms space for the larger sf_vec_size quant. @@ -841,17 +1152,10 @@ def run( ), ) - fc1_epi = SwapABFc1Epilogue( - self, - tidx, - epi_smem_storage, - sched_ext, - tma_atom_fc1_output, - fc1_output, - fc1_output_sf, - fc1_done_counter, - optional_epi_args, - ) + fc1_epi = SwapABFc1Epilogue(self, tidx, epi_smem_storage, sched_ext, + tma_atom_fc1_output, fc1_output, + fc1_output_sf, fc1_done_counter, + optional_epi_args) fc2_epi = SwapABFc2Epilogue(self, tidx, epi_smem_storage, fc2_output, token_comm_args, optional_epi_args) @@ -862,8 +1166,15 @@ def run( num_threads=32 * self._EpilogueWarpCnt, ) is_odd_turn = cutlass.Int32(1) - work_tile_info = sched_consumer.consume_work() + + flag_tracker = GpuReleaseFlagBatchTracker( + flag_addr=Int64(0), + cumulated_flags=cutlass.Int32(0), + phase=cutlass.Int32(work_tile_info.phase), + tid=tidx % (self._EpilogueWarpCnt * 32), + ) + while work_tile_info.is_valid_tile: if cutlass.const_expr(self.overlapping_accum): tmem_stage_idx = acc_consumer_state.phase @@ -904,20 +1215,37 @@ def run( if cur_was_linear1: cute.arch.cp_async_bulk_commit_group() cute.arch.cp_async_bulk_wait_group(0, read=True) - cute.arch.fence_acq_rel_gpu() - elif cutlass.const_expr(self.token_back_by_dispatch): - cute.arch.fence_acq_rel_gpu() wait_only_named_barrier.arrive_and_wait() # Publish completion for the work tile snapshotted above. if cur_was_linear1: - fc1_epi.signal_fc1_done(prev_work_tile_info) + flag_tracker = fc1_epi.signal_fc1_done(prev_work_tile_info, + work_tile_info, + flag_tracker) else: - fc2_epi.signal_fc2_done(prev_work_tile_info) + flag_tracker = fc2_epi.signal_fc2_done(prev_work_tile_info, + work_tile_info, + flag_tracker) + # Tail flush + flag_tracker.fire() + + +class _ImmutableAfterInit: + """Froze at the point calling `_freeze()`""" + + def __setattr__(self, name, value): + if self.__dict__.get("_frozen_", False): + raise AttributeError( + f"{type(self).__name__} is immutable after __init__ " + f"(cannot set {name!r}).") + object.__setattr__(self, name, value) + + def _freeze(self) -> None: + object.__setattr__(self, "_frozen_", True) # Device only object -class SwapABFc1Epilogue: +class SwapABFc1Epilogue(_ImmutableAfterInit): def __init__( self, @@ -953,6 +1281,7 @@ def __init__( self.fc1_output_sf = fc1_output_sf self.fc1_done_counter = fc1_done_counter self.optional_epi_args = optional_epi_args + self._freeze() def __getattr__(self, name): return getattr(object.__getattribute__(self, "base"), name) @@ -971,8 +1300,9 @@ def __new_from_mlir_values__(self, return self @cute.jit - def signal_fc1_done(self, work_tile_info): - # Only in-bound intermediate_downproj tiles signal + def signal_fc1_done(self, work_tile_info, next_work_tile_info, + flag_tracker): + # Only in-bound intermediate_downproj tiles signal; OOB -> null slot. if cutlass.const_expr(self.static_expert_shape is None or self.intermediate_downproj % self.cluster_tile_intermediate_downproj != 0): @@ -981,23 +1311,25 @@ def signal_fc1_done(self, work_tile_info): < self.fc1_output.shape[1]) else: in_bound = True + slot = (work_tile_info.cumulative_token_block_count + + work_tile_info.tile_n_idx) + flag_addr = Int64(0) if in_bound: - if self.tidx == 0: - slot = work_tile_info.cumulative_token_block_count + work_tile_info.tile_n_idx - _red_add_release_gpu_s32( - self.fc1_done_counter.iterator + slot, - cutlass.Int32(1), - ) + flag_addr = (self.fc1_done_counter.iterator + slot).toint() + return flag_tracker.accumulate( + next_work_tile_info.phase, + self.fc1_epi_flag_batch, + flag_addr, + ) @cute.jit def __call__( - self, - work_tile_info: MoEWorkTileInfo, - tmem_acc_tensor: cute.Tensor, # (cta_tile_m, cta_tile_n) - acc_pipeline, - acc_consumer_state, - is_odd_turn: cutlass.Int32, - ): + self, + work_tile_info: MoEWorkTileInfo, + tmem_acc_tensor: cute.Tensor, # (cta_tile_m, cta_tile_n) + acc_pipeline, + acc_consumer_state, + is_odd_turn: cutlass.Int32): # (tokens_this_expert, intermediate_down, 1) real_fc1_output, _ = self.sched_ext.get_gmem_tensor( "c", @@ -1024,14 +1356,12 @@ def __call__( norm_const = None # (cta_tile_m, cta_tile_n) -> (epi_tile_m, epi_tile_n, iters) tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( - tmem_acc_tensor, - (self._EpilogueFc1IntermediateGateUpTileSize, - self._EpilogueTokenTileSize), - )[None, None, 0, None] + tmem_acc_tensor, (self._EpilogueFc1IntermediateGateUpTileSize, + self._EpilogueTokenTileSize))[None, None, 0, None] acc_pipeline.consumer_wait(acc_consumer_state) iket.range_push("fc1_epi") - valid_tokens = work_tile_info.valid_tokens_in_tile + valid_tokens = work_tile_info.valid_tokens_in_cta_tile # Overlap path preloads two subtiles before releasing acc TMEM. unroll_tile_cnt = 2 if cutlass.const_expr(self.overlapping_accum) else 0 @@ -1055,9 +1385,9 @@ def __call__( # the tmem transpose consumes them. preload_subtile_first: Tuple[ cute.Tensor, cute.Tensor, cute.Tensor, - cute.Tensor] = (_TmemTranspose16x32Core.load_subtile_raw_acc( + cute.Tensor] = _TmemTranspose16x32Core.load_subtile_raw_acc( tmem_acc_tensor_tiled_by_epi_tile[None, None, - subtile_idx_first])) + subtile_idx_first]) # Release acc to next MMA unconditionally. cute.arch.fence_view_async_tmem_load() @@ -1068,9 +1398,9 @@ def __call__( # quadrant/offset invariants and opaque per-lane layout as above. preload_subtile_second: Tuple[ cute.Tensor, cute.Tensor, cute.Tensor, - cute.Tensor] = (_TmemTranspose16x32Core.load_subtile_raw_acc( + cute.Tensor] = _TmemTranspose16x32Core.load_subtile_raw_acc( tmem_acc_tensor_tiled_by_epi_tile[None, None, - subtile_idx_second])) + subtile_idx_second]) # Both unrolled subtiles borrow tmem_subtile_second as workspace. preload_pair = (preload_subtile_first, preload_subtile_second) @@ -1153,13 +1483,10 @@ def run_subtile( work_tile_info.tile_n_idx * self.cta_tile_n + subtile_idx * self._EpilogueTokenTileSize + self.lane_idx, work_tile_info.tile_n_idx * self.cta_tile_n + - subtile_idx * self._EpilogueTokenTileSize + self.lane_idx + 32, - ) + subtile_idx * self._EpilogueTokenTileSize + self.lane_idx + 32) if cutlass.const_expr(topk_score_tensor is not None): - topk_scores = ( - topk_score_tensor[current_two_token_idices[0]], - topk_score_tensor[current_two_token_idices[1]], - ) + topk_scores = (topk_score_tensor[current_two_token_idices[0]], + topk_score_tensor[current_two_token_idices[1]]) else: topk_scores = None @@ -1213,52 +1540,44 @@ def run_subtile( token_0_32_pre_quant_pre_trans = self.alpha_swiglu_clamp( gate_token_0_32, up_token_0_32, alpha_val) + # gate_token_32_64 / up_token_32_64 are already in the transpose input + # distribution (see TmemTranspose16x32 / load_subtile_raw_acc). token_32_64_tmem_trans = TmemTranspose32x32Inplace( tmem_subtile_tensor.iterator, - reg_tensor_top=TensorWithContract( - tensor=gate_token_32_64, - contract=TmemTranspose16x32.InputContract, - ), - reg_tensor_bot=TensorWithContract( - tensor=up_token_32_64, - contract=TmemTranspose16x32.InputContract, - ), + reg_tensor_top=gate_token_32_64, + reg_tensor_bot=up_token_32_64, ) - # (epi_tid, vid) -> (token_idx, intermediate_output_idx), each lane hold (token_1, intermediate_16) in this rmem tensor. - gate_token_32_64_trans_pre_act, up_token_32_64_trans_pre_act = ( - token_32_64_tmem_trans.from_r1_perm_until_last_store()) + # Transpose output: each lane holds (token_1, intermediate_16); tmem_dp + # = lane_idx (token), tmem_col = elem_idx (intermediate output idx). + gate_token_32_64_trans_pre_act, up_token_32_64_trans_pre_act = token_32_64_tmem_trans.from_r1_perm_until_last_store( + ) token_32_64_pre_quant = self.alpha_swiglu_clamp( - gate_token_32_64_trans_pre_act.tensor, - up_token_32_64_trans_pre_act.tensor, + gate_token_32_64_trans_pre_act, + up_token_32_64_trans_pre_act, alpha_val, ) token_0_32_tmem_trans = TmemTranspose16x32( tmem_subtile_tensor.iterator, Region.Top, - reg_tensor=TensorWithContract( - tensor=token_0_32_pre_quant_pre_trans, - contract=TmemTranspose16x32.InputContract, - ), + reg_tensor=token_0_32_pre_quant_pre_trans, ) token_0_32_pre_quant = token_0_32_tmem_trans.from_r1_perm_until_last_store( - ).tensor + ) # Step 2: Quant - self.nvfp4_quant( - work_tile_info=work_tile_info, - two_token=(token_0_32_pre_quant, token_32_64_pre_quant), - topk_scores=topk_scores, - norm_const=norm_const, - intermediate_output_size=cute.size(fc1_output, 1), - fc1_output_sf=fc1_output_sf, - subtile_idx=subtile_idx, - ) + self.nvfp4_quant(work_tile_info=work_tile_info, + two_token=(token_0_32_pre_quant, + token_32_64_pre_quant), + topk_scores=topk_scores, + norm_const=norm_const, + intermediate_output_size=cute.size(fc1_output, 1), + fc1_output_sf=fc1_output_sf, + subtile_idx=subtile_idx) # Step 3: TMASTG - cute.arch.fence_proxy("async.shared", space="cta") # (token_64, intermeidate_64) fc1_smem = self.smem_tensor[None, None, subtile_idx] # (token, intermediate_down, l=1) -> (cta_token, cta_intermediate_down) @@ -1270,8 +1589,8 @@ def run_subtile( fc1_gmem_subtile_view = cute.flat_divide( fc1_gmem_cta_view, (self._EpilogueTokenTileSize, - self._EpilogueFc1IntermediateDownTileSize), - )[None, None, subtile_idx, 0] + self._EpilogueFc1IntermediateDownTileSize))[None, None, + subtile_idx, 0] tma_smem_src, tma_gmem_dst = cpasync.tma_partition( self.fc1_tma_atom, 0, @@ -1286,9 +1605,11 @@ def run_subtile( barrier_id=subtile_bar_id, num_threads=self._EpilogueWarpCnt * 32, ) + cute.arch.fence_proxy("async.shared", space="cta") if self.warp_idx == subtile_idx: tma_ready_to_read_smem_named_barrier.arrive_and_wait() with cute.arch.elect_one(): + # if work_tile_info.tile_m_idx * (self.cta_tile_m // 2) < cute.size(fc1_output, 1): cute.copy(self.fc1_tma_atom, tma_smem_src, tma_gmem_dst) else: tma_ready_to_read_smem_named_barrier.arrive() @@ -1300,7 +1621,7 @@ def alpha_swiglu_clamp( Tensor, # Raw fc1 acc (pre-dequant); even-size 1D fp32 rmem up_rmem: cute. Tensor, # Raw fc1 acc (pre-dequant); even-size 1D fp32 rmem - alpha_val: Optional[cutlass.Float32], + alpha_val: Optional[cutlass.Float32] ) -> cute.Tensor: # ── Input contract checks (compile-time): fp32, 1D, even-count, rmem ── # Wrapped in const_expr so the DSL evaluates them at trace time and the @@ -1443,15 +1764,12 @@ def nvfp4_quant( intermediate_output_size: cutlass.Int32, fc1_output_sf: cute. Tensor, # MoE domain (token_this_rank, intermediate_down, 1) - subtile_idx: cutlass.Int32, - ): - _Nvfp4RcpLimit = 1.0 / 6.0 # 1 / max abs of Float4E2M1FN (= 6.0) - _Fp32Max = 3.40282346638528859812e38 + subtile_idx: cutlass.Int32): # ``two_token`` are the two post-swiglu, transposed token rmem tensors; # each lane holds one token's ``sf_vec_size`` (=16, one NVFP4 SF block) # intermediate-output values. half 0 -> token (lane), half 1 -> (lane+32). # - # Per token (ported from PostSwigluHalf._gen_sfc_quantize + stg_sfc + r2s): + # Per token: # 1. (Path A) pre-multiply topk weight into the values, if present. # 2. absmax over the (weighted) block. # 3. sfc = absmax * (1/6) * norm_const -> E4M3 scale factor. @@ -1464,8 +1782,10 @@ def nvfp4_quant( # norm_const is treated like alpha_val: None => behaves as 1.0 (factors # const-elided, not multiplied by 1.0). n = cute.size(two_token[0]) - rcp_limit = cutlass.Float32(_Nvfp4RcpLimit) - fp32_max = cutlass.Float32(_Fp32Max) + # The core block quant (amax + e4m3 sfc + capped/masked acc_scale + e2m1 + # cvt) is QuantImpl's job; this method still owns the topk pre-multiply, + # the sfc store, and the STS.64 into the shared output stage. + quant = QuantImpl("nvfp4", "regs_in_thread") intermediate_idx = (work_tile_info.tile_m_idx * (self.cta_tile_m // 2) + self.warp_idx * Nvfp4BlockSize) @@ -1505,39 +1825,12 @@ def nvfp4_quant( for i in cutlass.range_constexpr(0, n): weighted[i] = tok[i] - # 2) absmax over the block. - absmax = cutlass.Float32(0.0) - for i in cutlass.range_constexpr(0, n): - v = weighted[i] - absmax = cute.arch.fmax(absmax, cute.arch.fmax(v, -v)) - - # 3) scale factor. - if cutlass.const_expr(norm_const is not None): - sfc_fp32 = absmax * rcp_limit * norm_const - else: - sfc_fp32 = absmax * rcp_limit - sfc_e4m3 = sfc_fp32.to(self.fc1_output_sf_dtype) - sfc_rt = cutlass.Float32(sfc_e4m3) - - # 4) acc_scale = norm_const * rcp(sfc), capped, with sfc==0 guard. - if cutlass.const_expr(norm_const is not None): - acc_scale = norm_const * cute.arch.rcp_approx(sfc_rt) - else: - acc_scale = cute.arch.rcp_approx(sfc_rt) - acc_scale = cute.arch.fmin(acc_scale, fp32_max) - mask = cute.arch.fmin(sfc_rt * cutlass.Float32(1e30), - cutlass.Float32(1.0)) - acc_scale = acc_scale * mask - - scaled = cute.make_rmem_tensor((n, ), cutlass.Float32) - acc_scale_pair = (acc_scale, acc_scale) - for i in cutlass.range_constexpr(0, n, 2): - s0, s1 = cute.arch.mul_packed_f32x2( - (weighted[i], weighted[i + 1]), acc_scale_pair) - scaled[i] = s0 - scaled[i + 1] = s1 + # 2) Core block quant: amax + e4m3 sfc + capped/masked acc_scale + + # e2m1 cvt. One 16-wide block -> one e2m1 reg tensor + one sfc. + fp4_regs, sfc_regs = quant(weighted, norm_const=norm_const) + sfc_e4m3 = sfc_regs[0] - # 5) scale-factor store (predicate const-elided when statically + # 3) scale-factor store (predicate const-elided when statically # in-bound, mirroring signal_fc1_done's intermediate predicate). if cutlass.const_expr(self.static_expert_shape is None or self.intermediate_downproj % @@ -1549,9 +1842,7 @@ def nvfp4_quant( fc1_output_sf[token_idx_pair[half], intermediate_idx, 0] = sfc_e4m3 - # 6) NVFP4 cvt + STS.64 into this subtile's shared output stage. - fp4_regs = cute.make_rmem_tensor((n, ), cutlass.Float4E2M1FN) - fp4_regs.store(scaled.load().to(cutlass.Float4E2M1FN)) + # 4) STS.64 the e2m1 into this subtile's shared output stage. # ((1, 16), (token_tile_size, warp_cnt)) -> (16) smem_thread_row = smem_tiled[(0, None), (self.lane_idx + 32 * half, self.warp_idx)] @@ -1562,79 +1853,396 @@ def nvfp4_quant( ) -""" -Acc to pre-store process: - some kind of ldtm -> f2fp -> some kind of reorder - -Pre-store status: - (epi_tid, vid) -> rmem x (token, hidden) +@dataclasses.dataclass(frozen=True) +class Fc2ProcessPipeline(): + tmem_acc_load: Callable + f2fp: Callable + post_f2fp_reorder: Callable + store_function: Callable + # Kept as a finer-grained, elem-level reading aid for the store-out layout + # (never evaluated); ``store_out_mapping`` is the per-issue form that the + # router actually evaluates at runtime to drive metadata / pointer math. + fc2_cta_tile_contract: Contract + store_out_mapping: Contract # data plane, per-issue + require_tmem_trans: bool + # SF plane per-issue mapping; None for the bf16 (unquantized) paths. + sf_store_out_mapping: Optional[Contract] = None -Store process: - Mapping: (epi_tid, iter_idx) -> (token, topk, hidden). - This defines in each sending, which thread(s) send which part to which dst. However, this is impl-irrevalent. - Always starts at rmem x (token, hidden)? -""" +# Device only object +class SwapABFc2Epilogue(_ImmutableAfterInit): -def eval_function_mapping(contract: Contract, **domain_coord): - """Evaluate a FunctionMapping contract at runtime. + def __init__( + self, + base: SwapABSwigluFp4Epilogue, + tidx: cutlass.Int32, + epi_smem_storage, + fc2_output: cute.Tensor, # MoE domain (token, topk, hidden) + token_comm_args: TokenCommArgs, + optional_epi_args: NvFp4OptinalEpiArgs, + ): + self.base = base + self.tidx = tidx % (base._EpilogueWarpCnt * 32) + self.warp_idx = self.tidx // 32 + self.lane_idx = self.tidx % 32 + self.fc2_output = fc2_output + self.token_comm_args = token_comm_args + self.optional_epi_args = optional_epi_args + if cutlass.const_expr(base.fc2_use_bulk): + wire_dtype = base.combine_format.act_dtype + fc2_smem_rows = 32 + if cutlass.const_expr(fc2_smem_rows * + base._EpilogueFc2HiddenTileSize * + wire_dtype.width // 8 > base.epi_smem_bytes): + raise ValueError("fc2 UBLK data smem exceeds epi_smem budget.") + self.smem_tensor = cute.make_tensor( + cute.recast_ptr( + epi_smem_storage.epi_smem.data_ptr(), + dtype=wire_dtype, + ), + cute.make_layout( + (fc2_smem_rows, base._EpilogueFc2HiddenTileSize), + stride=(base._EpilogueFc2HiddenTileSize, 1), + ), + ) + self.process_pipeline = make_fc2_ublk_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ) + else: + self.smem_tensor = None + if cutlass.const_expr(base.reduce_topk_in_kernel): + self.process_pipeline = make_fc2_redg_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ) + else: + self.process_pipeline = make_fc2_stg_process_pipeline( + combine_format=base.combine_format, + cta_token_tile_size=base.cta_tile_n, + cta_hidden_tile_size=base.cta_tile_m, + ) + self._freeze() - This is intentionally local to the fc2 epilogue refactor for now. It only - supports FunctionMapping-backed contracts whose Python function can run in - CuTe tracing context; table-backed runtime eval can be designed later in - contract.py. - """ - if not isinstance(contract.mapping, FunctionMapping): - raise TypeError("runtime contract eval requires a FunctionMapping") - - result = contract.mapping.function(**domain_coord) - if isinstance(result, dict): - return result - if isinstance(result, (tuple, list)): - if len(result) != contract.codomain.rank: - raise ValueError( - "FunctionMapping result rank does not match codomain rank: " - f"{len(result)} vs {contract.codomain.rank}") - return { - name: result[i] - for i, name in enumerate(contract.codomain.names) - } - if contract.codomain.rank == 1: - return {contract.codomain.names[0]: result} - raise TypeError( - "FunctionMapping runtime eval must return dict/tuple/list, or scalar for rank-1 codomain" - ) + def __getattr__(self, name): + return getattr(object.__getattribute__(self, "base"), name) + def __extract_mlir_values__(self) -> List[ir.Value]: + # See SwapABFc1Epilogue.__extract_mlir_values__: this helper carries + # only loop-invariant Python context. It intentionally serializes no + # MLIR values, so changing it to store loop-carried state would be a + # correctness bug. + return [] -@dataclasses.dataclass(frozen=True) -class Fc2OutputRouter: - # (token, 3), 3 -> rank_idx, token_idx, top_k - # Later this will be changed to (token, 2), where top_k and rank_idx will be fused into 32bit. - # If metadata is None then this is a local write. - metadata: Optional[cute.Tensor] - direct_token_base_this_cta_tile: Optional[cutlass.Int32] - base_output: cute.Tensor # (token, topk, hidden) - hidden_base_this_cta_tile: Union[cutlass.Int32, int] - peer_rank_ptr_mapper: Optional[SymBufferDeviceBase] - valid_tokens_this_cta_tile: cutlass.Int32 - valid_hidden_this_cta_tile: Union[cutlass.Int32, int] - reduce_topk_in_kernel: bool - output_mapping: Contract # (epi_tid, iter_idx) -> (token_cta_tile, hidden_cta_tile). - epi_tid: cutlass.Int32 + def __new_from_mlir_values__(self, + values: List[ir.Value]) -> "SwapABFc2Epilogue": + assert len(values) == 0 + return self - # After metadata prefetch - dst_ptrs: Optional[cute.Tensor] = ( - None # i64 x (copy_iters_this_thread_cta_tile), fundamentally the pointers. - ) - valid: Optional[cute.Tensor] = None # (copy_iters_this_thread_cta_tile) + @cute.jit + def signal_fc2_done(self, work_tile_info, next_work_tile_info, + flag_tracker): + publish: cutlass.Constexpr = (self.token_back_by_dispatch + or self.combine_format.is_quantized) + if cutlass.const_expr(publish): + flag_addr = (self.token_comm_args.fc2_done_counter.iterator + + work_tile_info.expert_idx).toint() + else: + flag_addr = Int64(0) + no_fire: cutlass.Constexpr = not publish + return flag_tracker.accumulate(next_work_tile_info.phase, + self.fc2_epi_flag_batch, flag_addr, + no_fire) - def __post_init__(self) -> None: - if (self.metadata is None) == (self.direct_token_base_this_cta_tile - is None): - raise ValueError( - "Fc2OutputRouter requires exactly one of metadata or " - "direct_token_base_this_cta_tile.") - if (self.metadata is None) != (self.peer_rank_ptr_mapper is None): + @cute.jit + def _make_output_router( + self, + work_tile_info: MoEWorkTileInfo, + ) -> "Fc2OutputRouter": + task_tile_data_row_start = ( + work_tile_info.cumulative_data_physical_row + + work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n)) + hidden_base_this_cta_tile = (work_tile_info.tile_m_idx * + cutlass.Int32(self.cta_tile_m)) + valid_hidden_this_cta_tile = (cutlass.Int32(self.fc2_output.shape[2]) - + hidden_base_this_cta_tile) + if valid_hidden_this_cta_tile < 0: + valid_hidden_this_cta_tile = 0 + if valid_hidden_this_cta_tile > self._EpilogueFc2HiddenTileSize: + valid_hidden_this_cta_tile = self._EpilogueFc2HiddenTileSize + + metadata_u32 = None + peer_rank_ptr_mapper = None + data_token_base = task_tile_data_row_start + if cutlass.const_expr(self.token_comm_args is not None + and not self.token_back_by_dispatch): + metadata_u32 = cute.domain_offset( + (task_tile_data_row_start, 0), + cute.recast_tensor( + self.token_comm_args.token_src_metadata, + cutlass.Uint32, + ), + ) + peer_rank_ptr_mapper = self.token_comm_args.peer_rank_ptr_mapper + data_token_base = None + + if cutlass.const_expr(self.combine_format.is_quantized): + base_outputs = (self.fc2_output, self.token_comm_args.fc2_output_sf) + token_bases = (data_token_base, task_tile_data_row_start) + output_mappings = ( + self.process_pipeline.store_out_mapping, + self.process_pipeline.sf_store_out_mapping, + ) + else: + base_outputs = self.fc2_output + token_bases = data_token_base + output_mappings = self.process_pipeline.store_out_mapping + + return Fc2OutputRouter( + metadata=metadata_u32, + token_bases=token_bases, + base_outputs=base_outputs, + hidden_base_this_cta_tile=hidden_base_this_cta_tile, + peer_rank_ptr_mapper=peer_rank_ptr_mapper, + valid_tokens_this_cta_tile=work_tile_info.valid_tokens_in_cta_tile, + valid_hidden_this_cta_tile=valid_hidden_this_cta_tile, + reduce_topk_in_kernel=self.reduce_topk_in_kernel, + output_mappings=output_mappings, + epi_tid=self.tidx, + combine_format=self.combine_format, + ).prefetch() + + @cute.jit + def __call__( + self, + work_tile_info: MoEWorkTileInfo, + tmem_acc_tensor: cute.Tensor, + acc_pipeline, + acc_consumer_state, + is_odd_turn: cutlass.Int32, + ): + # subtile-irrelevant hoist: fc2 alpha scales raw fc2 accumulators before f2fp. + if cutlass.const_expr(self.optional_epi_args.fc2_alpha is not None): + alpha_val = self.optional_epi_args.fc2_alpha[ + work_tile_info.expert_idx] + else: + alpha_val = None + acc_ready = False + if not work_tile_info.peek_ready: + acc_ready = True + acc_pipeline.consumer_wait(acc_consumer_state) + fc2_output_router = self._make_output_router(work_tile_info) + # (cta_tile_m, cta_tile_n) -> (epi_tile_m, epi_tile_n, iters) + tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( + tmem_acc_tensor, (self._EpilogueFc2HiddenTileSize, + self._EpilogueTokenTileSize))[None, None, 0, None] + + acc_pipeline.consumer_wait(acc_consumer_state, acc_ready) + iket.range_push("fc2_epi") + valid_tokens = work_tile_info.valid_tokens_in_cta_tile + + # Overlap path preloads two subtiles before releasing acc TMEM. + unroll_tile_cnt = 2 if cutlass.const_expr( + self.overlapping_accum + and self.process_pipeline.require_tmem_trans) else 0 + remain_subtile_cnt = self.subtile_cnt - unroll_tile_cnt + + if cutlass.const_expr(unroll_tile_cnt > 0): + subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - + is_odd_turn) % cutlass.Int32(self.subtile_cnt) + subtile_idx_second = (cutlass.Int32(self.subtile_cnt + 1) - + is_odd_turn) % cutlass.Int32(self.subtile_cnt) + + # preload_subtile_first: subtile_idx_first's raw PRE-transpose acc, LDTM'd by + # all 128 epi threads into 4 reg tensors == the 4 quadrants of the subtile's + # (128 tmem_dp x 64 tmem_col) footprint. Only these raw-TMEM offsets are + # guaranteed: + # reg[0]/reg[1], reg[2]/reg[3] : top vs bot -> 16 apart in tmem_dp + # reg[0]/reg[2], reg[1]/reg[3] : 1st vs 2nd half -> 32 apart in tmem_col + # (so reg[0..1] = the first 128x32, reg[2..3] = the second 128x32 of the 128x64.) + # The per-lane (lane_idx, elem_idx) -> (tmem_dp, tmem_col) layout INSIDE each + # reg tensor is opaque -- do not assume it; it only becomes well-defined once + # the tmem transpose consumes them. + preload_subtile_first: Tuple[ + cute.Tensor, cute.Tensor, cute.Tensor, + cute.Tensor] = _TmemTranspose16x32Core.load_subtile_raw_acc( + tmem_acc_tensor_tiled_by_epi_tile[None, None, + subtile_idx_first]) + + # Release acc to next MMA unconditionally. + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + # preload_subtile_second: same 128 tmem_dp x 64 tmem_col footprint, but for + # subtile_idx_second (the other token subtile, not the 2nd col-half). Same + # quadrant/offset invariants and opaque per-lane layout as above. + preload_subtile_second: Tuple[ + cute.Tensor, cute.Tensor, cute.Tensor, + cute.Tensor] = _TmemTranspose16x32Core.load_subtile_raw_acc( + tmem_acc_tensor_tiled_by_epi_tile[None, None, + subtile_idx_second]) + + # Both unrolled subtiles borrow tmem_subtile_second as workspace. + preload_pair = (preload_subtile_first, preload_subtile_second) + subtile_idx_pair = (subtile_idx_first, subtile_idx_second) + for i in cutlass.range_constexpr(unroll_tile_cnt): + if subtile_idx_pair[i] * cutlass.Int32( + self._EpilogueTokenTileSize) < valid_tokens: + self.run_subtile( + subtile_idx=subtile_idx_pair[i], + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[ + None, None, subtile_idx_second], + preload_acc=preload_pair[i], + fc2_output_router=fc2_output_router, + alpha_val=alpha_val, + release_after_ldtm=False, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + ) + + if cutlass.const_expr(self.overlapping_accum and unroll_tile_cnt == 0): + release_after_ldtm = True + else: + release_after_ldtm = False + for i in cutlass.range(remain_subtile_cnt, unroll=1): + # for i in cutlass.range_constexpr(remain_subtile_cnt): + real_i = i + unroll_tile_cnt + if cutlass.const_expr(self.overlapping_accum): + subtile_idx = (cutlass.Int32(real_i + self.subtile_cnt) - + is_odd_turn) % cutlass.Int32(self.subtile_cnt) + else: + subtile_idx = cutlass.Int32(real_i) + + if subtile_idx * cutlass.Int32( + self._EpilogueTokenTileSize) < valid_tokens: + self.run_subtile( + subtile_idx=subtile_idx, + tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[ + None, None, subtile_idx], + preload_acc=None, + fc2_output_router=fc2_output_router, + alpha_val=alpha_val, + release_after_ldtm=release_after_ldtm, + acc_pipeline=acc_pipeline, + acc_consumer_state=acc_consumer_state, + ) + release_after_ldtm = False + + # Non-overlap-path release: at the natural task-tile boundary. + if cutlass.const_expr(not self.overlapping_accum): + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + + @cute.jit + def run_subtile( + self, + subtile_idx: cutlass.Int32, + # (hidden_tile, token_subtile), fundamentally (epi_tile_m, epi_tile_n) + tmem_subtile_tensor: cute.Tensor, + preload_acc: Optional[Tuple[cute.Tensor, cute.Tensor, cute.Tensor, + cute.Tensor]], + fc2_output_router: "Fc2OutputRouter", + alpha_val: Optional[cutlass.Float32], + release_after_ldtm: Union[cutlass.Boolean, bool], + acc_pipeline, + acc_consumer_state, + ): + process_pipeline = self.process_pipeline + if cutlass.const_expr(preload_acc is None): + loaded = process_pipeline.tmem_acc_load( + tmem_subtile_tensor=tmem_subtile_tensor, + epi=self, + ) + if release_after_ldtm: + cute.arch.fence_view_async_tmem_load() + acc_pipeline.consumer_release(acc_consumer_state) + else: + loaded = preload_acc + + casted = process_pipeline.f2fp( + *loaded, + alpha_val=alpha_val, + ) + # reorder returns a bare RMEM fragment in the store's expected pre-store + # distribution; reorder + store are paired 1:1 inside the pipeline. + pre_store = process_pipeline.post_f2fp_reorder( + casted=casted, + tmem_subtile_view=tmem_subtile_tensor, + ) + process_pipeline.store_function( + epi=self, + subtile=pre_store, + subtile_idx=subtile_idx, + fc2_output_router=fc2_output_router, + ) + + +@dataclasses.dataclass(frozen=True) +class Fc2OutputRouter: + # (token, 3), 3 -> rank_idx, token_idx, top_k + # Later this will be changed to (token, 2), where top_k and rank_idx will be fused into 32bit. + # If metadata is None then this is a local write. + metadata: Optional[cute.Tensor] + # token + possible sf + token_bases: Union[Tuple[Optional[cutlass.Int32], cutlass.Int32], + Optional[cutlass.Int32]] + base_outputs: Union[Tuple[cute.Tensor, cute.Tensor], + cute.Tensor] # (token, topk, hidden) + hidden_base_this_cta_tile: Union[cutlass.Int32, int] + peer_rank_ptr_mapper: Optional[SymBufferDeviceBase] + valid_tokens_this_cta_tile: cutlass.Int32 + valid_hidden_this_cta_tile: Union[cutlass.Int32, int] + reduce_topk_in_kernel: bool + # Per-issue (epi_tid, iter_idx) -> (token_cta_tile, hidden_cta_tile). Data + # mapping, or (data mapping, sf mapping) when quantized. + output_mappings: Union[Tuple[Contract, Contract], Contract] + epi_tid: cutlass.Int32 + combine_format: CombineFormat + # After metadata prefetch + dst_ptrs: Optional[ + cute. + Tensor] = None # i64 x (copy_iters_this_thread_cta_tile), fundamentally the pointers. + valid: Optional[cute.Tensor] = None # (copy_iters_this_thread_cta_tile) + + @property + def data_output(self) -> cute.Tensor: + return self.base_outputs[0] if isinstance(self.base_outputs, + tuple) else self.base_outputs + + @property + def sf_output(self) -> Optional[cute.Tensor]: + # Present iff quantized; (pool_token, 1, hidden // sf_vec) rank-local. + return self.base_outputs[1] if isinstance(self.base_outputs, + tuple) else None + + @property + def data_token_base(self) -> Optional[cutlass.Int32]: + return self.token_bases[0] if isinstance(self.token_bases, + tuple) else self.token_bases + + @property + def sf_token_base(self) -> Optional[cutlass.Int32]: + return self.token_bases[1] if isinstance(self.token_bases, + tuple) else None + + @property + def data_mapping(self) -> Contract: + return self.output_mappings[0] if isinstance( + self.output_mappings, tuple) else self.output_mappings + + @property + def sf_mapping(self) -> Optional[Contract]: + return self.output_mappings[1] if isinstance(self.output_mappings, + tuple) else None + + def __post_init__(self) -> None: + if (self.metadata is None) == (self.data_token_base is None): + raise ValueError( + "Fc2OutputRouter requires exactly one of metadata or " + "a (data) token base.") + if (self.metadata is None) != (self.peer_rank_ptr_mapper is None): raise ValueError( "Fc2OutputRouter requires peer_rank_ptr_mapper iff metadata is set." ) @@ -1645,8 +2253,14 @@ def __post_init__(self) -> None: @cute.jit def prefetch(self) -> "Fc2OutputRouter": - iter_axis = self.output_mapping.domain.names.index("iter_idx") - copy_iters: cutlass.Constexpr[int] = self.output_mapping.domain.sizes[ + # Only the metadata (comm) path prefetches a pointer array: its + # metadata-derived address has long-latency LDGs worth issuing early. + # The local (no-comm) path computes its affine address on demand in + # get_dst() -- no array, hence no runtime-indexed local-memory spill. + if cutlass.const_expr(self.metadata is None): + return self + iter_axis = self.data_mapping.domain.names.index("iter_idx") + copy_iters: cutlass.Constexpr[int] = self.data_mapping.domain.sizes[ iter_axis] valid = cute.make_rmem_tensor((copy_iters, ), cutlass.Int32) @@ -1656,7 +2270,7 @@ def prefetch(self) -> "Fc2OutputRouter": # We should check the SASS to ensure this happens. for iter_idx in cutlass.range_constexpr(copy_iters): coord = eval_function_mapping( - self.output_mapping, + self.data_mapping, epi_tid=self.epi_tid, iter_idx=iter_idx, ) @@ -1672,27 +2286,41 @@ def prefetch(self) -> "Fc2OutputRouter": if token_valid and hidden_valid: valid[iter_idx] = cutlass.Int32(1) if cutlass.const_expr(self.metadata is None): - dst_tokens = self.direct_token_base_this_cta_tile + token_in_tile + dst_tokens = self.data_token_base + token_in_tile dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile - dst_ptrs[iter_idx] = self.base_output[ - dst_tokens, None, dst_hidden].iterator.toint() + # Int64 token coord: dst_tokens*K*H overflows int32 once + # T*K*H exceeds 2^31 (data_output is (token, topk, hidden)). + dst_ptrs[iter_idx] = self.data_output[ + Int64(dst_tokens), None, dst_hidden].iterator.toint() else: - dst_rank = cutlass.Int32(self.metadata[token_in_tile, 0]) - dst_token = cutlass.Int32(self.metadata[token_in_tile, 1]) + md = TokenSrcMetadata.load(self.metadata.iterator.toint() + + Int64(token_in_tile) * + Int64(TokenSrcMetadata.nbytes)) + dst_rank = md.src_rank + dst_token = md.src_token dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile if cutlass.const_expr(not self.reduce_topk_in_kernel): - dst_topk = cutlass.Int32(self.metadata[token_in_tile, - 2]) + dst_topk = md.src_topk else: dst_topk = 0 + # Int64 token coord: domain_offset on (token, topk, hidden) + # computes dst_token*K*H, which overflows int32 once T*K*H > 2^31. + # byte_align mirrors get_data_dst: the STG vector is 16 B + # for e2m1 wires (odd warps at row base +16 B), 32 B for + # fp8/bf16. dst_ptrs[ iter_idx] = self.peer_rank_ptr_mapper.ptr_map_to_rank( cute.domain_offset( - (dst_token, dst_topk, dst_hidden), - self.base_output).iterator, + (Int64(dst_token), dst_topk, dst_hidden), + self.data_output).iterator, dst_rank, - ).toint() + byte_align=min( + 32, + (min(32, 256 // + self.data_output.element_type.width) * + self.data_output.element_type.width) // 8, + )).toint() return dataclasses.replace( self, @@ -1700,157 +2328,113 @@ def prefetch(self) -> "Fc2OutputRouter": valid=valid, ) - # Return a tuple of (src, dst) tensors, each represents copy_atom's one call. @cute.jit - def resolve( + def get_data_dst( self, - copy_src: cute.Tensor, # (v, rest...) - iters_per_subtile: int, - subtile_idx: Union[cutlass.Int32, int], - ) -> Tuple[Tuple[cute.Tensor, cute.Tensor], ...]: - if cutlass.const_expr(self.dst_ptrs is None or self.valid is None): - raise ValueError( - "Fc2OutputRouter.resolve requires prefetch() first.") - # Normalize any strategy-specific source view into the canonical copy - # iterator form: - # - # ((atom_v, rest_v), rests...) - # - # The trailing ``rest`` modes enumerate one copy-atom issue inside the - # current subtile; their product must equal ``iters_per_subtile``. - # Everything before those trailing modes is the copy atom payload - # (``atom_v`` elements). We intentionally flatten first because callers - # may hand us nested CuTe layouts whose hierarchy is meaningful to their - # local algorithm but irrelevant to the final copy issue schedule. After - # finding the trailing rest modes, two ``group_modes`` calls make rank 0 - # the payload and rank 1 the full rest/iter space; coalescing with - # ``target_profile=(1, 1)`` preserves that two-rank profile while - # simplifying each side's internal layout. - flat_src = cute.flatten(copy_src) - flat_rank: cutlass.Constexpr[int] = cute.rank(flat_src) - - rest_start = flat_rank - rest_size = 1 - for mode in cutlass.range_constexpr(flat_rank - 1, -1, -1): - rest_start = mode - rest_size *= cute.size(flat_src, mode=[mode]) - if cutlass.const_expr(rest_size == iters_per_subtile): - break - if cutlass.const_expr(rest_size != iters_per_subtile): - raise ValueError( - "Fc2OutputRouter.resolve: trailing rest modes must multiply " - f"to iters_per_subtile={iters_per_subtile}, got {rest_size}.") - if cutlass.const_expr(rest_start == 0): - raise ValueError( - "Fc2OutputRouter.resolve requires at least one atom payload mode " - "before the trailing rest modes.") - - atom_v: cutlass.Constexpr[int] = cute.size( - flat_src) // iters_per_subtile - atom_rest = cute.group_modes(flat_src, 0, rest_start) - atom_rest = cute.group_modes(atom_rest, 1, cute.rank(atom_rest)) - atom_rest = cute.coalesce(atom_rest, target_profile=(1, 1)) - - single_copy_layout = cute.make_layout( - ((atom_v, 1), ), - stride=((1, 0), ), - ) - subtile_iter_base = cutlass.Int32(subtile_idx) * cutlass.Int32( - iters_per_subtile) - - copy_pairs = () - for local_iter in cutlass.range_constexpr(iters_per_subtile): - global_iter = subtile_iter_base + cutlass.Int32(local_iter) - src_atom = atom_rest[None, local_iter] - copy_src_i = cute.make_tensor(src_atom.iterator, single_copy_layout) - dst_ptr = cute.make_ptr( - copy_src.element_type, - self.dst_ptrs[global_iter], - AddressSpace.gmem, - assumed_align=32, - ) - copy_dst_i = cute.make_tensor(dst_ptr, single_copy_layout) - copy_pairs = copy_pairs + ((copy_src_i, copy_dst_i), ) - - return copy_pairs + iter_idx: Union[int, cutlass.Int32], + ) -> Tuple[cute.Pointer, cutlass.Int32]: + """Per-issue DATA destination: gmem pointer + validity predicate. + The router owns ``data_output`` so the caller never re-assembles a + pointer from a raw int; it just builds its own copy tensor (STG) or + feeds the pointer to inline asm (REDG/UBLK). -@dataclasses.dataclass(frozen=True) -class Fc2ProcessPipeline: - tmem_acc_load: Callable - f2fp: Callable - post_f2fp_reorder: Callable - store_function: Callable - pre_store_contract: Contract - fc2_cta_tile_contract: Contract - store_out_mapping: Contract - require_tmem_trans: bool + Alignment is unified at 32 B: only STG feeds this pointer to a real + ``cute.copy`` (256 b vector store, genuinely 32 B aligned); REDG/UBLK + only ``ptrtoint`` it for inline-asm issue, where the hint is inert. + """ + if cutlass.const_expr(self.metadata is None): + # no-comm: on-demand affine address (no prefetched array). The + # invariant base hoists out of the caller's loop via CSE; a + # constexpr iter folds the per-issue offset into the store. + coord = eval_function_mapping( + self.data_mapping, + epi_tid=self.epi_tid, + iter_idx=iter_idx, + ) + token_in_tile = cutlass.Int32(coord["token_in_cta_tile"]) + hidden_in_tile = cutlass.Int32(coord["hidden_in_cta_tile"]) + pred = cutlass.Int32(0) + addr = cutlass.Int64(0) + if (token_in_tile < self.valid_tokens_this_cta_tile + and hidden_in_tile < cutlass.Int32( + self.valid_hidden_this_cta_tile)): + pred = cutlass.Int32(1) + dst_tokens = self.data_token_base + token_in_tile + dst_hidden = hidden_in_tile + self.hidden_base_this_cta_tile + # Int64 token coord: dst_tokens*K*H overflows int32 once T*K*H > 2^31. + addr = self.data_output[Int64(dst_tokens), None, + dst_hidden].iterator.toint() + else: + # comm: read the pointer / validity prefetched by prefetch(). + addr = self.dst_ptrs[iter_idx] + pred = self.valid[iter_idx] + # Alignment must match the real STG granularity: the store vector is + # min(32, 256 // width) elements, i.e. 32 B for fp8/bf16 wires but + # only 16 B for e2m1 (4-bit) -- odd warps then land on row base +16 B, + # so an align-32 promise would be UB for the fp4 combine format. + _stg_bytes = (min(32, 256 // self.data_output.element_type.width) * + self.data_output.element_type.width) // 8 + ptr = cute.make_ptr( + self.data_output.element_type, + addr, + AddressSpace.gmem, + assumed_align=min(32, _stg_bytes), + ) + return ptr, pred + @cute.jit + def get_sf_dst( + self, + iter_idx: Union[int, cutlass.Int32], + ) -> Tuple[cute.Pointer, cutlass.Int32]: + """Per-issue SF destination: rank-local gmem pointer + validity predicate. + + SF never goes to a peer (it is staged locally and pushed token-contiguously + by the dispatch / standalone warps), so this is always the affine local + address -- no metadata routing, no prefetch. ``sf_output`` is the broadcast + plane ``(pool_token, 1, (sf_vec, hidden//sf_vec)):(., ., (0, 1))``, so the + logical hidden coordinate folds to its scale block on indexing. + """ + coord = eval_function_mapping( + self.sf_mapping, + epi_tid=self.epi_tid, + iter_idx=iter_idx, + ) + token_in_tile = cutlass.Int32(coord["token_in_cta_tile"]) + hidden_in_tile = cutlass.Int32(coord["hidden_in_cta_tile"]) + pred = cutlass.Int32(0) + addr = cutlass.Int64(0) + if (token_in_tile < self.valid_tokens_this_cta_tile and hidden_in_tile + < cutlass.Int32(self.valid_hidden_this_cta_tile)): + pred = cutlass.Int32(1) + sf_row = self.sf_token_base + token_in_tile + sf_hidden = hidden_in_tile + self.hidden_base_this_cta_tile + addr = self.sf_output[Int64(sf_row), None, + sf_hidden].iterator.toint() + # Per-block scale offsets are element-granular; claim the scale dtype's + # natural element alignment (e8m0 1 B / bf16 2 B). An align-4 promise + # here is UB: adjacent warps land on row base +1/+2/+3 for 1-byte + # scales. + sf_ptr = cute.make_ptr( + self.sf_output.element_type, + addr, + AddressSpace.gmem, + assumed_align=self.sf_output.element_type.width // 8, + ) + return sf_ptr, pred -# ============================================================================= -# fc2 STG strategy callables (subtile granularity) -# -# Faithful port of the original transpose+STG path, re-cut into the four -# Fc2ProcessPipeline steps. Originals in epilogue.py: -# - load : _TmemTranspose16x32Core.load_subtile_raw_acc -# - pack : Fc2AccLoadAndPack.__init__ (L986-997) -# - transpose : TmemTranspose16x32Packed (+ from_r1_perm_until_last_store) -# - unpack : Fc2UnpackPermuteStg._init_direct (L1510-1520) -# - store : Fc2UnpackPermuteStg._stg_direct (L1522-1605) -# -# All callables take the unified kwargs + ``**_`` (extras ignored; missing -# required -> TypeError). ``epi`` is the SwapABFc2Epilogue device object. -# ============================================================================= -# Subtile pre-store contract C (the pivot): each lane holds 64 bf16 values; -# vid in [0,32) -> token=lane (half 0), vid in [32,64) -> token=lane+32 (half 1); -# hidden = vid % 32 (this warp's 32-hidden span, natural order). -_Fc2StgSubtilePreStoreContract = Contract( - domain=Space(("lane_idx", "vid"), (32, 64)), - codomain=Space(("token_idx", "hidden_idx"), (64, 32)), - mapping=FunctionMapping(lambda lane_idx, vid: { - "token_idx": lane_idx + 32 * (vid // 32), - "hidden_idx": vid % 32, - }), -) - -# UBLK pre-store contract: after f2fp (and before R2S), each lane owns one -# hidden element across the 64 token positions of the subtile. This is -# warp-local + subtile-local; the store-out contract below is CTA-level and -# describes the later bulk issue rows, not this RMEM distribution. -_Fc2UblkSubtilePreStoreContract = Contract( - domain=Space(("lane_idx", "vid"), (32, 64)), - codomain=Space(("token_idx", "hidden_idx"), (64, 32)), - mapping=FunctionMapping(lambda lane_idx, vid: { - "token_idx": vid, - "hidden_idx": lane_idx, - }), -) - -# REDG pre-store contract: after the extra STTM + LDTM(16x256b.x2) -# reshuffle, each lane owns bf16 scalar elements arranged so every 4 -# consecutive elem_idx values form one red.v2.bf16x2 issue payload. -_Fc2RedgSubtilePreStoreContract = Contract( - domain=Space(("lane_idx", "elem_idx"), (32, 64)), - codomain=Space(("token_idx", "hidden_idx"), (64, 32)), - mapping=FunctionMapping( - lambda lane_idx, elem_idx: { - "token_idx": - (((elem_idx // 2) // 16) * 32 + (((elem_idx // 2) // 8) % 2) * 16 + - (((elem_idx // 2) // 2) % 2) * 8 + lane_idx // 4), - "hidden_idx": ((lane_idx % 4) * 4 + (((elem_idx // 2) // 4) % 2) * - 16 + ((elem_idx // 2) % 2) * 2 + (elem_idx % 2)), - }), -) - - -# TODO: Enable for non-BF16 dtypes -def make_fc2_stg_cta_store_out_contract(fc2_output_dtype: Type[cutlass.Numeric], +def make_fc2_stg_cta_store_out_contract(combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int): - assert cta_hidden_tile_size == 128 - assert cta_token_tile_size % 64 == 0 - assert fc2_output_dtype.width == 16 + assert (cta_hidden_tile_size == 128) + assert (cta_token_tile_size % 64 == 0) + wire_dtype = combine_format.act_dtype + assert wire_dtype.width in (4, 8, + 16), "fc2 STG wire dtype must be fp4/fp8/bf16." + elems_per_stg = min(256 // wire_dtype.width, 32) + stgs_per_hidden32 = 32 // elems_per_stg fundamental_mapping = Contract( domain=Space(("epi_tid", "elem_idx"), (128, cta_token_tile_size)), codomain=Space(("token_in_cta_tile", "hidden_in_cta_tile"), @@ -1859,28 +2443,48 @@ def make_fc2_stg_cta_store_out_contract(fc2_output_dtype: Type[cutlass.Numeric], lambda epi_tid, elem_idx: { "token_in_cta_tile": epi_tid % 32 + elem_idx // 32 * 32, "hidden_in_cta_tile": elem_idx % 32 + epi_tid // 32 * 32, - }), - ) + })) store_out_mapping = Contract( domain=Space(("epi_tid", "iter_idx"), - (128, 32 // 16 * cta_token_tile_size // 32)), + (128, stgs_per_hidden32 * cta_token_tile_size // 32)), codomain=Space(("token_in_cta_tile", "hidden_in_cta_tile"), (cta_token_tile_size, cta_hidden_tile_size)), mapping=FunctionMapping( lambda epi_tid, iter_idx: { - "token_in_cta_tile": epi_tid % 32 + iter_idx // 2 * 32, - "hidden_in_cta_tile": (iter_idx % 2) * 16 + epi_tid // 32 * 32, - }), - ) - return store_out_mapping, fundamental_mapping + "token_in_cta_tile": + epi_tid % 32 + iter_idx // stgs_per_hidden32 * 32, + "hidden_in_cta_tile": (iter_idx % stgs_per_hidden32) * + elems_per_stg + epi_tid // 32 * 32, + })) + sf_store_out_mapping = None + if combine_format.is_quantized: + + def stg_sf_mapping(epi_tid, iter_idx): + lane = epi_tid % 32 + warp = epi_tid // 32 + return { + "token_in_cta_tile": lane + iter_idx * 32, + "hidden_in_cta_tile": warp * 32, + } + + sf_store_out_mapping = Contract( + domain=Space(("epi_tid", "iter_idx"), + (128, cta_token_tile_size // 32)), + codomain=Space(("token_in_cta_tile", "hidden_in_cta_tile"), + (cta_token_tile_size, cta_hidden_tile_size)), + mapping=FunctionMapping(stg_sf_mapping), + ) + return store_out_mapping, sf_store_out_mapping, fundamental_mapping -def make_fc2_redg_cta_store_out_contract( - fc2_output_dtype: Type[cutlass.Numeric], cta_token_tile_size: int, - cta_hidden_tile_size: int): - assert cta_hidden_tile_size == 128 - assert cta_token_tile_size % 64 == 0 - assert fc2_output_dtype.width == 16 +def make_fc2_redg_cta_store_out_contract(combine_format: CombineFormat, + cta_token_tile_size: int, + cta_hidden_tile_size: int): + assert (cta_hidden_tile_size == 128) + assert (cta_token_tile_size % 64 == 0) + # In-kernel reduce is bf16-only and never quantized, so there is no SF plane. + assert (combine_format.act_dtype.width == 16) + assert not combine_format.is_quantized fundamental_mapping = Contract( domain=Space(("epi_tid", "elem_idx"), (128, cta_token_tile_size)), codomain=Space(("token_in_cta_tile", "hidden_in_cta_tile"), @@ -1894,8 +2498,7 @@ def make_fc2_redg_cta_store_out_contract( "hidden_in_cta_tile": ( (epi_tid // 32) * 32 + (epi_tid % 4) * 4 + (( (elem_idx // 4) % 4) // 2) * 16 + elem_idx % 4), - }), - ) + })) # SIMT REDG emits one 8B red.v2.bf16x2 per 4 hidden elements. Each # 64-token subtile contributes two token rows per lane and 8 hidden # segments per token row. @@ -1911,61 +2514,102 @@ def make_fc2_redg_cta_store_out_contract( (iter_idx % 4) % 2) * 8 + (epi_tid % 32) // 4), "hidden_in_cta_tile": ((epi_tid // 32) * 32 + (epi_tid % 4) * 4 + ((iter_idx % 4) // 2) * 16), - }), - ) - return store_out_mapping, fundamental_mapping + })) + return store_out_mapping, None, fundamental_mapping -def make_fc2_ublk_store_out_contract(fc2_output_dtype: Type[cutlass.Numeric], +def make_fc2_ublk_store_out_contract(combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int): - assert cta_hidden_tile_size == 128 - assert cta_token_tile_size % 64 == 0 - assert fc2_output_dtype.width == 16 + assert (cta_hidden_tile_size == 128) + assert (cta_token_tile_size % 64 == 0) + # UBLK pushes whole hidden rows by byte count, so the token/hidden mapping + # is element-indexed and dtype-independent (wire dtype only sets copy bytes). + assert combine_format.act_dtype.width in ( + 4, 8, 16), "fc2 UBLK wire dtype must be fp4/fp8/bf16." + assert (cta_token_tile_size <= 256) + max_token_cta_tile = 256 fundamental_mapping = Contract( domain=Space(("epi_tid", "elem_idx"), (128, cta_token_tile_size)), codomain=Space(("token_in_cta_tile", "hidden_in_cta_tile"), - (cta_token_tile_size, cta_hidden_tile_size)), + (max_token_cta_tile, cta_hidden_tile_size)), mapping=FunctionMapping( lambda epi_tid, elem_idx: { "token_in_cta_tile": - epi_tid % 8 + epi_tid // 32 * 8 + ((epi_tid % 32) // 8) * 32 + - elem_idx // cta_hidden_tile_size * 128, + elem_idx // cta_hidden_tile_size * 32 + epi_tid % 8 + epi_tid // + 32 * 8 + ((epi_tid % 32) // 8) * 64, "hidden_in_cta_tile": elem_idx % cta_hidden_tile_size, - }), - ) + })) store_out_mapping = Contract( - domain=Space(("epi_tid", "iter_idx"), - (128, (cta_token_tile_size + 127) // 128)), + domain=Space(("epi_tid", "iter_idx"), (128, 2)), codomain=Space(("token_in_cta_tile", "hidden_in_cta_tile"), - (cta_token_tile_size, cta_hidden_tile_size)), + (max_token_cta_tile, cta_hidden_tile_size)), mapping=FunctionMapping( lambda epi_tid, iter_idx: { "token_in_cta_tile": - epi_tid % 8 + epi_tid // 32 * 8 + - ((epi_tid % 32) // 8) * 32 + iter_idx * 128, + iter_idx * 32 + epi_tid % 8 + epi_tid // 32 * 8 + + ((epi_tid % 32) // 8) * 64, "hidden_in_cta_tile": 0, - }), - ) - return store_out_mapping, fundamental_mapping - + })) + # SF mapping: each lane owns one hidden across a 64-token subtile, and the + # sf_vec lanes of a block share the CREDUX scale -- so they split the block's + # tokens. sf_iter flattens (subtile, slot) to keep the whole cta-tile domain. + sf_store_out_mapping = None + if combine_format.is_quantized: + sf_vec = combine_format.scale_block + lanes_per_block = sf_vec + subtile_tokens = SwapABSwigluFp4Epilogue._EpilogueTokenTileSize + iters_per_subtile = subtile_tokens // lanes_per_block + n_subtiles = cta_token_tile_size // subtile_tokens + + def ublk_sf_mapping(epi_tid, iter_idx): + subtile_idx = iter_idx // iters_per_subtile + slot = iter_idx % iters_per_subtile + lane = epi_tid % 32 + warp = epi_tid // 32 + lane_in_block = lane % lanes_per_block + block_in_warp = lane // lanes_per_block + return { + "token_in_cta_tile": + subtile_idx * subtile_tokens + lane_in_block + + slot * lanes_per_block, + "hidden_in_cta_tile": + warp * 32 + block_in_warp * sf_vec, + } + + sf_store_out_mapping = Contract( + domain=Space(("epi_tid", "iter_idx"), + (128, n_subtiles * iters_per_subtile)), + codomain=Space(("token_in_cta_tile", "hidden_in_cta_tile"), + (max_token_cta_tile, cta_hidden_tile_size)), + mapping=FunctionMapping(ublk_sf_mapping), + ) + return store_out_mapping, sf_store_out_mapping, fundamental_mapping + + +# (...) -> ((atom_v, 1)) +@cute.jit +def wrap_into_copy_standard_layout(tensor: cute.Tensor): + tensor = cute.coalesce(cute.flatten(tensor)) + tensor = cute.append_ones(tensor, cute.rank(tensor) + 1) + tensor = cute.group_modes(tensor, 0, cute.rank(tensor) - 1) + tensor = cute.group_modes(tensor, 0, cute.rank(tensor)) + return tensor + @cute.jit def fc2_f2fp( *tensors, - fc2_output_dtype: Type[cutlass.Numeric], alpha_val: Optional[cutlass.Float32] = None, **_, ) -> cute.Tensor: - # cvt every input fp32 rmem -> fc2_output_dtype and concatenate, in order, - # into one flat rmem tensor. Each block is stored contiguously (no scalar - # element copy) at its running offset. + reorder_dtype = cutlass.BFloat16 total_size = 0 for t in tensors: total_size += cute.size(t) - converted_acc = cute.make_rmem_tensor((total_size, ), fc2_output_dtype) + converted_acc = cute.make_rmem_tensor((total_size, ), reorder_dtype) elems_processed = 0 for t in tensors: current_tensor_size = cute.size(t) @@ -1974,7 +2618,7 @@ def fc2_f2fp( cute.make_layout((current_tensor_size, )), ) if cutlass.const_expr(alpha_val is None): - dst.store(t.load().to(fc2_output_dtype)) + dst.store(t.load().to(reorder_dtype)) else: if cutlass.const_expr(current_tensor_size % 2 != 0): raise ValueError( @@ -1987,14 +2631,16 @@ def fc2_f2fp( (alpha_val, alpha_val)) scaled[i] = s0 scaled[i + 1] = s1 - dst.store(scaled.load().to(fc2_output_dtype)) + dst.store(scaled.load().to(reorder_dtype)) elems_processed += current_tensor_size return converted_acc @cute.jit -def post_f2fp_reorder_identity(*, casted: cute.Tensor, contract: Contract, **_): - return TensorWithContract(tensor=casted, contract=contract) +def post_f2fp_reorder_identity(*, casted: cute.Tensor, **_): + # UBLK: the f2fp output is already in the pre-store distribution (each lane + # owns one hidden element across the 64 subtile tokens); no reorder needed. + return casted @cute.jit @@ -2026,12 +2672,11 @@ def fc2_ublk_tmem_acc_load(*, tmem_subtile_tensor: cute.Tensor, epi, **_): @cute.jit def fc2_stg_post_f2fp_reorder( - *, - casted: cute.Tensor, # (subtile_cnt,) - fc2_output_dtype: Type[cutlass.Numeric], - tmem_subtile_view: cute.Tensor, # (epi_tile_m, epi_tile_n) - **_, -): + *, + casted: cute.Tensor, # (subtile_cnt,) + tmem_subtile_view: cute.Tensor, # (epi_tile_m, epi_tile_n) + **_): + if cutlass.const_expr(cute.size(casted) != 64): raise NotImplementedError( "fc2 stg pass expects 64 fp32 regs in total before store reorder.") @@ -2042,19 +2687,17 @@ def fc2_stg_post_f2fp_reorder( # read casted through (t, hidden, half) -> casted[t*16 + hidden + half*32] # in (t fastest) order -> [top0,bot0,top1,bot1,...] per half = packed bf16x2. # scatter: de-interleave the transposed natural-hidden regs back to the - # (token, hidden) pre-store order declared by _Fc2StgSubtilePreStoreContract. + # STG pre-store order (token = lane + 32*(vid//32), hidden = vid % 32). gather_top_bot_map = ((2, 16, 2), (16, 1, 32)) scatter_top_bot_map = ((16, 2, 2), (2, 1, 32)) - dtype = fc2_output_dtype + dtype = cutlass.BFloat16 packed = cute.make_rmem_tensor((64, ), dtype) cute.autovec_copy( cute.composition( casted, cute.make_layout((gather_top_bot_map[0], ), - stride=(gather_top_bot_map[1], ))), - packed, - ) + stride=(gather_top_bot_map[1], ))), packed) # Although this works... # packed.store( # cute.make_tensor( @@ -2065,62 +2708,42 @@ def fc2_stg_post_f2fp_reorder( packed_i32 = cute.recast_tensor(packed, cutlass.Float32) # (32,): 16 i32 per half - token_0_32_pre_scatter_back = TmemTranspose16x32Packed( + # Reuse the 32-bit transpose: each i32 slot carries one packed bf16x2 pair. + token_0_32_pre_scatter_back = TmemTranspose16x32( tmem_subtile_view.iterator, Region.Top, - reg_tensor=TensorWithContract( - tensor=cute.composition(packed_i32, (16, )), - contract=TmemTranspose16x32Packed.InputContract, - ), + reg_tensor=cute.composition(packed_i32, (16, )), ).from_r1_perm_until_last_store() - token_32_64_pre_scatter_back = TmemTranspose16x32Packed( + token_32_64_pre_scatter_back = TmemTranspose16x32( tmem_subtile_view.iterator + 32, Region.Top, - reg_tensor=TensorWithContract( - tensor=cute.composition(cute.domain_offset(16, packed_i32), (16, )), - contract=TmemTranspose16x32Packed.InputContract, - ), + reg_tensor=cute.composition(cute.domain_offset(16, packed_i32), (16, )), ).from_r1_perm_until_last_store() - cute.autovec_copy(token_0_32_pre_scatter_back.tensor, + cute.autovec_copy(token_0_32_pre_scatter_back, cute.zipped_divide(packed_i32, (16, ))[None, 0]) - cute.autovec_copy(token_32_64_pre_scatter_back.tensor, + cute.autovec_copy(token_32_64_pre_scatter_back, cute.zipped_divide(packed_i32, (16, ))[None, 1]) out = cute.make_rmem_tensor((64, ), dtype) cute.autovec_copy( cute.composition( packed, cute.make_layout((scatter_top_bot_map[0], ), - stride=(scatter_top_bot_map[1], ))), - out, - ) - return TensorWithContract(tensor=out, - contract=_Fc2StgSubtilePreStoreContract) - - -# (...) -> ((atom_v, 1)) -@cute.jit -def wrap_into_copy_standard_layout(tensor: cute.Tensor): - tensor = cute.coalesce(cute.flatten(tensor)) - tensor = cute.append_ones(tensor, cute.rank(tensor) + 1) - tensor = cute.group_modes(tensor, 0, cute.rank(tensor) - 1) - tensor = cute.group_modes(tensor, 0, cute.rank(tensor)) - return tensor + stride=(scatter_top_bot_map[1], ))), out) + return out @cute.jit def fc2_redg_post_f2fp_reorder( *, casted: cute.Tensor, - fc2_output_dtype: Type[cutlass.Numeric], tmem_subtile_view: cute.Tensor, **_, ): # (epi_tid, elem_idx) -> (token_64, hidden_128), each thread hold token_2 x hidden_32 natural = fc2_stg_post_f2fp_reorder( casted=casted, - fc2_output_dtype=fc2_output_dtype, tmem_subtile_view=tmem_subtile_view, - ).tensor + ) core_matrix_reorder_sttm_atom = cute.make_copy_atom( tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32, @@ -2151,90 +2774,108 @@ def fc2_redg_post_f2fp_reorder( wrap_into_copy_standard_layout(current_sttm_src), wrap_into_copy_standard_layout( tmem_subtile_divided_by_token_group_divided_by_16dp[None, None, - None, i]), - ) + None, i])) cute.copy( core_matrix_reorder_ldtm_atom, wrap_into_copy_standard_layout( tmem_subtile_divided_by_token_group_divided_by_16dp[None, None, 0, i]), - wrap_into_copy_standard_layout(out_as_i32[(None, 0), i]), - ) + wrap_into_copy_standard_layout(out_as_i32[(None, 0), i])) cute.copy( core_matrix_reorder_ldtm_atom, wrap_into_copy_standard_layout( tmem_subtile_divided_by_token_group_divided_by_16dp[None, None, 1, i]), - wrap_into_copy_standard_layout(out_as_i32[(None, 1), i]), - ) + wrap_into_copy_standard_layout(out_as_i32[(None, 1), i])) - return TensorWithContract(tensor=cute.coalesce(out), - contract=_Fc2RedgSubtilePreStoreContract) + return cute.coalesce(out) @cute.jit def fc2_stg_store_function( - *, - epi, - subtile: TensorWithContract, - subtile_idx: cutlass.Int32, - fc2_output_router: Fc2OutputRouter, - **_, -): - assert_contract_equivalent( - subtile.contract, - _Fc2StgSubtilePreStoreContract, - context="fc2 STG store input", - ) - copy_atom_256b = cute.make_copy_atom( + *, + epi, + subtile: cute.Tensor, # Always bf16 pre quant tesnor + subtile_idx: cutlass.Int32, + fc2_output_router: Fc2OutputRouter, + **_): + if cutlass.const_expr(epi.combine_format.is_quantized): + data_subtile, sf_regs = QuantImpl(epi.combine_format, + "regs_in_thread")(subtile) + else: + data_subtile = subtile + sf_regs = None + stg_width_elems: cutlass.Constexpr[int] = min( + 32, 256 // data_subtile.element_type.width) + stg_bits: cutlass.Constexpr[ + int] = stg_width_elems * data_subtile.element_type.width + copy_atom_vec = cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), - subtile.tensor.element_type, - num_bits_per_copy=256, + cutlass.Int32, + num_bits_per_copy=stg_bits, ) - stg_width_elems: cutlass.Constexpr[ - int] = 256 // subtile.tensor.element_type.width - elem_axis: cutlass.Constexpr[int] = subtile.contract.domain.names.index( - "vid") - elems_per_thread: cutlass.Constexpr[int] = subtile.contract.domain.sizes[ - elem_axis] + elems_per_thread: cutlass.Constexpr[int] = cute.size(data_subtile) if cutlass.const_expr(elems_per_thread % stg_width_elems != 0): raise ValueError( "fc2 STG store requires pre-store elems per thread to be divisible " f"by STG issue width, got {elems_per_thread} and {stg_width_elems}." ) + + if cutlass.const_expr(sf_regs is not None): + sf_scales_per_stg: cutlass.Constexpr[ + int] = 32 // epi.combine_format.scale_block + token_groups_per_subtile: cutlass.Constexpr[ + int] = epi._EpilogueTokenTileSize // 32 + for token_group in cutlass.range_constexpr(token_groups_per_subtile): + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32( + token_groups_per_subtile) + token_group + sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) + if sf_pred != cutlass.Int32(0): + sf_dst = cute.make_tensor( + sf_ptr, cute.make_layout((sf_scales_per_stg, ))) + for k in cutlass.range_constexpr(sf_scales_per_stg): + sf_dst[k] = sf_regs[token_group * sf_scales_per_stg + k] + iters_per_subtile: cutlass.Constexpr[ int] = elems_per_thread // stg_width_elems - copy_src = cute.zipped_divide(subtile.tensor, (stg_width_elems, )) - copy_pairs = fc2_output_router.resolve( - copy_src, - iters_per_subtile, - subtile_idx, - ) + copy_src = cute.zipped_divide(data_subtile, (stg_width_elems, )) + single_copy_layout = cute.make_layout(((stg_width_elems, 1), ), + stride=((1, 0), )) subtile_iter_base = cutlass.Int32(subtile_idx) * cutlass.Int32( iters_per_subtile) for local_iter in cutlass.range_constexpr(iters_per_subtile): global_iter = subtile_iter_base + cutlass.Int32(local_iter) - if fc2_output_router.valid[global_iter] != cutlass.Int32(0): - copy_src_i, copy_dst_i = copy_pairs[local_iter] - cute.copy(copy_atom_256b, copy_src_i, copy_dst_i) + dst_ptr, pred = fc2_output_router.get_data_dst(global_iter) + if pred != cutlass.Int32(0): + src_i = cute.make_tensor(copy_src[None, local_iter].iterator, + single_copy_layout) + dst_i = cute.make_tensor(dst_ptr, single_copy_layout) + cute.copy(copy_atom_vec, cute.recast_tensor(src_i, cutlass.Int32), + cute.recast_tensor(dst_i, cutlass.Int32)) @cute.jit def fc2_ublk_store_function_impl( *, epi, - subtile: TensorWithContract, + subtile: cute.Tensor, # Always bf16 pre-quant tensor subtile_idx: cutlass.Int32, fc2_output_router: Fc2OutputRouter, + **_, ): - assert_contract_equivalent( - subtile.contract, - epi.process_pipeline.pre_store_contract, - context="fc2 UBLK store input", - ) smem_tensor = epi.smem_tensor if cutlass.const_expr(smem_tensor is None): raise ValueError("fc2 UBLK store requires epi.smem_tensor.") + quantized: cutlass.Constexpr[bool] = epi.combine_format.is_quantized + if cutlass.const_expr(quantized): + data_subtile, selected_sf = QuantImpl( + epi.combine_format, + "threads_with_the_same_reg", + lane_idx=epi.lane_idx, + )(subtile) + else: + data_subtile = subtile + selected_sf = None smem_read_write_bar = pipeline.NamedBarrier( barrier_id=SwapABSwigluFp4Epilogue._EpilogueSyncWaitBarId, @@ -2244,10 +2885,7 @@ def fc2_ublk_store_function_impl( lane_idx = epi.lane_idx warp_hidden_base = cutlass.Int32(warp_idx * 32) - vid_axis: cutlass.Constexpr[int] = subtile.contract.domain.names.index( - "vid") - regs_per_thread: cutlass.Constexpr[int] = subtile.contract.domain.sizes[ - vid_axis] + regs_per_thread: cutlass.Constexpr[int] = cute.size(data_subtile) tokens_per_smem_slice: cutlass.Constexpr[int] = cute.size(smem_tensor, mode=[0]) if cutlass.const_expr(regs_per_thread % tokens_per_smem_slice != 0): @@ -2257,55 +2895,59 @@ def fc2_ublk_store_function_impl( ) loop_cnt: cutlass.Constexpr[int] = regs_per_thread // tokens_per_smem_slice - for loop_idx in cutlass.range_constexpr(loop_cnt): - if cutlass.const_expr(loop_idx > 0): + # SF straight out (no smem): the sf_vec lanes of a block share the CREDUX + # result, so each emits its slice -- one scale per slot. The sf mapping owns + # the (lane, slot) -> (token, hidden) layout; selected_sf[slot] is this + # lane's slot-th scale, aligned to sf_iter = subtile_idx*iters_per_subtile+slot. + if cutlass.const_expr(quantized): + # A scale block = sf_vec hidden = sf_vec lanes (UBLK), and those lanes + # split the subtile's tokens, so each emits subtile_tokens // lanes_per_block. + lanes_per_block: cutlass.Constexpr[int] = epi.combine_format.scale_block + iters_per_subtile: cutlass.Constexpr[ + int] = epi._EpilogueTokenTileSize // lanes_per_block + for slot in cutlass.range_constexpr(iters_per_subtile): + sf_iter = cutlass.Int32(subtile_idx) * cutlass.Int32( + iters_per_subtile) + slot + sf_ptr, sf_pred = fc2_output_router.get_sf_dst(sf_iter) + if sf_pred != cutlass.Int32(0): + sf_dst = cute.make_tensor(sf_ptr, cute.make_layout((1, ))) + sf_dst[0] = selected_sf[slot] + + for token32_group_idx in cutlass.range_constexpr(loop_cnt): + if cutlass.const_expr(token32_group_idx > 0): cute.arch.cp_async_bulk_wait_group(0, read=True) - cute.arch.sync_warp() smem_read_write_bar.arrive_and_wait() - # R2S: materialize this loop's token slice into the fixed - # (token_rows, hidden_128) scratch tile. The pre-store contract says - # each lane owns one hidden column over all subtile tokens, so loop_idx - # simply selects the next contiguous token_rows chunk from RMEM. + # R2S transpose: each lane writes its hidden column's 32 token rows for + # this group. (SF already went straight out above; only DATA transposes + # through smem here.) for token_i in cutlass.range_constexpr(tokens_per_smem_slice): - src_reg = token_i + tokens_per_smem_slice * loop_idx + src_reg = token_i + tokens_per_smem_slice * token32_group_idx smem_tensor[token_i, - warp_hidden_base + lane_idx] = subtile.tensor[src_reg] + warp_hidden_base + lane_idx] = data_subtile[src_reg] cute.arch.fence_proxy("async.shared", space="cta") - cute.arch.sync_warp() smem_read_write_bar.arrive_and_wait() - iter_idx = subtile_idx // cutlass.Int32(2) - store_coord = eval_function_mapping( - fc2_output_router.output_mapping, - epi_tid=epi.tidx, - iter_idx=iter_idx, - ) - token_in_cta_tile = cutlass.Int32(store_coord["token_in_cta_tile"]) - slice_token_start = subtile_idx * cutlass.Int32( - epi._EpilogueTokenTileSize) + cutlass.Int32( - loop_idx * tokens_per_smem_slice) - slice_token_end = slice_token_start + cutlass.Int32( - tokens_per_smem_slice) - - if (fc2_output_router.valid[iter_idx] != cutlass.Int32(0) - and token_in_cta_tile >= slice_token_start - and token_in_cta_tile < slice_token_end): - scratch_row = token_in_cta_tile - slice_token_start + # ublk_iter_idx is constexpr so get_dst indexes a constexpr slot (no spill). + # Gate / scratch_row specialize the store-out mapping: lane_idx//8 picks the + # subtile, warp_idx*8 + lane_idx%8 is the token's row within the 32-group. + ublk_iter_idx = token32_group_idx + # SF already went straight out above; here only the DATA row is pushed. + dst_ptr, pred = fc2_output_router.get_data_dst(ublk_iter_idx) + if pred != cutlass.Int32(0) and (lane_idx // + cutlass.Int32(8)) == subtile_idx: + scratch_row = warp_idx * cutlass.Int32( + 8) + lane_idx % cutlass.Int32(8) copy_elems = cutlass.Int32(128) if cutlass.const_expr(epi.fc2_hidden_needs_predicate): copy_elems = cutlass.Int32( fc2_output_router.valid_hidden_this_cta_tile) - copy_bytes = copy_elems * epi.fc2_output_dtype.width // 8 + # smem_tensor holds the combine wire dtype, so the bulk byte count + # scales with the wire width, not the bf16 compute dtype. + copy_bytes = copy_elems * epi.combine_format.act_dtype.width // 8 src_row = cute.slice_(smem_tensor, (scratch_row, None)) - dst_ptr = cute.make_ptr( - src_row.element_type, - fc2_output_router.dst_ptrs[iter_idx], - AddressSpace.gmem, - assumed_align=16, - ) if cutlass.const_expr(epi.reduce_topk_in_kernel): _cp_reduce_async_bulk_add_noftz_bf16_s2g( dst_ptr, @@ -2331,21 +2973,13 @@ def fc2_ublk_store_function_impl( def fc2_redg_store_function( *, epi, - subtile: TensorWithContract, + subtile: cute.Tensor, # Always bf16; in-kernel reduce never quantizes subtile_idx: cutlass.Int32, fc2_output_router: Fc2OutputRouter, **_, ): - assert_contract_equivalent( - subtile.contract, - _Fc2RedgSubtilePreStoreContract, - context="fc2 REDG store input", - ) redg_width_elems: cutlass.Constexpr[int] = 4 - elem_axis: cutlass.Constexpr[int] = subtile.contract.domain.names.index( - "elem_idx") - elems_per_thread: cutlass.Constexpr[int] = subtile.contract.domain.sizes[ - elem_axis] + elems_per_thread: cutlass.Constexpr[int] = cute.size(subtile) if cutlass.const_expr(elems_per_thread % redg_width_elems != 0): raise ValueError( "fc2 REDG store requires pre-store elems per thread to be divisible " @@ -2355,20 +2989,14 @@ def fc2_redg_store_function( int] = elems_per_thread // redg_width_elems subtile_iter_base = cutlass.Int32(subtile_idx) * cutlass.Int32( iters_per_subtile) - subtile_by_redg_issue = cute.zipped_divide(subtile.tensor, - (redg_width_elems, )) + subtile_by_redg_issue = cute.zipped_divide(subtile, (redg_width_elems, )) for local_iter in cutlass.range_constexpr(iters_per_subtile): global_iter = subtile_iter_base + cutlass.Int32(local_iter) - if fc2_output_router.valid[global_iter] != cutlass.Int32(0): + dst_ptr, pred = fc2_output_router.get_data_dst(global_iter) + if pred != cutlass.Int32(0): bf16x4 = subtile_by_redg_issue[None, local_iter] packed_bf16x2 = cute.recast_tensor(bf16x4, cutlass.Float32) - dst_ptr = cute.make_ptr( - fc2_output_router.base_output.element_type, - fc2_output_router.dst_ptrs[global_iter], - AddressSpace.gmem, - assumed_align=8, - ) _red_add_relaxed_sys_v2_bf16x2( dst_ptr, cutlass.Float32(packed_bf16x2[0]), @@ -2378,12 +3006,12 @@ def fc2_redg_store_function( def make_fc2_stg_process_pipeline( *, - fc2_output_dtype: Type[cutlass.Numeric], + combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int, ) -> Fc2ProcessPipeline: - store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_contract( - fc2_output_dtype, + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_stg_cta_store_out_contract( + combine_format, cta_token_tile_size, cta_hidden_tile_size, ) @@ -2392,21 +3020,21 @@ def make_fc2_stg_process_pipeline( f2fp=fc2_f2fp, post_f2fp_reorder=fc2_stg_post_f2fp_reorder, store_function=fc2_stg_store_function, - pre_store_contract=_Fc2StgSubtilePreStoreContract, fc2_cta_tile_contract=fundamental_mapping, store_out_mapping=store_out_mapping, + sf_store_out_mapping=sf_store_out_mapping, require_tmem_trans=True, ) def make_fc2_redg_process_pipeline( *, - fc2_output_dtype: Type[cutlass.Numeric], + combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int, ) -> Fc2ProcessPipeline: - store_out_mapping, fundamental_mapping = make_fc2_redg_cta_store_out_contract( - fc2_output_dtype, + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_redg_cta_store_out_contract( + combine_format, cta_token_tile_size, cta_hidden_tile_size, ) @@ -2415,21 +3043,21 @@ def make_fc2_redg_process_pipeline( f2fp=fc2_f2fp, post_f2fp_reorder=fc2_redg_post_f2fp_reorder, store_function=fc2_redg_store_function, - pre_store_contract=_Fc2RedgSubtilePreStoreContract, fc2_cta_tile_contract=fundamental_mapping, store_out_mapping=store_out_mapping, + sf_store_out_mapping=sf_store_out_mapping, require_tmem_trans=True, ) def make_fc2_ublk_process_pipeline( *, - fc2_output_dtype: Type[cutlass.Numeric], + combine_format: CombineFormat, cta_token_tile_size: int, cta_hidden_tile_size: int, ) -> Fc2ProcessPipeline: - store_out_mapping, fundamental_mapping = make_fc2_ublk_store_out_contract( - fc2_output_dtype, + store_out_mapping, sf_store_out_mapping, fundamental_mapping = make_fc2_ublk_store_out_contract( + combine_format, cta_token_tile_size, cta_hidden_tile_size, ) @@ -2438,306 +3066,8 @@ def make_fc2_ublk_process_pipeline( f2fp=fc2_f2fp, post_f2fp_reorder=post_f2fp_reorder_identity, store_function=fc2_ublk_store_function_impl, - pre_store_contract=_Fc2UblkSubtilePreStoreContract, fc2_cta_tile_contract=fundamental_mapping, store_out_mapping=store_out_mapping, + sf_store_out_mapping=sf_store_out_mapping, require_tmem_trans=False, ) - - -# Device only object -class SwapABFc2Epilogue: - - def __init__( - self, - base: SwapABSwigluFp4Epilogue, - tidx: cutlass.Int32, - epi_smem_storage, - fc2_output: cute.Tensor, # MoE domain (token, topk, hidden) - token_comm_args: TokenCommArgs, - optional_epi_args: NvFp4OptinalEpiArgs, - ): - self.base = base - self.tidx = tidx % (base._EpilogueWarpCnt * 32) - self.warp_idx = self.tidx // 32 - self.lane_idx = self.tidx % 32 - self.fc2_output = fc2_output - self.token_comm_args = token_comm_args - self.optional_epi_args = optional_epi_args - if cutlass.const_expr(base.fc2_use_bulk): - fc2_smem_rows = (base.epi_smem_bytes * 8 // - base._EpilogueFc2HiddenTileSize // - base.fc2_output_dtype.width) - if cutlass.const_expr(fc2_smem_rows != 32): - raise NotImplementedError( - "Remember to adjust fc2 smem structure if switch to non-bf16 combine." - ) - self.smem_tensor = cute.make_tensor( - cute.recast_ptr( - epi_smem_storage.epi_smem.data_ptr(), - dtype=base.fc2_output_dtype, - ), - cute.make_layout( - (fc2_smem_rows, base._EpilogueFc2HiddenTileSize), - stride=(base._EpilogueFc2HiddenTileSize, 1), - ), - ) - self.process_pipeline = make_fc2_ublk_process_pipeline( - fc2_output_dtype=base.fc2_output_dtype, - cta_token_tile_size=base.cta_tile_n, - cta_hidden_tile_size=base.cta_tile_m, - ) - else: - self.smem_tensor = None - if cutlass.const_expr(base.reduce_topk_in_kernel): - self.process_pipeline = make_fc2_redg_process_pipeline( - fc2_output_dtype=base.fc2_output_dtype, - cta_token_tile_size=base.cta_tile_n, - cta_hidden_tile_size=base.cta_tile_m, - ) - else: - self.process_pipeline = make_fc2_stg_process_pipeline( - fc2_output_dtype=base.fc2_output_dtype, - cta_token_tile_size=base.cta_tile_n, - cta_hidden_tile_size=base.cta_tile_m, - ) - - def __getattr__(self, name): - return getattr(object.__getattribute__(self, "base"), name) - - def __extract_mlir_values__(self) -> List[ir.Value]: - # See SwapABFc1Epilogue.__extract_mlir_values__: this helper carries - # only loop-invariant Python context. It intentionally serializes no - # MLIR values, so changing it to store loop-carried state would be a - # correctness bug. - return [] - - def __new_from_mlir_values__(self, - values: List[ir.Value]) -> "SwapABFc2Epilogue": - assert len(values) == 0 - return self - - @cute.jit - def signal_fc2_done(self, work_tile_info): - if cutlass.const_expr(self.token_back_by_dispatch): - if self.tidx == 0: - _red_add_release_gpu_s32( - self.token_comm_args.fc2_done_counter.iterator + - work_tile_info.expert_idx, - cutlass.Int32(1), - ) - - @cute.jit - def _make_output_router( - self, - work_tile_info: MoEWorkTileInfo, - ) -> Fc2OutputRouter: - task_tile_data_row_start = ( - work_tile_info.cumulative_data_physical_row + - work_tile_info.tile_n_idx * cutlass.Int32(self.cta_tile_n)) - hidden_base_this_cta_tile = work_tile_info.tile_m_idx * cutlass.Int32( - self.cta_tile_m) - valid_hidden_this_cta_tile = (cutlass.Int32(self.fc2_output.shape[2]) - - hidden_base_this_cta_tile) - if valid_hidden_this_cta_tile < 0: - valid_hidden_this_cta_tile = 0 - if valid_hidden_this_cta_tile > self._EpilogueFc2HiddenTileSize: - valid_hidden_this_cta_tile = self._EpilogueFc2HiddenTileSize - - metadata_u32 = None - peer_rank_ptr_mapper = None - direct_token_base_this_cta_tile = task_tile_data_row_start - if cutlass.const_expr(self.token_comm_args is not None - and not self.token_back_by_dispatch): - metadata_u32 = cute.domain_offset( - (task_tile_data_row_start, 0), - cute.recast_tensor( - self.token_comm_args.token_src_metadata, - cutlass.Uint32, - ), - ) - peer_rank_ptr_mapper = self.token_comm_args.peer_rank_ptr_mapper - direct_token_base_this_cta_tile = None - - return Fc2OutputRouter( - metadata=metadata_u32, - direct_token_base_this_cta_tile=direct_token_base_this_cta_tile, - base_output=self.fc2_output, - hidden_base_this_cta_tile=hidden_base_this_cta_tile, - peer_rank_ptr_mapper=peer_rank_ptr_mapper, - valid_tokens_this_cta_tile=work_tile_info.valid_tokens_in_tile, - valid_hidden_this_cta_tile=valid_hidden_this_cta_tile, - reduce_topk_in_kernel=self.reduce_topk_in_kernel, - output_mapping=self.process_pipeline.store_out_mapping, - epi_tid=self.tidx, - ).prefetch() - - @cute.jit - def __call__( - self, - work_tile_info: MoEWorkTileInfo, - tmem_acc_tensor: cute.Tensor, - acc_pipeline, - acc_consumer_state, - is_odd_turn: cutlass.Int32, - ): - # subtile-irrelevant hoist: fc2 alpha scales raw fc2 accumulators before f2fp. - if cutlass.const_expr(self.optional_epi_args.fc2_alpha is not None): - alpha_val = self.optional_epi_args.fc2_alpha[ - work_tile_info.expert_idx] - else: - alpha_val = None - acc_ready = False - if not work_tile_info.peek_ready: - acc_ready = True - acc_pipeline.consumer_wait(acc_consumer_state) - fc2_output_router = self._make_output_router(work_tile_info) - # (cta_tile_m, cta_tile_n) -> (epi_tile_m, epi_tile_n, iters) - tmem_acc_tensor_tiled_by_epi_tile = cute.flat_divide( - tmem_acc_tensor, (self._EpilogueFc2HiddenTileSize, - self._EpilogueTokenTileSize))[None, None, 0, None] - - acc_pipeline.consumer_wait(acc_consumer_state, acc_ready) - iket.range_push("fc2_epi") - valid_tokens = work_tile_info.valid_tokens_in_tile - - # Overlap path preloads two subtiles before releasing acc TMEM. - unroll_tile_cnt = (2 if cutlass.const_expr( - self.overlapping_accum and self.process_pipeline.require_tmem_trans) - else 0) - remain_subtile_cnt = self.subtile_cnt - unroll_tile_cnt - - if cutlass.const_expr(unroll_tile_cnt > 0): - subtile_idx_first = (cutlass.Int32(self.subtile_cnt) - - is_odd_turn) % cutlass.Int32(self.subtile_cnt) - subtile_idx_second = (cutlass.Int32(self.subtile_cnt + 1) - - is_odd_turn) % cutlass.Int32(self.subtile_cnt) - - # preload_subtile_first: subtile_idx_first's raw PRE-transpose acc, LDTM'd by - # all 128 epi threads into 4 reg tensors == the 4 quadrants of the subtile's - # (128 tmem_dp x 64 tmem_col) footprint. Only these raw-TMEM offsets are - # guaranteed: - # reg[0]/reg[1], reg[2]/reg[3] : top vs bot -> 16 apart in tmem_dp - # reg[0]/reg[2], reg[1]/reg[3] : 1st vs 2nd half -> 32 apart in tmem_col - # (so reg[0..1] = the first 128x32, reg[2..3] = the second 128x32 of the 128x64.) - # The per-lane (lane_idx, elem_idx) -> (tmem_dp, tmem_col) layout INSIDE each - # reg tensor is opaque -- do not assume it; it only becomes well-defined once - # the tmem transpose consumes them. - preload_subtile_first: Tuple[ - cute.Tensor, cute.Tensor, cute.Tensor, - cute.Tensor] = (_TmemTranspose16x32Core.load_subtile_raw_acc( - tmem_acc_tensor_tiled_by_epi_tile[None, None, - subtile_idx_first])) - - # Release acc to next MMA unconditionally. - cute.arch.fence_view_async_tmem_load() - acc_pipeline.consumer_release(acc_consumer_state) - - # preload_subtile_second: same 128 tmem_dp x 64 tmem_col footprint, but for - # subtile_idx_second (the other token subtile, not the 2nd col-half). Same - # quadrant/offset invariants and opaque per-lane layout as above. - preload_subtile_second: Tuple[ - cute.Tensor, cute.Tensor, cute.Tensor, - cute.Tensor] = (_TmemTranspose16x32Core.load_subtile_raw_acc( - tmem_acc_tensor_tiled_by_epi_tile[None, None, - subtile_idx_second])) - - # Both unrolled subtiles borrow tmem_subtile_second as workspace. - preload_pair = (preload_subtile_first, preload_subtile_second) - subtile_idx_pair = (subtile_idx_first, subtile_idx_second) - for i in cutlass.range_constexpr(unroll_tile_cnt): - if subtile_idx_pair[i] * cutlass.Int32( - self._EpilogueTokenTileSize) < valid_tokens: - self.run_subtile( - subtile_idx=subtile_idx_pair[i], - tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[ - None, None, subtile_idx_second], - preload_acc=preload_pair[i], - fc2_output_router=fc2_output_router, - alpha_val=alpha_val, - release_after_ldtm=False, - acc_pipeline=acc_pipeline, - acc_consumer_state=acc_consumer_state, - ) - - if cutlass.const_expr(self.overlapping_accum and unroll_tile_cnt == 0): - release_after_ldtm = True - else: - release_after_ldtm = False - for i in cutlass.range(remain_subtile_cnt, unroll=1): - # for i in cutlass.range_constexpr(remain_subtile_cnt): - real_i = i + unroll_tile_cnt - if cutlass.const_expr(self.overlapping_accum): - subtile_idx = (cutlass.Int32(real_i + self.subtile_cnt) - - is_odd_turn) % cutlass.Int32(self.subtile_cnt) - else: - subtile_idx = cutlass.Int32(real_i) - - if subtile_idx * cutlass.Int32( - self._EpilogueTokenTileSize) < valid_tokens: - self.run_subtile( - subtile_idx=subtile_idx, - tmem_subtile_tensor=tmem_acc_tensor_tiled_by_epi_tile[ - None, None, subtile_idx], - preload_acc=None, - fc2_output_router=fc2_output_router, - alpha_val=alpha_val, - release_after_ldtm=release_after_ldtm, - acc_pipeline=acc_pipeline, - acc_consumer_state=acc_consumer_state, - ) - release_after_ldtm = False - - # Non-overlap-path release: at the natural task-tile boundary. - if cutlass.const_expr(not self.overlapping_accum): - cute.arch.fence_view_async_tmem_load() - acc_pipeline.consumer_release(acc_consumer_state) - - @cute.jit - def run_subtile( - self, - subtile_idx: cutlass.Int32, - # (hidden_tile, token_subtile), fundamentally (epi_tile_m, epi_tile_n) - tmem_subtile_tensor: cute.Tensor, - preload_acc: Optional[Tuple[cute.Tensor, cute.Tensor, cute.Tensor, - cute.Tensor]], - fc2_output_router: Fc2OutputRouter, - alpha_val: Optional[cutlass.Float32], - release_after_ldtm: Union[cutlass.Boolean, bool], - acc_pipeline, - acc_consumer_state, - ): - process_pipeline = self.process_pipeline - if cutlass.const_expr(preload_acc is None): - loaded = process_pipeline.tmem_acc_load( - tmem_subtile_tensor=tmem_subtile_tensor, - epi=self, - ) - if release_after_ldtm: - cute.arch.fence_view_async_tmem_load() - acc_pipeline.consumer_release(acc_consumer_state) - else: - loaded = preload_acc - - casted = process_pipeline.f2fp( - *loaded, - fc2_output_dtype=self.fc2_output_dtype, - alpha_val=alpha_val, - ) - pre_store = process_pipeline.post_f2fp_reorder( - casted=casted, - contract=process_pipeline.pre_store_contract, - fc2_output_dtype=self.fc2_output_dtype, - tmem_subtile_view=tmem_subtile_tensor, - ) - assert_contract_equivalent( - pre_store.contract, - process_pipeline.pre_store_contract, - context="fc2 process pipeline pre-store", - ) - process_pipeline.store_function( - epi=self, - subtile=pre_store, - subtile_idx=subtile_idx, - fc2_output_router=fc2_output_router, - ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py index 4602434b292f..09412f1690a3 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/fc1_fc2_fuse_sched.py @@ -1,5 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause """Fused fc1 + fc2 MegaMoE scheduler.""" from enum import IntEnum @@ -657,19 +659,31 @@ def create( num_fc2_hidden_blocks = (hidden + params.cluster_tile_n - 1) // params.cluster_tile_n - # current_work init must use ext.WorkTileInfo (8 fields) to match the - # shape that gen_next_work writes; otherwise MLIR serialization slot - # Scheduler emits the final 8-field work tile directly. - current_work = ext.WorkTileInfo( - expert_idx=Int32(WorkTileState.DONE), - tile_m_idx=Int32(0), - tile_n_idx=Int32(0), - cumulative_data_physical_row=Int32(0), - cumulative_sf_physical_row=Int32(0), - cumulative_token_block_count=Int32(0), - valid_tokens_in_tile=Int32(0), - phase_and_peek=Int32(BlockPhase.None_), - ) + # current_work init must use ext.WorkTileInfo to match the shape that + # gen_next_work writes; otherwise MLIR serialization slots would differ. + if const_expr(params.is_swap_ab): + current_work = ext.WorkTileInfo( + expert_idx=Int32(WorkTileState.DONE), + tile_m_idx=Int32(0), + tile_n_idx=Int32(0), + cumulative_data_physical_row=Int32(0), + cumulative_sf_physical_row=Int32(0), + cumulative_token_block_count=Int32(0), + valid_tokens_in_cta_tile=Int32(0), + phase_and_peek=Int32(BlockPhase.None_), + ) + else: + current_work = ext.WorkTileInfo( + expert_idx=Int32(WorkTileState.DONE), + tile_m_idx=Int32(0), + tile_n_idx=Int32(0), + cumulative_data_physical_row=Int32(0), + cumulative_sf_physical_row=Int32(0), + cumulative_token_block_count=Int32(0), + valid_tokens_in_cta_cluster_tile=Int32(0), + phase_and_peek=Int32(BlockPhase.None_), + fc1_counter_index=Int32(0), + ) sched_producer_group = pipeline.CooperativeGroup( pipeline.Agent.Thread, 32) @@ -855,9 +869,7 @@ def _advance_work_linear_tile_idx_dynamic( result across the sched warp. Cluster-internal protocol (DSMEM broadcast + cluster_pipeline mbar - wait) is identical between the two paths. Mirrors - ``cute_dsl_kernel_library/dsl_kernels/moe/moe_persistent_scheduler.py`` - ``_fetch_next_cluster_idx`` (lines 838-894). + wait) is identical between the two paths. """ ds = self._dynamic_state cluster_pipeline = self._cluster_pipeline @@ -1176,14 +1188,14 @@ def _decode_inside_expert( cluster_intermediate_or_hidden_block_idx * params.cluster_shape_mn[1] + self.cta_id_in_cluster[1]) - # valid_tokens_in_tile: clip cta_tile_m tokens at the current expert + # valid_tokens_in_cta_tile: clip cta_tile_m tokens at the current expert # right boundary. token_idx_start_in_expert = cta_token_block_idx * Int32(cta_tile_m) remaining_in_expert = (state.current_this_expert_token_cnt - token_idx_start_in_expert) remaining_in_expert = cutlass.max(remaining_in_expert, Int32(0)) - valid_tokens_in_tile = cutlass.min(remaining_in_expert, - Int32(cta_tile_m)) + valid_tokens_in_cta_tile = cutlass.min(remaining_in_expert, + Int32(cta_tile_m)) # Swap scheduler-internal M/N back to GEMM-domain M/N on output. if const_expr(params.is_swap_ab): @@ -1194,16 +1206,41 @@ def _decode_inside_expert( tile_n_idx = cta_intermediate_or_hidden_block_idx # ext.enrich_work_tile_info may OR the peek bit into phase_and_peek. - return self._ext.WorkTileInfo( - expert_idx=state.current_expert_idx, - tile_m_idx=tile_m_idx, - tile_n_idx=tile_n_idx, - cumulative_data_physical_row=state.current_data_cumul, - cumulative_sf_physical_row=state.current_sf_cumul, - cumulative_token_block_count=state.current_token_block_cumul, - valid_tokens_in_tile=valid_tokens_in_tile, - phase_and_peek=state.current_phase, - ) + if const_expr(params.is_swap_ab): + return self._ext.WorkTileInfo( + expert_idx=state.current_expert_idx, + tile_m_idx=tile_m_idx, + tile_n_idx=tile_n_idx, + cumulative_data_physical_row=state.current_data_cumul, + cumulative_sf_physical_row=state.current_sf_cumul, + cumulative_token_block_count=state.current_token_block_cumul, + valid_tokens_in_cta_tile=valid_tokens_in_cta_tile, + phase_and_peek=state.current_phase, + ) + else: + fc1_counter_index = cluster_token_block_idx + cluster_tile_m = params.cluster_shape_mn[0] * cta_tile_m + cluster_start = cluster_token_block_idx * Int32(cluster_tile_m) + remaining_cluster = cutlass.max( + state.current_this_expert_token_cnt - cluster_start, Int32(0)) + valid_tokens_in_cluster_tile = cutlass.min(remaining_cluster, + Int32(cluster_tile_m)) + # Pack: high 16b = per-CTA tile count, low 16b = cluster-level count. + valid_tokens_in_cta_cluster_tile = ( + (valid_tokens_in_cta_tile << Int32(16)) + | valid_tokens_in_cluster_tile) + return self._ext.WorkTileInfo( + expert_idx=state.current_expert_idx, + tile_m_idx=tile_m_idx, + tile_n_idx=tile_n_idx, + cumulative_data_physical_row=state.current_data_cumul, + cumulative_sf_physical_row=state.current_sf_cumul, + cumulative_token_block_count=state.current_token_block_cumul, + valid_tokens_in_cta_cluster_tile= + valid_tokens_in_cta_cluster_tile, + phase_and_peek=state.current_phase, + fc1_counter_index=fc1_counter_index, + ) @dsl_user_op @cute.jit @@ -1218,16 +1255,29 @@ def _gen_work_from_cluster_idx( state = self._fused_state # Sentinel-by-default work tile; conditionally overwritten by decode. - base_work = self._ext.WorkTileInfo( - expert_idx=Int32(WorkTileState.DONE), - tile_m_idx=Int32(0), - tile_n_idx=Int32(0), - cumulative_data_physical_row=Int32(0), - cumulative_sf_physical_row=Int32(0), - cumulative_token_block_count=Int32(0), - valid_tokens_in_tile=Int32(0), - phase_and_peek=Int32(BlockPhase.None_), - ) + if const_expr(self.params.is_swap_ab): + base_work = self._ext.WorkTileInfo( + expert_idx=Int32(WorkTileState.DONE), + tile_m_idx=Int32(0), + tile_n_idx=Int32(0), + cumulative_data_physical_row=Int32(0), + cumulative_sf_physical_row=Int32(0), + cumulative_token_block_count=Int32(0), + valid_tokens_in_cta_tile=Int32(0), + phase_and_peek=Int32(BlockPhase.None_), + ) + else: + base_work = self._ext.WorkTileInfo( + expert_idx=Int32(WorkTileState.DONE), + tile_m_idx=Int32(0), + tile_n_idx=Int32(0), + cumulative_data_physical_row=Int32(0), + cumulative_sf_physical_row=Int32(0), + cumulative_token_block_count=Int32(0), + valid_tokens_in_cta_cluster_tile=Int32(0), + phase_and_peek=Int32(BlockPhase.None_), + fc1_counter_index=Int32(0), + ) # DSL carry for mutated self and while-condition fields. outer_group_end = state.current_group_end diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/flag_batch.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/flag_batch.py new file mode 100644 index 000000000000..ffc1a3c8b169 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/flag_batch.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Rotating-lane delayed release helper for done-counter publishing.""" + +import dataclasses +from typing import Any + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import Int64 + +from .ptx_helpers import red_add_release_gpu_s32 + + +@dataclasses.dataclass(frozen=True) +class GpuReleaseFlagBatchTracker: + """Batched done-counter publisher using GPU-scope release reductions. + + Carries the loop-carried per-thread accumulation state. + """ + + flag_addr: Int64 # per-lane counter-slot address (0 == null) + cumulated_flags: cutlass.Int32 # current batch fill count (uniform) + phase: cutlass.Int32 # current accumulated phase (uniform) + tid: cutlass.Int32 # lane/thread id within the rotating group + + @cute.jit + def _make( + self, + flag_addr: Int64, + cumulated_flags: cutlass.Int32, + phase: cutlass.Int32, + ) -> "GpuReleaseFlagBatchTracker": + return GpuReleaseFlagBatchTracker( + flag_addr=flag_addr, + cumulated_flags=cumulated_flags, + phase=phase, + tid=self.tid, + ) + + @cute.jit + def fire(self) -> None: + """Publish this lane's pending slot.""" + if self.flag_addr != Int64(0): + ptr = cute.make_ptr( + cutlass.Int32, + self.flag_addr, + AddressSpace.gmem, + assumed_align=4, + ) + red_add_release_gpu_s32(ptr, cutlass.Int32(1)) + + @cute.jit + def accumulate( + self, + next_phase: Any, + flush_threshold: int, + flag_addr: Int64, + no_fire: bool = False, + ) -> "GpuReleaseFlagBatchTracker": + if cutlass.const_expr(flush_threshold == 1): + if cutlass.const_expr(not no_fire): + per_lane_addr = Int64(0) + if self.tid == 0: + per_lane_addr = flag_addr + self._make( + flag_addr=per_lane_addr, + cumulated_flags=cutlass.Int32(1), + phase=self.phase, + ).fire() + return self._make( + flag_addr=Int64(0), + cumulated_flags=cutlass.Int32(0), + phase=cutlass.Int32(next_phase), + ) + + cur_addr = self.flag_addr + cumulated = self.cumulated_flags + if self.tid == cumulated: + cur_addr = flag_addr + cumulated = cumulated + 1 + + if cumulated == flush_threshold or next_phase != self.phase: + if not no_fire: + self._make( + flag_addr=cur_addr, + cumulated_flags=cumulated, + phase=self.phase, + ).fire() + cumulated = cutlass.Int32(0) + cur_addr = Int64(0) + + return self._make( + flag_addr=cur_addr, + cumulated_flags=cumulated, + phase=cutlass.Int32(next_phase), + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py index 96f9f27e5a02..e5c21708e1c6 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/grid_sync.py @@ -119,11 +119,12 @@ def software_grid_sync(counter_ptr, "selp.b32 %delta, $2, $3, %is_sm0;\n\t" "atom.release.gpu.global.add.u32 %old, [$0], %delta;\n\t" "SPIN:\n\t" - "ld.acquire.gpu.global.b32 %cur, [$0];\n\t" + "ld.relaxed.gpu.global.b32 %cur, [$0];\n\t" "xor.b32 %cur, %cur, %old;\n\t" "and.b32 %cur, %cur, 0x80000000;\n\t" "setp.eq.u32 %waiting, %cur, 0;\n\t" "@%waiting bra SPIN;\n\t" + "fence.acq_rel.gpu;\n\t" "DONE:\n\t" "}"), "l,r,r,r,r", diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py index 8682e566b020..ef06de0d1f03 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py @@ -1,5 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause """Fused fc1+fc2 swap-AB SwiGLU NVFP4 kernel for SM100.""" from typing import Literal, Optional, Tuple, Type @@ -33,6 +35,7 @@ from .megamoe_constants import (Nvfp4BlockSize, SupportedMmaTileM, SupportedMmaTileN) from .moe_utils import spin_wait +from .token_comm import CombineFormat # token_comm_args is an opaque subclass-owned bundle. The base only forwards it # to hook methods; ``None`` keeps the lean fc1+fc2 path free of token-comm IR. @@ -55,31 +58,32 @@ class Sm100SwapABSwigluFp4Fc12Kernel: _SmemMiscBudget = 1024 def __init__( - self, - # Geometry. - mma_tiler_mnk: Tuple[int, int, int], - cluster_shape_mnk: Tuple[int, int, int], - use_2cta_instrs: bool, - # Fused fc1+fc2 scheduler knobs. - group_hint: int, - token_padding_block: int, - sf_padding_block: int, - load_balance_mode: Literal["static", "atomic_counter"] = "static", - # Optional scheduler/codegen knobs. - static_expert_shape: Optional[Tuple[int, int, int]] = None, - force_static_sched: bool = True, - clc_bundle_size: Optional[int] = None, - num_sched_stages: Optional[int] = None, - acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, - sf_vec_size: int = 16, - scenario: Literal["2Dx3D"] = "2Dx3D", - *, - fc2_output_dtype: Type[cutlass.Numeric], - non_ubulk_fc2_store: bool = True, - in_kernel_fc2_reduce: bool = False, - token_back_by_dispatch: bool = False, - apply_topk_in_fc1: bool = True, - gate_up_clamp: Optional[float] = None, + self, + # Geometry. + mma_tiler_mnk: Tuple[int, int, int], + cluster_shape_mnk: Tuple[int, int, int], + use_2cta_instrs: bool, + # Fused fc1+fc2 scheduler knobs. + group_hint: int, + token_padding_block: int, + sf_padding_block: int, + load_balance_mode: Literal["static", "atomic_counter"] = "static", + # Optional scheduler/codegen knobs. + static_expert_shape: Optional[Tuple[int, int, int]] = None, + force_static_sched: bool = True, + clc_bundle_size: Optional[int] = None, + num_sched_stages: Optional[int] = None, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + sf_vec_size: int = 16, + scenario: Literal["2Dx3D"] = "2Dx3D", + *, + fc2_output_dtype: Type[cutlass.Numeric], + non_ubulk_fc2_store: bool = True, + in_kernel_fc2_reduce: bool = False, + token_back_by_dispatch: bool = False, + apply_topk_in_fc1: bool = True, + gate_up_clamp: Optional[float] = None, + epi_flag_batch: Optional[Tuple[int, int]] = (1, 1), ) -> None: if not force_static_sched: raise NotImplementedError( @@ -123,6 +127,7 @@ def __init__( self.token_back_by_dispatch = token_back_by_dispatch self.apply_topk_in_fc1 = apply_topk_in_fc1 self.gate_up_clamp = gate_up_clamp + self.epi_flag_batch = epi_flag_batch self._validate_mma_tiler_and_cluster_shape() self.mma_tiler = mma_tiler_mnk @@ -142,6 +147,8 @@ def __init__( self.sched_warp_id = 7 # Installed by token-comm subclasses. self.dispatch_warp_id: Optional[Tuple[int, int, int, int]] = None + self.token_back_warp_id: Optional[Tuple[int, int, int, int]] = None + self.token_back_standalone: bool = False self.threads_per_cta = 32 * len(( self.mma_warp_id, self.tma_a_warp_id, @@ -161,11 +168,39 @@ def __init__( # register allocation because setmaxnreg emission is gated by # ``self.enable_token_comm`` inside the device kernel. self.epi_reg_cnt = 256 - self.task_reg_cnt = 96 + self.task_reg_cnt = 72 self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols(self.arch) + def name(self) -> str: + """Canonical encoding of the codegen-affecting constexpr -- the compiled- + kernel cache key. ``local_rank`` excluded (deployment detail); only + ``force_static_sched`` / ``clc_bundle_size`` / ``num_sched_stages`` / + ``scenario`` are dropped -- every other constexpr is kept.""" + m, n, k = self.mma_tiler_mnk + cm, cn = self.cluster_shape_mn + exp = "x".join(map( + str, + self.static_expert_shape)) if self.static_expert_shape else "dyn" + epiflag = "x".join(map( + str, self.epi_flag_batch)) if self.epi_flag_batch else "none" + cta = "2_cta" if self.use_2cta_instrs else "1_cta" + fc2store = "fc2store_stg" if self.non_ubulk_fc2_store else "fc2store_ublk" + inkred = "inkernel_redg" if self.in_kernel_fc2_reduce else "no_inkernel_redg" + apply_topk = "apply_topk_fc1_pre_quant" if self.apply_topk_in_fc1 else "apply_topk_after_fc2" + # token-back is a token-communication concept -- it lives in the MegaMoE + # subclass name(), not the lean base. + return ( + "moe_fc12_fuse_nvfp4" + f"_mmatiler_{m}x{n}x{k}_cluster_{cm}x{cn}_{cta}_sched_{self.load_balance_mode}" + f"_expert_shape_{exp}_grouphint_{self.group_hint}" + f"_padding_{self.token_padding_block}x{self.sf_padding_block}" + f"_{fc2store}_{inkred}_{apply_topk}" + f"_fc2out{self.fc2_output_dtype.__name__}_sfvec{self.sf_vec_size}" + f"_acc{self.acc_dtype.__name__}_clamp{self.gate_up_clamp}_epiflag{epiflag}" + ) + def _validate_mma_tiler_and_cluster_shape(self) -> None: """Validate user-provided geometry against v1 fused-fc12 constraints. @@ -201,7 +236,7 @@ def _validate_mma_tiler_and_cluster_shape(self) -> None: f"cluster_shape M ({cm}) must be even when use_2cta_instrs=True" ) - def is_pow2(x): + def is_pow2(x: int) -> bool: return x > 0 and (x & (x - 1)) == 0 if cm * cn > 16 or not is_pow2(cm) or not is_pow2( @@ -312,16 +347,23 @@ def _setup_attributes(self) -> None: # dtype that lives in sC). fc2 output dtype is hard-coded as # ``BFloat16`` inside the epilogue's ``Fc2UnpackPermuteStg`` and # does not flow through this knob. + # combine_format is comm-side state: the MegaMoE subclass injects + # ``self.combine_format`` before super().__init__, so it is already on + # ``self`` here; the standalone base defaults to the bf16 (no-quant) + # combine. Only hooks / token_comm_args otherwise cross into this base. + if not hasattr(self, "combine_format"): + self.combine_format = CombineFormat.parse("bf16") self.epilogue = SwapABSwigluFp4Epilogue( mma_tiler_mnk=self.mma_tiler, cluster_shape_mn=self.cluster_shape_mn, use_2cta_instrs=self.use_2cta_instrs, sf_vec_size=self.sf_vec_size, fc1_output_dtype=self.fc1_output_dtype, - fc2_output_dtype=self.fc2_output_dtype, + combine_format=self.combine_format, non_ubulk_fc2_store=self.non_ubulk_fc2_store, in_kernel_fc2_reduce=self.in_kernel_fc2_reduce, token_back_by_dispatch=self.token_back_by_dispatch, + epi_flag_batch=self.epi_flag_batch, acc_dtype=self.acc_dtype, allow_overlap_acc=True, static_expert_shape=self.static_expert_shape, @@ -351,6 +393,13 @@ def _setup_attributes(self) -> None: self.num_sched_stages, self._smem_misc_budget_bytes(), ) + # print( + # f"[fc12 stages] num_ab_stage={self.num_ab_stage} " + # f"num_acc_stage={self.num_acc_stage} " + # f"misc_budget={self._smem_misc_budget_bytes()} " + # f"c_bytes_total={c_bytes_total} smem_cap={self.smem_capacity} " + # f"token_back_standalone={self.token_back_standalone}" + # ) self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( tiled_mma, @@ -581,7 +630,7 @@ def token_comm_hook_fc1_tma_b_predispatch_spin( """Emitted on the TMA-B warp at the head of each fc1-phase task tile, before its K-loop. Default: no-op. MegaMoE: blocking spin on the dispatch->fc1 release counter at ``cumulative_token_block_count + - tile_n_idx`` until it reaches ``work_tile_info.valid_tokens_in_tile``, + tile_n_idx`` until it reaches ``work_tile_info.valid_tokens_in_cta_tile``, unless ``work_tile_info.peek_ready`` already saturated it. Skipping this in the lean path is correct because in the lean path the per-tile input is already resident in GMEM at launch time.""" @@ -598,6 +647,18 @@ def token_comm_hook_dispatch_warp_body( ): """Subclass dispatch warp body; no-op in the lean kernel.""" + @cute.jit + def token_comm_hook_token_back_warp_body( + self, + token_comm_args, + token_comm_storage, + *, + warp_idx, + lane_idx, + tidx, + ): + """Subclass standalone token-back warp body; no-op in the lean kernel.""" + @cute.jit def token_comm_hook_kernel_tail( self, @@ -735,13 +796,20 @@ def __call__( c1 = cutlass.Int32(1) cutlass.Int32(0) + def round_up(value, multiple): + return ((value + multiple - 1) // multiple) * multiple + + def is_constexpr_int(*values) -> bool: + return all(isinstance(value, int) for value in values) + # A_gemm (fc1 weights): (experts, hidden, intermediate_gateup) # -> (M=intermediate_gateup, K=hidden, L=experts). experts, hidden_b, intermediate_gateup = fc1_weight.shape fc1_weight_gemm = cute.make_tensor( fc1_weight.iterator, cute.make_layout( - (intermediate_gateup, hidden_b, experts), + (cutlass.Int32(intermediate_gateup), cutlass.Int32(hidden_b), + cutlass.Int32(experts)), stride=(fc1_weight.stride[2], fc1_weight.stride[1], fc1_weight.stride[0]), ), @@ -752,7 +820,7 @@ def __call__( activation_gemm = cute.make_tensor( activation.iterator, cute.make_layout( - (tokens_sum, hidden, 1), + (tokens_sum, cutlass.Int32(hidden), c1), stride=(activation.stride[0], activation.stride[1], 0), ), ) @@ -762,7 +830,7 @@ def __call__( fc1_output_gemm = cute.make_tensor( fc1_output.iterator, cute.make_layout( - (tokens_sum, intermediate_downproj, 1), + (tokens_sum, cutlass.Int32(intermediate_downproj), c1), stride=(fc1_output.stride[0], fc1_output.stride[1], 0), ), ) @@ -775,16 +843,30 @@ def __call__( activation_sf_gemm = cute.make_tensor( activation_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - (tokens_sum_padded, hidden_padded, 1), self.sf_vec_size), + (tokens_sum_padded, cutlass.Int32(hidden_padded), c1), + self.sf_vec_size), ) - intermediate_gateup_padded_mul_hidden_padded = fc1_weight_sf.shape[1] - intermediate_gateup_padded = ( - intermediate_gateup_padded_mul_hidden_padded * - self.sf_vec_size) // hidden_padded + # M-side SF atom is 128 rows (same as the fc2 path below); the + # K-side ``sf_vec_size * 4`` granularity would understate the padded + # row count for non-128-aligned intermediate sizes. + intermediate_gateup_padded = round_up(intermediate_gateup, 128) + expected_fc1_weight_sf_cols = (intermediate_gateup_padded * + hidden_padded // self.sf_vec_size) + if cutlass.const_expr( + is_constexpr_int(fc1_weight_sf.shape[1], + expected_fc1_weight_sf_cols)): + if cutlass.const_expr( + fc1_weight_sf.shape[1] != expected_fc1_weight_sf_cols): + raise ValueError( + f"fc1_weight_sf.shape[1] ({fc1_weight_sf.shape[1]}) does not " + f"match intermediate_padded({intermediate_gateup_padded}) * " + f"hidden_padded({hidden_padded}) / sf_vec_size({self.sf_vec_size}) " + f"= {expected_fc1_weight_sf_cols}.") fc1_weight_sf_gemm = cute.make_tensor( fc1_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - (intermediate_gateup_padded, hidden_padded, experts), + (cutlass.Int32(intermediate_gateup_padded), + cutlass.Int32(hidden_padded), cutlass.Int32(experts)), self.sf_vec_size, ), ) @@ -799,7 +881,9 @@ def __call__( fc2_weight_gemm = cute.make_tensor( fc2_weight.iterator, cute.make_layout( - (hidden_b2, intermediate_downproj_b2, experts2), + (cutlass.Int32(hidden_b2), + cutlass.Int32(intermediate_downproj_b2), + cutlass.Int32(experts2)), stride=(fc2_weight.stride[2], fc2_weight.stride[1], fc2_weight.stride[0]), ), @@ -825,18 +909,36 @@ def __call__( fc1_output_sf_gemm_for_fc2_load = cute.make_tensor( fc1_output_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - (tokens_sum_padded_sf, intermediate_downproj_padded, 1), + (tokens_sum_padded_sf, + cutlass.Int32(intermediate_downproj_padded), c1), self.sf_vec_size, ), ) - hidden_padded_fc2_mul_intermediate_downproj_padded = fc2_weight_sf.shape[ - 1] - hidden_padded_fc2 = (hidden_padded_fc2_mul_intermediate_downproj_padded - * self.sf_vec_size) // intermediate_downproj_padded + hidden_padded_fc2 = round_up(hidden_b2, 128) + intermediate_downproj_padded = round_up( + intermediate_downproj_b2, + self.sf_vec_size * 4, + ) + expected_fc2_weight_sf_cols = (hidden_padded_fc2 * + intermediate_downproj_padded // + self.sf_vec_size) + if cutlass.const_expr( + is_constexpr_int(fc2_weight_sf.shape[1], + expected_fc2_weight_sf_cols)): + if cutlass.const_expr( + fc2_weight_sf.shape[1] != expected_fc2_weight_sf_cols): + raise ValueError( + f"fc2_weight_sf.shape[1] ({fc2_weight_sf.shape[1]}) does not " + f"match hidden_padded({hidden_padded_fc2}) * " + f"intermediate_padded({intermediate_downproj_padded}) / " + f"sf_vec_size({self.sf_vec_size}) = {expected_fc2_weight_sf_cols}." + ) fc2_weight_sf_gemm = cute.make_tensor( fc2_weight_sf.iterator, blockscaled_utils.tile_atom_to_shape_SF( - (hidden_padded_fc2, intermediate_downproj_padded, experts2), + (cutlass.Int32(hidden_padded_fc2), + cutlass.Int32(intermediate_downproj_padded), + cutlass.Int32(experts2)), self.sf_vec_size, ), ) @@ -907,7 +1009,7 @@ def __call__( self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape, - internal_type=cutlass.Uint64, + internal_type=cutlass.Uint16, ) # TMA load SFB1 (= activation_sf, fc1 activation SFs) @@ -922,7 +1024,7 @@ def __call__( self.mma_tiler_sfb, tiled_mma_sfb, self.cluster_layout_sfb_vmnk.shape, - internal_type=cutlass.Uint64, + internal_type=cutlass.Uint16, ) # TMA store for fc1 NVFP4 output (via SMEM-staged bulk store). @@ -980,7 +1082,7 @@ def __call__( self.mma_tiler, tiled_mma, self.cluster_layout_vmnk.shape, - internal_type=cutlass.Uint64, + internal_type=cutlass.Uint16, ) tma_atom_fc1_output_sf_as_fc2_input, tma_tensor_fc1_output_sf_as_fc2_input = cute.nvgpu.make_tiled_tma_atom_B( sfb_op, @@ -989,7 +1091,7 @@ def __call__( self.mma_tiler_sfb, tiled_mma_sfb, self.cluster_layout_sfb_vmnk.shape, - internal_type=cutlass.Uint64, + internal_type=cutlass.Uint16, ) # ── Scheduler params + grid + launch ── @@ -1731,7 +1833,7 @@ class SharedStorage: if not is_leader_cta: load_shift = dynamic_mainloop.compute_non_leader_cta_load_shift( valid_tokens_in_tile=work_tile_info. - valid_tokens_in_tile, + valid_tokens_in_cta_tile, mma_tiler_n=self.mma_tiler[1], ) real_b = cute.domain_offset((load_shift, 0, 0), @@ -1834,7 +1936,7 @@ class SharedStorage: spin_wait( counter_ptr, lambda v: v >= fc2_spin_threshold, - fail_sleep_cycles=20, + fail_sleep_cycles=500, ) iket.range_pop() @@ -1857,7 +1959,7 @@ class SharedStorage: if not is_leader_cta: load_shift = dynamic_mainloop.compute_non_leader_cta_load_shift( valid_tokens_in_tile=work_tile_info. - valid_tokens_in_tile, + valid_tokens_in_cta_tile, mma_tiler_n=self.mma_tiler[1], ) real_b = cute.domain_offset((load_shift, 0, 0), @@ -2023,8 +2125,10 @@ class SharedStorage: ab_consumer.reset() peek_ab_full_status = cutlass.Boolean(1) if k_tile_cnt > 0: + iket.range_push("mma_acquire") peek_ab_full_status = ab_consumer.try_wait() acc_pipeline.producer_acquire(acc_producer_state) + iket.range_pop() # Apply TMEM pointer offset hack when mma_tiler_n == 64. tCtSFB_mma = tCtSFB @@ -2071,7 +2175,7 @@ class SharedStorage: sfb_tensor=tCtSFB_mma, k_tile_idx=k_tile, valid_tokens_in_tile=work_tile_info. - valid_tokens_in_tile, + valid_tokens_in_cta_tile, mma_tiler_mnk=self.mma_tiler_mnk, ) handle.release() @@ -2151,13 +2255,31 @@ class SharedStorage: cute.arch.warpgroup_reg_dealloc(self.task_reg_cnt) lane_idx_for_dispatch = cute.arch.lane_idx() - self.token_comm_hook_dispatch_warp_body( - token_comm_args, - token_comm_storage, - warp_idx=warp_idx, - lane_idx=lane_idx_for_dispatch, - tidx=tidx, - ) + if cutlass.const_expr(self.token_back_standalone): + if warp_idx < self.token_back_warp_id[0]: + self.token_comm_hook_dispatch_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + else: + self.token_comm_hook_token_back_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) + else: + self.token_comm_hook_dispatch_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx_for_dispatch, + tidx=tidx, + ) # ════════════════════════════════════════════════════════════════════ # Kernel tail hook (MegaMoE-only path; lean base = no-op) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py index 2cdaa13a57bc..1a43b5682a9d 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_constants.py @@ -1,6 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Shared constants for the fused fc1+fc2 MegaMoE path.""" +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Shared constants for the fused fc1+fc2 MegaMoE path. + +``Fp32Max`` is the only CuteDSL-typed constant here; it is resolved +lazily via module ``__getattr__`` (PEP 562) so importing this module -- +and therefore the package ``__init__`` and the host-side tactic +enumeration in ``cute_dsl_megamoe_custom_op.py``, which only need the +plain-Python constants -- does not require a cutlass-dsl install. The +kernel modules that consume ``Fp32Max`` (``epilogue_refactor.py``) +import cutlass themselves, so the lazy resolution always succeeds +wherever the constant is actually used. +""" Nvfp4BlockSize = 16 SfPaddingBlock = 128 @@ -9,5 +21,20 @@ Nvfp4E2M1Max = 6.0 Fp8E4M3FNMax = 448.0 +Nvfp4E2M1RcpLimit = 1.0 / Nvfp4E2M1Max +Fp8E4M3RcpLimit = 1.0 / Fp8E4M3FNMax + SupportedMmaTileM = (128, 256) SupportedMmaTileN = (64, 128, 256) + + +def __getattr__(name: str): + if name == "Fp32Max": + from cutlass.cutlass_dsl import Float32 + + value = Float32(3.40282346638528859812e38) + # Cache so later lookups (and ``from ... import Fp32Max``) bind the + # SAME object instead of re-wrapping per access. + globals()[name] = value + return value + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py index 112c441c11ed..0ea3f48ae9ac 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py @@ -1,5 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause """MegaMoE fused dispatch + fc1 + fc2 + combine kernel. The base class owns the local fc1/fc2 GEMM pipeline. This subclass owns the @@ -38,7 +40,7 @@ # quoted explicitly. import dataclasses -from typing import Any, Dict, List, Optional, Tuple, Type +from typing import Any, Dict, List, Literal, Optional, Tuple, Type import cutlass import cutlass.cute as cute @@ -46,8 +48,10 @@ from cutlass.cutlass_dsl import Int64 from .kernel_fc12 import Sm100SwapABSwigluFp4Fc12Kernel +from .token_comm import CombineFormat from .token_comm import TokenCommArgs as ExtractedTokenCommArgs -from .token_comm import TokenInPullTokenBackPush +from .token_comm import TokenInPullTokenBackPush, TokenSrcMetadata +from .topk_reduce import TopkReduce # ============================================================================= # Module-level constants. @@ -60,9 +64,9 @@ # Dispatch warp count. _DispatchWarpCount = 4 -# Per-pool-slot provenance record consumed by combine STG redirect (S3). -# Three packed Uint32 fields = 12 bytes: ``{src_rank, src_token, src_topk}``. -_TokenMetadataBytes = 12 +# Per-pool-slot provenance record consumed by combine STG redirect (S3) and +# token-back; one i64 = {src_rank, src_token, src_topk} (see TokenSrcMetadata). +_TokenMetadataBytes = TokenSrcMetadata.nbytes # NVLink signal slots used by the DeepGEMM-style phase/sign barrier. # A separate local counter selects phase/sign; the signal slots are not reset @@ -168,17 +172,42 @@ def __init__( # MegaMoE-specific independent constants. *, world_size: int, - local_rank: int, num_topk: int, max_tokens_per_rank: int, hidden: int, fc2_output_dtype: Type[cutlass.Numeric], + combine_format: CombineFormat = CombineFormat.parse("bf16"), non_ubulk_fc2_store: bool = True, in_kernel_fc2_reduce: bool = False, - token_back_by_dispatch: bool = False, + token_back_mode: Literal["epi_warps", "standalone_warps", + "reuse_dispatch_warps"] = "epi_warps", apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, + epi_flag_batch: Optional[Tuple[int, int]] = (1, 1), + flag_batch: int = 1, ) -> None: + # The combine wire format drives the fc2 epilogue encoder, token_comm + # push, and the combine_quant/combine_sf workspace sizing. The dataflow + # (workspace carve, views, arg threading, epilogue encode, topk_reduce + # receiver) is wired for every format, quantized included -- see the + # closed-loop note above the fused launch in __call__. The guards below + # only reject combinations the kernel cannot express + # (in_kernel_fc2_reduce is bf16-only; FP4 needs non_ubulk_fc2_store). + self.combine_format = combine_format + if in_kernel_fc2_reduce and combine_format.is_quantized: + raise ValueError( + f"in_kernel_fc2_reduce requires a non-quantized (bf16) combine " + f"format; got {combine_format}.") + if in_kernel_fc2_reduce and not apply_topk_in_fc1: + raise ValueError( + "in_kernel_fc2_reduce requires apply_topk_in_fc1=True; " + "the REDG path can only atomic-add terms whose topk score " + "was already absorbed before fc2.") + if (combine_format.act_dtype is cutlass.Float4E2M1FN + and not non_ubulk_fc2_store): + raise ValueError( + f"{combine_format} combine requires non_ubulk_fc2_store=True " + "(the UBLK fc2 store path cannot scalar-dereference FP4).") if static_expert_shape is None: raise NotImplementedError( "Sm100MegaMoEKernel currently requires " @@ -191,6 +220,20 @@ def __init__( f"hidden ({hidden}) must equal " f"static_expert_shape[2] ({static_expert_shape[2]}).") + # token_back_mode selects where the cross-rank fc2 push-back runs: + # epi_warps -> epilogue warps STG directly to the peer + # standalone_warps -> dedicated warp group 12-15, concurrent + # with dispatch_pull + # reuse_dispatch_warps -> dispatch warps 8-11 push after dispatch_pull + # The two non-epi modes both stage fc2 to a local workspace first, i.e. + # token_back_by_dispatch=True; epi_warps keeps the epilogue STG redirect. + if token_back_mode not in ("epi_warps", "standalone_warps", + "reuse_dispatch_warps"): + raise ValueError( + f"token_back_mode must be 'epi_warps', 'standalone_warps', " + f"or 'reuse_dispatch_warps'; got {token_back_mode!r}.") + token_back_by_dispatch = token_back_mode != "epi_warps" # nosec B105 + super().__init__( mma_tiler_mnk=mma_tiler_mnk, cluster_shape_mnk=cluster_shape_mnk, @@ -212,23 +255,33 @@ def __init__( token_back_by_dispatch=token_back_by_dispatch, apply_topk_in_fc1=apply_topk_in_fc1, gate_up_clamp=gate_up_clamp, + epi_flag_batch=epi_flag_batch, ) self.enable_token_comm = True self.dispatch_warp_id = (8, 9, 10, 11) + # Standalone token-back: a dedicated 4-warp group (12-15) doing + # token_back_by_push concurrently with dispatch_pull, selected by the + # user-facing token_back_mode knob ("standalone_warps"). + self.token_back_mode = token_back_mode + self.token_back_standalone = token_back_mode == "standalone_warps" # nosec B105 + self.token_back_warp_id = (12, 13, 14, + 15) if self.token_back_standalone else None + num_token_back_warps = (len(self.token_back_warp_id) + if self.token_back_standalone else 0) self.threads_per_cta = 32 * ( len(self.epilogue_warp_id) + 1 # mma + 1 # tma_a + 1 # tma_b + 1 # sched - + len(self.dispatch_warp_id)) + + len(self.dispatch_warp_id) + num_token_back_warps) # Independent MegaMoE-specific constants. self.world_size = world_size - self.local_rank = local_rank self.num_topk = num_topk self.max_tokens_per_rank = max_tokens_per_rank self.hidden = hidden + self.flag_batch = flag_batch # stored so name() can encode it # static_expert_shape = (num_experts_per_rank, intermediate_gateup, hidden). self.num_experts_per_rank = static_expert_shape[0] @@ -245,6 +298,15 @@ def __init__( # Cross-rank totals: per-rank count * world_size. self.num_total_experts = world_size * self.num_experts_per_rank + # Per-(token, topk) SF block padded to a 16 B multiple so the token-back + # cp.async.bulk push moves one aligned block per shot; only the first + # hidden//scale_block entries of each block are valid (the rest is the + # alignment gap). ``None`` for the bf16 baseline (no SF plane). + self.sf_block_pad = (_round_up( + self.hidden // self.combine_format.scale_block, + 16 // (int(self.combine_format.scale_dtype.width) // 8), + ) if self.combine_format.is_quantized else None) + # Per-task-tile release-counter granularity used by dispatch_pull. self.cluster_tile_tokens = mma_tiler_mnk[1] * cluster_shape_mnk[1] @@ -275,18 +337,32 @@ def __init__( (self.hidden + cluster_fc2_tile_hidden - 1) // cluster_fc2_tile_hidden) * self.cluster_shape_mn[0] + # Token-back warps run when they push the DATA plane (dispatch modes) OR + # the SF plane (any quantized combine, including the epi_warps data path + # where the epilogue STGs the data straight to the peer but SF still + # needs the staged token-contiguous push). + self.token_back_enabled = (self.token_back_by_dispatch + or self.combine_format.is_quantized) + # Homomorphic to the fc1+fc2 scheduler: atomic_counter token-back only + # nets a win with enough tokens, the same condition that selects the + # atomic_counter fc1+fc2 scheduler. Static when token-back is off. + self.token_back_schedule_mode = (self.load_balance_mode if + self.token_back_enabled else "static") + self.token_comm = TokenInPullTokenBackPush( world_size=self.world_size, - local_rank=self.local_rank, num_topk=self.num_topk, num_experts_per_rank=self.num_experts_per_rank, num_total_experts=self.num_total_experts, hidden=self.hidden, fc1_token_dtype=cutlass.Float4E2M1FN, - fc2_output_dtype=(self.fc2_output_dtype - if self.token_back_by_dispatch else None), + combine_format=self.combine_format, + token_back_by_dispatch=self.token_back_by_dispatch, fc2_publishes_per_token_cluster_tile= fc2_publishes_per_token_cluster_tile, + token_back_reduce_topk=(self.token_back_by_dispatch + and self.in_kernel_fc2_reduce), + token_back_standalone=self.token_back_standalone, sf_uint32_per_token=self.sf_uint32_per_token, token_padding_block=self.token_padding_block, sf_padding_block=self.sf_padding_block, @@ -294,6 +370,9 @@ def __init__( cluster_shape_mn=self.cluster_shape_mn, dispatch_warp_start=self.dispatch_warp_id[0], num_other_warps=num_other_warps, + flag_batch=flag_batch, + is_swap_ab=True, + token_back_schedule_mode=self.token_back_schedule_mode, ) # Region layout (same call drives both get_workspace_sizes() and @@ -313,6 +392,22 @@ def __init__( for r in self._shared_region_specs } + # Counter-prefix byte extents: the accumulating counters are front-placed + # (see _build_*_region_specs), so the leading bytes up to the first data + # region cover exactly the per-launch zero set. ``tail_reset_counters`` + # bulk-zeros these each launch; only the first launch needs a caller-zeroed + # workspace. 128B-aligned offsets => multiples of 4 (Int32 zeroing exact). + local_leading = self._local_offsets[ + "l1_token_buffer"] # first data region + shared_leading = self._shared_offsets[ + "src_token_topk_idx"] # first data region + self.require_zero_workspace_leading_bytes: Tuple[int, int] = ( + local_leading, + shared_leading, + ) + self.local_zero_i32_count = local_leading // 4 + self.shared_zero_i32_count = shared_leading // 4 + # ========================================================================= # SMEM budget hook (base override) # ========================================================================= @@ -322,9 +417,13 @@ def _dispatch_smem_bytes(self) -> int: pull_mbar_bytes = _DispatchWarpCount * 8 expert_count_bytes = self.num_total_experts * 4 pull_buffer_bytes = _DispatchWarpCount * self.hidden_bytes - return (_round_up(pull_mbar_bytes, 16) + - _round_up(expert_count_bytes, 16) + - _round_up(pull_buffer_bytes, 128)) + total = (_round_up(pull_mbar_bytes, 16) + + _round_up(expert_count_bytes, 16) + + _round_up(pull_buffer_bytes, 128)) + if self.token_back_standalone: + total += (_round_up(_DispatchWarpCount * 8, 16) + _round_up( + _DispatchWarpCount * self.token_comm.tb_chunk_bytes, 128)) + return total def _smem_misc_budget_bytes(self) -> int: """Base misc reservation plus dispatch-warp SMEM.""" @@ -409,65 +508,108 @@ def _build_local_region_specs(self) -> List[_RegionSpec]: (pool_token_capacity + mma_tiler_n - 1) // mma_tiler_n + num_experts_per_rank) + # === Accumulating-counter prefix =========================================== + # Front-placed so the per-launch reset is a single bulk zero of + # ``[0:local_leading]`` (the bytes up to the first data region). These hold + # spin thresholds / write cursors / phase-flip counters that the kernel + # accumulates across the launch and that MUST start at 0 each launch; the + # kernel tail (``tail_reset_counters``) bulk-zeros this prefix, so only the + # FIRST launch relies on a caller-zeroed workspace. specs: List[_RegionSpec] = [ - # L1 input pool (dispatch_pull writes -> fc1 reads). Stored - # as Uint8 bytes; the NVFP4 view at the same offset is - # built inside ``__call__``. _RegionSpec( - "l1_token_buffer", - cutlass.Uint8, - (pool_token_capacity, hidden_bytes), - 128, + "l1_arrival_count", + cutlass.Int32, + (pool_task_tile_capacity, ), + 16, ), - # Stored as Int32 (dispatch_pull's 32 b read/write); the FP8 - # view for activation_sf is built at the same offset. - # 1D Int32 atom-flat buffer. Total Int32 count = pool_sf_capacity - # (M-axis token positions) * sf_uint32_per_token (K-atom count), - # laid out atom-by-atom per cute SFA layout. dispatch writes - # individual Int32 slots via the linear offset returned by - # ``src/sf_swizzle.py:sf_atom_int32_offset``; the mma side - # re-views this same byte buffer through ``tile_atom_to_shape_SF`` - # which reads back the atom-swizzled bytes. _RegionSpec( - "l1_sf_buffer", - cutlass.Int32, - (pool_sf_capacity * sf_uint32_per_token, ), + "expert_send_count", + cutlass.Int64, + (num_total_experts, ), 16, ), _RegionSpec( - "l1_topk_weights_buffer", - cutlass.Float32, - (pool_token_capacity, ), + "grid_sync_counter", + cutlass.Int32, + (_GridSyncSlotCount, ), 16, ), _RegionSpec( - "l1_arrival_count", + "fc1_done_counter", cutlass.Int32, - (pool_task_tile_capacity, ), + (fc1_done_slots, ), 16, ), + ] + if self.token_back_enabled: + # Per-expert fc2 completion gate consumed by the token-back push + # (DATA and/or SF). Published by the fc2 epilogue for every enabled + # token-back path, including the epi_warps SF-only push. + specs.append( + _RegionSpec( + "fc2_done_counter", + cutlass.Int32, + (num_experts_per_rank, ), + 16, + )) + if self.token_back_schedule_mode == "atomic_counter": # nosec B105 + specs.append( + _RegionSpec( + "token_back_schedule_counter", + cutlass.Int32, + (1, ), + 16, + )) + if self.load_balance_mode == "atomic_counter": + specs.append( + _RegionSpec( + "load_balance_counter", + cutlass.Int32, + (1, ), + 16, + )) + + # === Data buffers (overwritten each launch; NOT zeroed) ==================== + # ``l1_token_buffer`` MUST be the first data region: ``__init__`` derives + # ``local_leading`` from its offset (= end of the counter prefix). + specs += [ + # L1 input pool (dispatch_pull writes -> fc1 reads), Uint8 bytes; the + # NVFP4 view at the same offset is built inside ``__call__``. _RegionSpec( - "token_src_metadata", + "l1_token_buffer", cutlass.Uint8, - (pool_token_capacity, _TokenMetadataBytes), - 16, + (pool_token_capacity, hidden_bytes), + 128, ), + # Persisted across launches (deliberately OUT of the zero prefix): the + # sense-reversing nvlink barrier rides this phase counter across launch + # boundaries (non-ncu back-to-back), and ncu kernel replay restores it + # via its local-memory snapshot. Only the FIRST launch relies on the + # caller-zeroed workspace. _RegionSpec( - "expert_send_count", - cutlass.Int64, - (num_total_experts, ), + "nvlink_barrier_counter", + cutlass.Int32, + (1, ), 16, ), + # Int32 atom-flat SF buffer (dispatch_pull 32b read/write; FP8 view at + # the same offset). Count = pool_sf_capacity * sf_uint32_per_token. _RegionSpec( - "grid_sync_counter", + "l1_sf_buffer", cutlass.Int32, - (_GridSyncSlotCount, ), + (pool_sf_capacity * sf_uint32_per_token, ), 16, ), _RegionSpec( - "nvlink_barrier_counter", - cutlass.Int32, - (1, ), + "l1_topk_weights_buffer", + cutlass.Float32, + (pool_token_capacity, ), + 16, + ), + _RegionSpec( + "token_src_metadata", + cutlass.Uint8, + (pool_token_capacity, _TokenMetadataBytes), 16, ), _RegionSpec( @@ -482,36 +624,33 @@ def _build_local_region_specs(self) -> List[_RegionSpec]: (sf_total_rows_upper, sf_block_cols), 128, ), - _RegionSpec( - "fc1_done_counter", - cutlass.Int32, - (fc1_done_slots, ), - 16, - ), ] if self.token_back_by_dispatch: + # Local fc2 DATA staging (token_back_by_dispatch modes only); the + # wire-format dtype sizes the plane. specs.append( _RegionSpec( "fc2_output_workspace", - self.fc2_output_dtype, + self.combine_format.act_dtype, (pool_token_capacity, 1, self.hidden), 128, )) + # The per-block SF plane is ALWAYS staged locally (then pushed + # token-contiguously by the dispatch / standalone warps), independent of + # whether the DATA path goes local or straight to a peer: writing SF + # per-token to a peer would scatter one warp's 32 lanes across up to 32 + # ranks and explode the NVLink request count. So it is allocated for + # every quantized format, not only the dispatch data path. + if self.combine_format.is_quantized: + # Flat padded capacity (pool_token * sf_block_pad scale entries); the + # (pool_token, 1, hidden//scale_block) logical shape + 16 B-aligned + # per-block stride is assembled where the view is built in __call__. specs.append( _RegionSpec( - "fc2_done_counter", - cutlass.Int32, - (num_experts_per_rank, ), - 16, - )) - - if self.load_balance_mode == "atomic_counter": - specs.append( - _RegionSpec( - "load_balance_counter", - cutlass.Int32, - (1, ), - 16, + "fc2_output_sf", + self.combine_format.scale_dtype, + (pool_token_capacity * self.sf_block_pad, ), + 128, )) return specs @@ -536,13 +675,13 @@ def _build_shared_region_specs(self) -> List[_RegionSpec]: # edge any peer might publish for this rank's local experts. max_slot = max_tokens_per_rank * num_topk - return [ - _RegionSpec( - "src_token_topk_idx", - cutlass.Int32, - (num_experts_per_rank, world_size, max_slot), - 16, - ), + # Accumulating counters first (zero-prefix -> tail bulk-zeros + # ``[0:shared_leading]``); then the data/signal regions that must persist + # across launches (``src_token_topk_idx`` overwritten each launch; + # ``nvlink_barrier_signal`` is phase-flip and must NOT be zeroed). + # ``src_token_topk_idx`` MUST be the first non-counter region: ``__init__`` + # derives ``shared_leading`` from its offset (= end of the counter prefix). + specs: List[_RegionSpec] = [ _RegionSpec( "expert_recv_count", cutlass.Int64, @@ -555,6 +694,12 @@ def _build_shared_region_specs(self) -> List[_RegionSpec]: (num_experts_per_rank, ), 16, ), + _RegionSpec( + "src_token_topk_idx", + cutlass.Int32, + (num_experts_per_rank, world_size, max_slot), + 16, + ), _RegionSpec( "nvlink_barrier_signal", cutlass.Int32, @@ -562,6 +707,37 @@ def _build_shared_region_specs(self) -> List[_RegionSpec]: 16, ), ] + # separate-kernel-reduce only: the per-topk fc2 combine staging buffer is + # internalized here instead of being a caller tensor, so the public output + # is the 2D (T, hidden) reduce result. It is the cross-rank combine STG + # target, hence must live on the symmetric heap (= this shared workspace); + # appended after the data regions so it stays out of the per-launch zero + # prefix (only the first launch relies on the caller-zeroed workspace). + if not self.in_kernel_fc2_reduce: + # combine_quant: the cross-rank combine data plane, one cell per + # (token, topk). dtype follows the wire format -- the bf16 baseline is + # byte-identical to the old ``combine_partial``; fp4/e4m3 shrink it. + specs.append( + _RegionSpec( + "combine_quant", + self.combine_format.act_dtype, + (max_tokens_per_rank, num_topk, self.hidden), + 128, + )) + # combine_sf: per-block scale plane; only quantized formats carry one. + if self.combine_format.is_quantized: + # Flat padded capacity (max_tokens * num_topk * sf_block_pad scale + # entries); the (max_tokens, num_topk, hidden//scale_block) logical + # shape + 16 B-aligned per-block stride is assembled where the view + # is built in __call__. + specs.append( + _RegionSpec( + "combine_sf", + self.combine_format.scale_dtype, + (max_tokens_per_rank * num_topk * self.sf_block_pad, ), + 128, + )) + return specs # ========================================================================= # Public: workspace size query @@ -581,19 +757,24 @@ def get_workspace_sizes(self) -> Tuple[int, int]: @staticmethod def _make_typed_view( - byte_workspace: cute.Tensor, + byte_base: cute.Pointer, byte_offset: int, cute_dtype: Any, shape: Tuple[int, ...], stride: Optional[Tuple[int, ...]], assumed_align: int, ) -> cute.Tensor: - """Build a typed cute view at ``byte_offset`` of the opaque workspace.""" - # Large MegaMoE problems can place later workspace regions above the - # 2 GiB / 4 GiB boundary. Keep the base adjustment in 64-bit pointer - # arithmetic so region starts such as fc1_output_sf / counters do not - # wrap before the typed view is built. - byte_ptr = byte_workspace.iterator + Int64(byte_offset) + """Build a typed cute view at ``byte_offset`` of the opaque workspace. + + The workspace is a raw ``cute.Pointer`` (uint8 gmem base), not a tensor: + the kernel only ever needs the base address + its own byte-offset table, so + a tensor's shape would be both ignored AND, for >2 GiB workspaces (the + internalized combine staging), overflow cute's 32-bit memref shape field. + """ + # Large MegaMoE problems place later regions above the 2 GiB / 4 GiB + # boundary; keep the base adjustment in 64-bit pointer arithmetic so region + # starts (fc1_output_sf / counters) do not wrap before the typed view. + byte_ptr = byte_base + Int64(byte_offset) typed_iter = cute.make_ptr( cute_dtype, byte_ptr.toint(), @@ -605,7 +786,7 @@ def _make_typed_view( def _view_local( self, - local_workspace: cute.Tensor, + local_workspace: cute.Pointer, name: str, *, cute_dtype: Optional[Any] = None, @@ -628,7 +809,7 @@ def _view_local( def _view_shared( self, - shared_workspace: cute.Tensor, + shared_workspace: cute.Pointer, name: str, *, cute_dtype: Optional[Any] = None, @@ -646,7 +827,7 @@ def _view_shared( def _partition_region( self, - byte_workspace: cute.Tensor, + byte_workspace: cute.Pointer, offsets: Dict[str, int], spec: _RegionSpec, *, @@ -680,6 +861,39 @@ def _partition_region( # __call__ # ========================================================================= + def name(self) -> str: + """Full compiled-kernel cache key. Self-contained on purpose (no shared + helper): mirrors the base fc12 fields and appends the MegaMoE-specific + ones -- keep the shared part in sync with the base ``name()``. + ``local_rank`` excluded (deployment detail); same dropped set as base.""" + m, n, k = self.mma_tiler_mnk + cm, cn = self.cluster_shape_mn + exp = "x".join(map( + str, + self.static_expert_shape)) if self.static_expert_shape else "dyn" + epiflag = "x".join(map( + str, self.epi_flag_batch)) if self.epi_flag_batch else "none" + cta = "2_cta" if self.use_2cta_instrs else "1_cta" + fc2store = "fc2store_stg" if self.non_ubulk_fc2_store else "fc2store_ublk" + inkred = "inkernel_redg" if self.in_kernel_fc2_reduce else "no_inkernel_redg" + apply_topk = "apply_topk_fc1_pre_quant" if self.apply_topk_in_fc1 else "apply_topk_after_fc2" + token_back = { + "epi_warps": "epiwarps", + "standalone_warps": "standalone", + "reuse_dispatch_warps": "reuse_dispatch", + }.get(self.token_back_mode, self.token_back_mode) + return ( + "megamoe_nvfp4" + f"_mmatiler_{m}x{n}x{k}_cluster_{cm}x{cn}_{cta}_sched_{self.load_balance_mode}" + f"_expert_shape_{exp}_grouphint_{self.group_hint}" + f"_padding_{self.token_padding_block}x{self.sf_padding_block}" + f"_{fc2store}_{inkred}_token_back_by_{token_back}_{apply_topk}" + f"_fc2out{self.fc2_output_dtype.__name__}_combine{self.combine_format}_sfvec{self.sf_vec_size}" + f"_acc{self.acc_dtype.__name__}_clamp{self.gate_up_clamp}_epiflag{epiflag}" + # MegaMoE-specific constexpr: + f"_ep_{self.world_size}_topk_{self.num_topk}_maxtoken_{self.max_tokens_per_rank}" + f"_flagbatch_{self.flag_batch}") + @cute.jit def __call__( self, @@ -697,12 +911,14 @@ def __call__( fc1_alpha: cute.Tensor, fc2_alpha: cute.Tensor, fc1_norm_const: cute.Tensor, - # Combine destination (peer write target under S3; local fc2 - # output region under S2 -- same memory, same caller). - combine_output: cute.Tensor, # (T, num_topk, hidden) BF16 + # Final combined output the caller consumes: 2D (T, hidden). Under + # in_kernel_reduce it is the cross-rank REDG target (symmetric heap); + # under separate_kernel_reduce it is the local tail-reduce destination + # while the per-topk staging lives in the internal ``combine_quant``. + output_activation: cute.Tensor, # (T, hidden) BF16 # Opaque workspaces. - local_workspace: cute.Tensor, # (local_ws_bytes,) Uint8 - shared_workspace: cute.Tensor, # (shared_ws_bytes,) Uint8 + local_workspace: cute.Pointer, # uint8 gmem base of (local_ws_bytes,) + shared_workspace: cute.Pointer, # uint8 gmem base of (shared_ws_bytes,) # Runtime host payload; packed into ``SymBuffer{world_size}`` # before entering the device kernel. peer_rank_ptr_mapper_host, @@ -722,10 +938,13 @@ def __call__( unconstrained (cuda local or sym heap). * ``fc1_weight`` / ``fc1_weight_sf`` / ``fc2_weight`` / ``fc2_weight_sf`` are local-only. - * ``combine_output`` is the per-rank S3 combine STG target; - under S2 it acts as the rank's local BF16 fc2 output. - Placement: sym heap (peer write target) or local in the - single-rank degenerate case. + * ``output_activation`` is the 2D (T, hidden) result. Under + ``in_kernel_reduce`` it is the per-rank cross-rank combine STG + target and MUST be reachable via the peer mapper (sym heap, or + local in the single-rank degenerate case). Under + ``separate_kernel_reduce`` the cross-rank target is the internal + ``combine_quant`` staging region and ``output_activation`` only + receives the local tail reduce, so it may be plain local memory. Workspace zero-init contract: caller is currently expected to zero ``shared_workspace`` before launch (the dispatch @@ -853,34 +1072,129 @@ def __call__( stride=(2, ), ) + # MoE-domain ``(token_max, topk, hidden)`` cross-rank combine STG target. + # * in_kernel_reduce: REDG collapses topk on the fly, so this is a + # ``(T, 1, hidden)`` view of the caller's 2D ``output_activation`` + # (the epilogue's topk index is a constexpr 0 in this mode). + # * separate_kernel_reduce: peers write one cell per (token, topk) into + # the internal ``combine_quant`` staging; the tail reduce below + # collapses topk into ``output_activation``. + if cutlass.const_expr(self.in_kernel_fc2_reduce): + combine_target = cute.make_tensor( + output_activation.iterator, + cute.make_layout( + (self.max_tokens_per_rank, 1, self.hidden), + stride=(self.hidden, self.hidden, 1), + ), + ) + else: + combine_target = self._view_shared(shared_workspace, + "combine_quant") + + # Per-block scale plane parallel to combine_quant; quantized formats only + # (bf16 carries no SF, and in_kernel_reduce is bf16-by-construction). + sf_blocks = (self.hidden // self.combine_format.scale_block + if self.combine_format.is_quantized else 0) + if cutlass.const_expr(self.combine_format.is_quantized + and not self.in_kernel_fc2_reduce): + # Assemble the (token, topk, valid_blocks) logical view + 16 B-aligned + # per-block stride over the flat combine_sf capacity. + combine_sf = self._view_shared( + shared_workspace, + "combine_sf", + shape=(self.max_tokens_per_rank, self.num_topk, sf_blocks), + stride=(self.num_topk * self.sf_block_pad, self.sf_block_pad, + 1), + ) + else: + combine_sf = None + + if cutlass.const_expr(self.combine_format.is_quantized): + # Same: logical (pool_token, 1, valid_blocks) + padded per-block stride + # over the flat fc2_output_sf capacity. + fc2_output_sf_phys = self._view_local( + local_workspace, + "fc2_output_sf", + shape=(self.pool_token_capacity, 1, sf_blocks), + stride=(self.sf_block_pad, self.sf_block_pad, 1), + ) + sf_vec = self.combine_format.scale_block + sf_layout = fc2_output_sf_phys.layout + fc2_output_sf = cute.make_tensor( + fc2_output_sf_phys.iterator, + cute.make_layout( + (sf_layout.shape[0], sf_layout.shape[1], + (sf_vec, sf_layout.shape[2])), + stride=(sf_layout.stride[0], sf_layout.stride[1], + (0, sf_layout.stride[2])), + ), + ) + else: + fc2_output_sf = None + if cutlass.const_expr(self.token_back_by_dispatch): fc2_output_workspace_native = self._view_local( local_workspace, "fc2_output_workspace", ) + # Byte count = elements * dtype.width // 8 (multiply before divide so + # the 4-bit fp4 data plane is not truncated to zero bytes/element). fc2_output_workspace_u8 = self._make_typed_view( local_workspace, self._local_offsets["fc2_output_workspace"], cutlass.Uint8, - (pool_token_capacity * self.hidden * - (int(self.fc2_output_dtype.width) // 8), ), + ((pool_token_capacity * self.hidden * + int(self.combine_format.act_dtype.width)) // 8, ), None, self._local_region_by_name["fc2_output_workspace"].align, ) - fc2_done_counter = self._view_local( - local_workspace, - "fc2_done_counter", - ) combine_output_u8 = cute.recast_tensor( - combine_output, + combine_target, cutlass.Uint8, ) else: fc2_output_workspace_native = None fc2_output_workspace_u8 = None + combine_output_u8 = combine_target + + # fc2 completion gate: present whenever token-back runs (DATA and/or SF + # push), so the epi_warps SF-only push gates per expert just like the + # dispatch DATA path. + if cutlass.const_expr(self.token_back_enabled): + fc2_done_counter = self._view_local( + local_workspace, + "fc2_done_counter", + ) + else: fc2_done_counter = None - combine_output_u8 = combine_output + if cutlass.const_expr(self.token_back_schedule_mode == + "atomic_counter"): # nosec B105 + token_back_schedule_counter = self._view_local( + local_workspace, + "token_back_schedule_counter", + ).iterator + else: + token_back_schedule_counter = None + + # Int32 views over each workspace's front counter prefix; the kernel tail + # bulk-zeros them every launch (tail_reset_counters). + local_zero_prefix = self._make_typed_view( + local_workspace, + 0, + cutlass.Int32, + (self.local_zero_i32_count, ), + (1, ), + 16, + ) + shared_zero_prefix = self._make_typed_view( + shared_workspace, + 0, + cutlass.Int32, + (self.shared_zero_i32_count, ), + (1, ), + 16, + ) token_comm_args = ExtractedTokenCommArgs( input_token_buffer=activation, input_sf_buffer=activation_sf, @@ -896,14 +1210,19 @@ def __call__( fc1_ready_counter=l1_arrival_count, token_src_metadata=token_src_metadata, combine_output=combine_output_u8, + combine_sf=combine_sf, fc2_output_workspace=fc2_output_workspace_u8, + fc2_output_sf=fc2_output_sf, fc2_done_counter=fc2_done_counter, + token_back_schedule_counter=token_back_schedule_counter, nvlink_barrier_signal=nvlink_barrier_signal, nvlink_barrier_counter=nvlink_barrier_counter, grid_sync_counter=grid_sync_counter, + local_zero_prefix=local_zero_prefix, + shared_zero_prefix=shared_zero_prefix, peer_rank_ptr_mapper=peer_rank_ptr_mapper, world_size=self.world_size, - local_rank=self.local_rank, + local_rank=peer_rank_ptr_mapper_host.rank_idx, num_total_experts=self.num_total_experts, num_experts_per_rank=self.num_experts_per_rank, num_topk=self.num_topk, @@ -918,16 +1237,22 @@ def __call__( # sf_padding_block == "sf_block_m") so the pool layout and the # sched cumulative-row offsets align by construction. # - # ``combine_output`` is MoE-domain storage. Non-reduce modes use - # ``(max_tokens_per_rank, num_topk, hidden)`` and host-reduce topk; - # REDG modes use ``(max_tokens_per_rank, 1, hidden)`` and reduce in - # kernel. The epilogue return tile maps local pool rows back to the - # source rank's token row through ``token_comm_args``. + # ``combine_target`` is MoE-domain storage. separate_kernel_reduce uses + # ``(max_tokens_per_rank, num_topk, hidden)`` (internal staging) and a + # tail reduce; in_kernel_reduce uses ``(max_tokens_per_rank, 1, hidden)`` + # and reduces in kernel. The epilogue return tile maps local pool rows + # back to the source rank's token row through ``token_comm_args``. if cutlass.const_expr(self.token_back_by_dispatch): fc2_output_target = fc2_output_workspace_native else: - fc2_output_target = combine_output - + fc2_output_target = combine_target + + # Quantized combine loop is closed: the fc2 epilogue encodes the data + # plane (STG-to-peer in epi_warps, or local staging in the dispatch + # modes) and writes per-block scales to the rank-local SF plane + # (fc2_output_sf); the token-back warps push that SF plane + # token-contiguously to the peers' combine_sf (and the data plane too in + # the dispatch modes); TopkReduce below dequantizes and reduces. super().__call__( activation=l1_token_buffer_nvfp4, fc1_weight=fc1_weight, @@ -951,6 +1276,26 @@ def __call__( token_comm_args=token_comm_args, ) + # separate_kernel_reduce: collapse the per-topk ``combine_quant`` staging + # into the public 2D output via the shared TopkReduce launcher (sized from + # codegen-time hidden / num_topk / combine_format -- it dequantizes per + # format and reduces over topk). Same stream, so it is ordered strictly + # after the cross-rank combine writes landed (the mega kernel's nvlink + # barrier guarantees all peer combine STGs to this rank completed before it + # exits). Weighting follows the compute graph: deepgemm (apply_topk_in_fc1) + # folded the routing weight into fc1 -> plain K-sum; transformers applies + # topk_weights here. + if cutlass.const_expr(not self.in_kernel_fc2_reduce): + score = (topk_weights if + cutlass.const_expr(not self.apply_topk_in_fc1) else None) + TopkReduce(self.hidden, self.num_topk, self.combine_format)( + combine_target, + combine_sf, + output_activation, + score, + stream, + ) + # ========================================================================= # TokenComm delegation surface consumed by the fc1/fc2 base kernel # ========================================================================= @@ -994,6 +1339,24 @@ def token_comm_hook_dispatch_warp_body( tidx=tidx, ) + @cute.jit + def token_comm_hook_token_back_warp_body( + self, + token_comm_args, + token_comm_storage, + *, + warp_idx, + lane_idx, + tidx, + ): + self.token_comm.token_back_warp_body( + token_comm_args, + token_comm_storage, + warp_idx=warp_idx, + lane_idx=lane_idx, + tidx=tidx, + ) + @cute.jit def token_comm_hook_tail_reset_shared_counters( self, @@ -1003,8 +1366,9 @@ def token_comm_hook_tail_reset_shared_counters( local_warp_idx, lane_idx, ): - self.token_comm.tail_reset_shared_counters( + self.token_comm.tail_reset_counters( token_comm_args, + token_comm_args.shared_zero_prefix, cta_linear_id=cta_linear_id, local_warp_idx=local_warp_idx, lane_idx=lane_idx, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py index 037bb71b3101..a4e223eb0fcf 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/moe_utils.py @@ -62,7 +62,8 @@ def spin_wait( ip: Optional[ir.InsertionPoint] = None, ) -> Boolean: """Spin until condition is true, or do one condition check with peek_only.""" - current = cute.arch.load(ptr, ptr.dtype, cop="cg", loc=loc, ip=ip) + # The first load must be acquire to build the mem order. + current = cute.arch.load(ptr, ptr.dtype, sem="acquire", scope="gpu") if cutlass.const_expr(peek_only): # One-shot peek: forward the condition Boolean to the caller. return Boolean(condition(current)) @@ -70,7 +71,7 @@ def spin_wait( # Load with L1 cache bypass (ld.global.cg) if cutlass.const_expr(fail_sleep_cycles > 0): _nanosleep(fail_sleep_cycles, loc=loc, ip=ip) - current = cute.arch.load(ptr, ptr.dtype, cop="cg", loc=loc, ip=ip) + current = cute.arch.load(ptr, ptr.dtype, sem="acquire", scope="gpu") # Spin-path: condition was satisfied; uniformize return type with the # peek path so callers always see a Boolean. return Boolean(True) @@ -80,8 +81,7 @@ def spin_wait( # Cluster-DSMEM helpers (for atomic_counter dynamic scheduler) # ============================================================================= # -# Ported from cute_dsl_kernel_library/dsl_kernels/moe/moe_persistent_scheduler.py -# (lines 79-145). Used by the fused fc1+fc2 mega scheduler when +# Used by the fused fc1+fc2 mega scheduler when # load_balance_mode == 'atomic_counter' to # implement the leader-CTA atom.add + DSMEM broadcast cluster-tile-idx # fetch protocol. ``atom.add`` itself uses cute.arch.atomic_add (the diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py index 0a6aafd61097..adcb04c8c1d7 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/ptx_helpers.py @@ -1,7 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Inline-PTX wrappers (TMA 1D load/store, fns.b32, raw-int64 peer ops) for the cuTeDSL dispatch kernel.""" +"""Inline-PTX wrappers (TMA 1D load/store, fns.b32, raw-int64 peer ops) for the cuTeDSL kernels.""" +from typing import Optional + +from cutlass._mlir import ir from cutlass._mlir.dialects import llvm from cutlass.cutlass_dsl import Float32, Int32, Int64, T, dsl_user_op @@ -92,6 +95,63 @@ def tma_store_1d(dst_gmem, src_smem, num_bytes, *, loc=None, ip=None): ) +@dsl_user_op +def cp_async_bulk_s2g(dst_gmem, + src_smem, + size_bytes, + *, + loc=None, + ip=None) -> None: + """Issue descriptor-free ``cp.async.bulk`` SMEM->GMEM. + + The caller owns ``cp_async_bulk_commit_group`` so copy and reduce bulk + paths can share the same group boundary. + """ + llvm.inline_asm( + None, + [ + dst_gmem.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip), + src_smem.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip), + size_bytes.ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.global.shared::cta.bulk_group [$0], [$1], $2;", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def cp_reduce_async_bulk_add_noftz_bf16_s2g( + dst_gmem, + src_smem, + size_bytes, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + """Issue descriptor-free ``cp.reduce.async.bulk`` for BF16 add.""" + llvm.inline_asm( + None, + [ + dst_gmem.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip), + src_smem.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip), + size_bytes.ir_value(loc=loc, ip=ip), + ], + "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.noftz.bf16 " + "[$0], [$1], $2;", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + @dsl_user_op def fns_b32(mask: Int32, base: Int32, n: Int32, *, loc=None, ip=None) -> Int32: return Int32( @@ -204,6 +264,31 @@ def red_add_release_sys_u64_raw(addr: Int64, ) +@dsl_user_op +def red_add_relaxed_sys_u64_raw(addr: Int64, + val: Int64, + *, + loc=None, + ip=None) -> None: + """``red.relaxed.sys.global.add.u64`` via raw int64 byte address. + + Relaxed (no release fence) sibling of ``red_add_release_sys_u64_raw``. Use + when an enclosing ``bar.sync`` + a later release fence (e.g. the trailing + ``nvlink_barrier`` publish) already provides cross-rank ordering, so the + per-op release fence would just emit a redundant ``membar.sys`` drain. + """ + llvm.inline_asm( + None, + [addr.ir_value(), val.ir_value()], + "red.relaxed.sys.global.add.u64 [$0], $1;", + "l,l", + has_side_effects=True, + asm_dialect=0, + loc=loc, + ip=ip, + ) + + @dsl_user_op def red_add_release_sys_s32_raw(addr: Int64, val: Int32, @@ -227,6 +312,58 @@ def red_add_release_sys_s32_raw(addr: Int64, ) +@dsl_user_op +def red_add_relaxed_sys_v2_bf16x2( + addr, + val0_packed_bf16x2, + val1_packed_bf16x2, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + """Issue ``red.relaxed.sys.global.add.v2.bf16x2 [addr], {v0, v1};``.""" + llvm.inline_asm( + None, + [ + addr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip), + val0_packed_bf16x2.ir_value(loc=loc, ip=ip), + val1_packed_bf16x2.ir_value(loc=loc, ip=ip), + ], + "red.relaxed.sys.global.add.noftz.v2.bf16x2 [$0], {$1, $2};", + "l,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def red_add_release_gpu_s32( + counter_ptr, + value, + *, + loc: Optional[ir.Location] = None, + ip: Optional[ir.InsertionPoint] = None, +) -> None: + """Issue ``red.release.gpu.global.add.s32`` to a GMEM int32 location.""" + llvm.inline_asm( + None, + [ + counter_ptr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip), + value.ir_value(loc=loc, ip=ip), + ], + "red.release.gpu.global.add.s32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + @dsl_user_op def red_async_add_release_sys_u32_raw(addr: Int64, val: Int32, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py index 33af2cfc5380..c63ac97d5fa0 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/sym_buffer.py @@ -2,21 +2,28 @@ # SPDX-License-Identifier: Apache-2.0 """Symmetric-heap peer pointer mapper. -``SymBufferHost`` is the runtime payload that crosses the Python -> -generated-host-code boundary. Inside the generated host wrapper, it packs -the runtime base address and per-rank offsets into a device-side -``SymBuffer{N}`` native struct: - - { i64 base, vector offsets, i32 rank_idx } - -Device code only sees that struct and calls ``.map`` / -``.ptr_map_to_rank``. The vector field is deliberate: LLVM supports -runtime-indexed ``extractelement`` on vectors, which NVPTX lowers to an -indexed param-bank load (``LDC.U64``). +The device-side SymBuffer carries ONLY the per-rank offset table. Layout is +chosen at trace time by ``num_max_ranks`` (a constexpr): + +* ``<= 16`` ranks: a ``cute_nvgpu.grid_constant`` / ``llvm.byval`` pointer to a + ``struct<(array<16 x i64>)>`` (exactly 128B). ``map`` does a runtime-indexed + ``getelementptr`` + ``load`` that lowers to a const-bank ``LDC c[bnk][Ra]``. + The 128B byval size matches the host-side ``createKernelArgs`` copy (hardcoded + to 128B today), so it is corruption-safe up to EP16. + +* ``> 16`` ranks: a by-value ``vector`` read with ``extractelement``. + Correct but spills the runtime index -- the fallback until the byval-size fix + reaches public cutedsl. + +``base`` is intentionally dropped (``get_base_ptr`` had no callers) and +``rank_idx`` lives outside the offset table: ``_SymBufferHostAdapter`` +marshals it as a separate i32 scalar and the kernel consumes it as the device +``local_rank`` (see ``megamoe_kernel.py``). Freeing those struct slots lets +all 16 byval lanes hold offsets (EP16, not EP14). """ from dataclasses import dataclass -from typing import Any, Tuple +from typing import Any, Optional import cutlass import cutlass.cute as cute @@ -24,16 +31,64 @@ from cutlass._mlir.dialects import arith, llvm from cutlass.base_dsl.dsl import (extract_mlir_values, get_mlir_types, new_from_mlir_values) -from cutlass.base_dsl.native_struct import native_struct from cutlass.base_dsl.runtime.jit_arg_adapters import JitArgAdapterRegistry from cutlass.base_dsl.typing import get_c_pointers from cutlass.cute.typing import AddressSpace from cutlass.cutlass_dsl import Int32, Int64, dsl_user_op +try: + # GEP encodes a runtime index by storing this sentinel in rawConstantIndices; + # it is MLIR's LLVM::GEPOp::kDynamicIndex (== INT32_MIN), frozen by the IR + # encoding ABI. Track the canonical constant rather than re-hardcoding it. + from cutlass.base_dsl.typing import MLIR_DYNAMIC_INDEX +except ImportError: # older wheels: value is fixed by the GEP encoding ABI + MLIR_DYNAMIC_INDEX = -(2**31) + +_BYVAL_RANK_LIMIT = 16 # struct<(array<16 x i64>)> == exactly 128B + + +# TODO: Remove once the compiler is fixed. Workaround for a cuda-to-llvm bug: +# any kernel arg marked `grid_constant + byval` is treated as a tma_desc. +def _byval_struct_ty() -> Any: + """128B byval pointee shared by alloca / GEP / the ``llvm.byval`` attr.""" + return ir.Type.parse(f"!llvm.struct<(array<{_BYVAL_RANK_LIMIT} x i64>)>") + @dataclass(frozen=True) class SymBufferDeviceBase: - """Device-side methods shared by all generated ``SymBuffer{N}`` types.""" + """Device-side SymBuffer: the per-rank offset table only. + + ``val`` is a ``!llvm.ptr`` to a byval/grid_constant ``struct<(array<16 x i64>)>`` + when ``num_max_ranks <= 16`` (``map`` -> GEP + load -> ``LDC``), else a by-value + ``vector`` (``map`` -> ``extractelement``, spills). + """ + + val: Any + num_max_ranks: cutlass.Constexpr[int] + + def __extract_mlir_values__(self) -> list: + return [self.val] + + def __new_from_mlir_values__(self, values: list) -> "SymBufferDeviceBase": + return SymBufferDeviceBase(val=values[0], + num_max_ranks=self.num_max_ranks) + + def __get_mlir_types__(self) -> list: + if self.num_max_ranks <= _BYVAL_RANK_LIMIT: + return [ir.Type.parse("!llvm.ptr")] + return [ir.Type.parse(f"vector<{self.num_max_ranks}xi64>")] + + def __extract_mlir_attributes__(self) -> list: + if self.num_max_ranks <= _BYVAL_RANK_LIMIT: + return [ + ir.DictAttr.get({ + "cute_nvgpu.grid_constant": + ir.UnitAttr.get(), + "llvm.byval": + ir.TypeAttr.get(_byval_struct_ty()), + }) + ] + return [ir.DictAttr.get({})] @cute.jit def map( @@ -42,84 +97,103 @@ def map( dst_rank_idx: Int32, byte_off: Int64 = Int64(0), ) -> Int64: - off = Int64(llvm.extractelement(self.offsets, dst_rank_idx.ir_value())) + if cutlass.const_expr(self.num_max_ranks <= _BYVAL_RANK_LIMIT): + # Opaque ptr -> the offsets array sits at byte 0 of the byval struct, + # so a flat ``gep i64, ptr, dst_rank`` reaches offsets[dst_rank] + # directly; the byval struct type only governs the 128B const-bank copy. + i64_ty = ir.Type.parse("i64") + off_ptr = llvm.getelementptr( + ir.Type.parse("!llvm.ptr"), + self.val, + [dst_rank_idx.ir_value()], + [MLIR_DYNAMIC_INDEX], + i64_ty, + no_wrap_flags="None", + ) + off = Int64(llvm.load(i64_ty, off_ptr)) + else: + off = Int64(llvm.extractelement(self.val, dst_rank_idx.ir_value())) return local_ptr + off + byte_off @cute.jit - def get_base_ptr(self) -> Int64: - return self.base - - @cute.jit - def ptr_map_to_rank(self, ptr, dst_rank_idx: Int32): + def ptr_map_to_rank(self, + ptr, + dst_rank_idx: Int32, + byte_align: Optional[int] = None): if cutlass.const_expr(ptr.memspace != AddressSpace.gmem): raise ValueError( f"ptr_map_to_rank: source pointer must live in GMEM " f"(NVSHMEM symmetric heap), got memspace={ptr.memspace}.") + if cutlass.const_expr(byte_align is None): + byte_align = ptr.max_alignment peer_addr = self.map(ptr.toint(), dst_rank_idx, Int64(0)) return cute.make_ptr( ptr.dtype, peer_addr, ptr.memspace, - assumed_align=ptr.max_alignment, + assumed_align=byte_align, ) @dataclass(frozen=True) class SymBufferHost: - """Runtime launch payload for a device-side ``SymBuffer{N}``.""" + """Runtime launch payload for a device-side ``SymBuffer{N}``. + + Marshalled across the JIT boundary by ``_SymBufferHostAdapter`` (registered + below), which re-wraps the scalar fields with ``Int64(...)``.""" - base_addr: int - offsets: Tuple[int, ...] - rank_idx: int + offsets: tuple + rank_idx: Int32 num_max_ranks: cutlass.Constexpr[int] @staticmethod def _as_int64(value) -> Int64: return value if isinstance(value, Int64) else Int64(int(value)) - @staticmethod - def _as_int32(value) -> Int32: - return value if isinstance(value, Int32) else Int32(int(value)) - - @staticmethod - def _make_device_type(num_max_ranks: int) -> type: - if num_max_ranks <= 0: - raise ValueError( - f"num_max_ranks must be positive, got {num_max_ranks}") - - vec_ty_str = f"vector<{num_max_ranks}xi64>" - - class _OffsetsT: - - @staticmethod - def mlir_type() -> ir.Type: - return ir.Type.parse(vec_ty_str) - - @native_struct - class _SymBufferDevice(SymBufferDeviceBase): - base: Int64 - offsets: _OffsetsT - rank_idx: Int32 - - cls = _SymBufferDevice - cls.__name__ = f"SymBuffer{num_max_ranks}" - cls.__qualname__ = cls.__name__ - cls.NUM_MAX_RANKS = num_max_ranks - return cls - @dsl_user_op def make_device_obj(self, *, loc=None, ip=None) -> Any: + """Build the offsets-only device obj (see module docstring for layout).""" offsets = tuple(self.offsets) num_max_ranks = self.num_max_ranks if len(offsets) != num_max_ranks: - raise ValueError( - f"len(offsets)={len(offsets)} must equal " - f"num_max_ranks={num_max_ranks}; SymBuffer requires its " - f"runtime payload length to match the compiled vector type.") + raise ValueError(f"len(offsets)={len(offsets)} must equal " + f"num_max_ranks={num_max_ranks}.") + + if num_max_ranks <= _BYVAL_RANK_LIMIT: + ptr_ty = ir.Type.parse("!llvm.ptr") + st_ty = _byval_struct_ty() + i64_ty = ir.Type.parse("i64") + one = arith.constant( + value=ir.IntegerAttr.get(i64_ty, 1), + result=i64_ty, + loc=loc, + ip=ip, + ) + buf = llvm.alloca( + res=ptr_ty, + elem_type=st_ty, + array_size=one, + alignment=64, + loc=loc, + ip=ip, + ) + for i, off in enumerate(offsets): + slot = llvm.getelementptr( + ptr_ty, + buf, + [], + [i], + i64_ty, + no_wrap_flags="None", + loc=loc, + ip=ip, + ) + llvm.store(self._as_int64(off).ir_value(), slot, loc=loc, ip=ip) + return SymBufferDeviceBase(val=buf, num_max_ranks=num_max_ranks) + i32_ty = ir.Type.parse("i32") vec_ty = ir.Type.parse(f"vector<{num_max_ranks}xi64>") vec = llvm.mlir_zero(vec_ty, loc=loc, ip=ip) - i32_ty = ir.Type.parse("i32") for i, off in enumerate(offsets): idx = arith.constant( value=ir.IntegerAttr.get(i32_ty, i), @@ -134,14 +208,7 @@ def make_device_obj(self, *, loc=None, ip=None) -> Any: loc=loc, ip=ip, ) - - return self._make_device_type(num_max_ranks)( - base=self._as_int64(self.base_addr), - offsets=vec, - rank_idx=self._as_int32(self.rank_idx), - loc=loc, - ip=ip, - ) + return SymBufferDeviceBase(val=vec, num_max_ranks=num_max_ranks) @JitArgAdapterRegistry.register_jit_arg_adapter(SymBufferHost) @@ -150,7 +217,7 @@ class _SymBufferHostAdapter: Python-side ``SymBufferHost`` stays pure host data (ints + tuple). The adapter is the only place that maps it to DSL scalar arguments: - base/offsets are i64, rank_idx is i32, and num_max_ranks remains a + offsets are i64, rank_idx is i32, and num_max_ranks remains a constexpr carried through reconstruction. """ @@ -161,7 +228,6 @@ def __init__(self, arg: SymBufferHost) -> None: f"len(offsets)={len(tuple(arg.offsets))} must equal " f"num_max_ranks={int(arg.num_max_ranks)}.") self._fields = ( - Int64(arg.base_addr), *(Int64(x) for x in arg.offsets), Int32(arg.rank_idx), ) @@ -187,13 +253,8 @@ def __extract_mlir_values__(self) -> list[ir.Value]: def __new_from_mlir_values__(self, values: list[ir.Value]) -> SymBufferHost: idx = 0 - base_n = len(get_mlir_types(self._fields[0])) - base_addr = new_from_mlir_values(self._fields[0], - values[idx:idx + base_n]) - idx += base_n - offsets = [] - for field in self._fields[1:-1]: + for field in self._fields[:-1]: n = len(get_mlir_types(field)) offsets.append(new_from_mlir_values(field, values[idx:idx + n])) idx += n @@ -208,7 +269,6 @@ def __new_from_mlir_values__(self, values: list[ir.Value]) -> SymBufferHost: ) obj = object.__new__(SymBufferHost) - object.__setattr__(obj, "base_addr", base_addr) object.__setattr__(obj, "offsets", tuple(offsets)) object.__setattr__(obj, "rank_idx", rank_idx) object.__setattr__(obj, "num_max_ranks", self._arg.num_max_ranks) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/token_comm.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/token_comm.py index d0c70b730ddf..fd67dd4828bf 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/token_comm.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/token_comm.py @@ -6,11 +6,14 @@ ``dispatch_kernel`` uses the same object methods as the fused MegaMoE kernel. """ -from typing import Any, Dict, List +import dataclasses +from typing import Any, ClassVar, Dict, List, Literal, Optional, Union import cutlass import cutlass.cute as cute import cutlass.pipeline as pipeline +from cutlass.base_dsl.dsl import extract_mlir_attributes +from cutlass.cute.typing import AddressSpace from cutlass.cutlass_dsl import (Float32, Int32, Int64, Uint8, Uint32, extract_mlir_values, new_from_mlir_values) @@ -27,27 +30,160 @@ from cutlass._mlir import ir +from .flag_batch import GpuReleaseFlagBatchTracker from .grid_sync import software_grid_sync -from .moe_utils import spin_wait -from .ptx_helpers import (fns_b32, ldg_b32_raw, ldg_f32_raw, - red_add_release_sys_s32_raw, - red_add_release_sys_u64_raw, stg_b32_raw, stg_b64_raw, +from .moe_utils import _nanosleep, spin_wait +from .ptx_helpers import (cp_reduce_async_bulk_add_noftz_bf16_s2g, fns_b32, + ldg_b32_raw, ldg_f32_raw, read_clock64, + red_add_relaxed_sys_u64_raw, + red_add_release_sys_s32_raw, stg_b32_raw, stg_b64_raw, tma_load_1d_raw, tma_store_1d) from .sf_swizzle import sf_atom_int32_offset +# --------------------------------------------------------------------------- +# Low-precision combine wire format (the central driver for the token-back +# quantized combine path: fc2 epilogue encoder, token_comm push, and the +# topk_reduce receiver all describe themselves through one CombineFormat). +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class CombineFormat: + """Wire format of the cross-rank combine (token-back) payload. + + The fc2 epilogue quantizes each token's hidden vector into a packed data + plane (``combine_quant``) plus a per-block scale plane (``combine_sf``); the + receiver (topk_reduce) dequantizes and reduces over topk. The consumer is + ALWAYS a software dequant -- never a tensor-core MMA -- so the format is + fully self-defined here, with no hardware scale-layout constraint. + + Canonical string ``"{scale_block}{act}x{scale}"`` (leading number = scale + block size in hidden elements), e.g. ``"16e2m1xbf16"`` (per-16 bf16 amax + + fp4 data) or ``"32e4m3xe8m0"`` (standard MXFP8). ``"bf16"`` is the + no-staging baseline: bf16 fc2 terms reduced directly, no scale plane. + """ + + # Element/scale tag <-> cuTe dtype. Only the dtypes a real format uses + # today (the bf16 baseline is ``scale_dtype is None``); extend when a new + # format is actually added rather than ahead of need. + _act_by_tag: ClassVar[Dict[str, type]] = { + "e2m1": cutlass.Float4E2M1FN, + "e4m3": cutlass.Float8E4M3FN, + } + _scale_by_tag: ClassVar[Dict[str, type]] = { + "bf16": cutlass.BFloat16, + "e8m0": cutlass.Float8E8M0FNU, + } + + act_dtype: type # cuTe dtype of the packed data-plane element + scale_dtype: Optional[ + type] # cuTe dtype of a scale entry; None == bf16 baseline + scale_block: Optional[ + int] # hidden elements per scale entry; None == baseline + + def __post_init__(self): + allowed_act = {cutlass.BFloat16, *self._act_by_tag.values()} + if self.act_dtype not in allowed_act: + raise ValueError( + f"combine act_dtype {self.act_dtype} not in {allowed_act}.") + allowed_scale = {None, *self._scale_by_tag.values()} + if self.scale_dtype not in allowed_scale: + raise ValueError( + f"combine scale_dtype {self.scale_dtype} not in {allowed_scale}." + ) + if self.scale_dtype is None: # bf16 no-staging baseline + if self.act_dtype is not cutlass.BFloat16 or self.scale_block is not None: + raise ValueError( + "baseline must be bf16 act with scale_block=None.") + return + if self.act_dtype is cutlass.BFloat16: + raise ValueError("a quantized combine cannot use a bf16 act dtype.") + # The scale dtype pins the block: per-16 bf16 amax / per-32 e8m0 power-of-two. + if self.scale_dtype is cutlass.BFloat16 and self.scale_block != 16: + raise ValueError("bf16 amax scale requires scale_block == 16.") + if self.scale_dtype is cutlass.Float8E8M0FNU and self.scale_block != 32: + raise ValueError("e8m0 scale requires scale_block == 32.") -def _store_token_src_metadata_u32x3( - token_src_metadata, - pool_token_idx, - src_rank: Uint32, - src_token: Uint32, - src_topk: Uint32, -) -> None: - """Store `{src_rank, src_token, src_topk}` as three 32-bit fields.""" - base_ptr = token_src_metadata.iterator + (pool_token_idx * Int32(12)) - cute.arch.store(base_ptr, src_rank, scope="gpu") - cute.arch.store(base_ptr + Int32(4), src_token, scope="gpu") - cute.arch.store(base_ptr + Int32(8), src_topk, scope="gpu") + @property + def is_quantized(self) -> bool: + """``False`` for the bf16 (no-staging) baseline.""" + return self.scale_dtype is not None + + @property + def name(self) -> str: + if not self.is_quantized: + return "bf16" + act_tag = next(t for t, d in self._act_by_tag.items() + if d is self.act_dtype) + scale_tag = next(t for t, d in self._scale_by_tag.items() + if d is self.scale_dtype) + return f"{self.scale_block}{act_tag}x{scale_tag}" + + def __str__(self) -> str: + return self.name + + @classmethod + def parse(cls, text: str) -> "CombineFormat": + """Build a CombineFormat from its canonical string (the argparser entry). + + Only the handful of supported wire formats are accepted; each key is the + exact string ``name`` produces (so ``parse(str(fmt)) == fmt``). + """ + # (act_dtype, scale_dtype, scale_block); None scale == bf16 baseline. + specs = { + "bf16": (cutlass.BFloat16, None, None), + "16e2m1xbf16": (cutlass.Float4E2M1FN, cutlass.BFloat16, 16), + "32e4m3xe8m0": (cutlass.Float8E4M3FN, cutlass.Float8E8M0FNU, 32), + } + token = text.strip().lower() + if token not in specs: + raise ValueError( + f"invalid combine_format {text!r}: expected one of {tuple(specs)}." + ) + act_dtype, scale_dtype, scale_block = specs[token] + return cls(act_dtype=act_dtype, + scale_dtype=scale_dtype, + scale_block=scale_block) + + +@dataclasses.dataclass(frozen=True) +class TokenSrcMetadata: + """Per pool-token routing record: written by token-in, read by token-back + and the fc2 combine-redirect epilogue. + + Wire format is one i64: low 32b = ``src_token`` (needs full width); high 32b + = ``(src_rank << 16) | src_topk`` (``src_rank < world_size`` and + ``src_topk < num_topk`` both fit in 16b). ``load`` / ``store`` accept either + a ``cute.Pointer`` or a raw ``Int64`` byte address. + """ + + src_rank: Int32 + src_token: Int32 + src_topk: Int32 + + nbytes: ClassVar[int] = 8 + + def _pack(self) -> Int64: + hi = (Int64(self.src_rank) << Int64(16)) | Int64(self.src_topk) + return (hi << Int64(32)) | (Int64(self.src_token) & Int64(0xFFFFFFFF)) + + @staticmethod + def _i64_ptr(addr: Union[cute.Pointer, Int64]) -> cute.Pointer: + addr_i = addr if isinstance(addr, Int64) else addr.toint() + return cute.make_ptr(Int64, addr_i, AddressSpace.gmem, assumed_align=8) + + def store(self, addr: Union[cute.Pointer, Int64]) -> None: + cute.arch.store(self._i64_ptr(addr), self._pack(), scope="gpu") + + @classmethod + def load(cls, addr: Union[cute.Pointer, Int64]) -> "TokenSrcMetadata": + v = Int64(cute.arch.load(cls._i64_ptr(addr), Int64, scope="gpu")) + hi = v >> Int64(32) + return cls( + src_rank=Int32((hi >> Int64(16)) & Int64(0xFFFF)), + src_token=Int32(v & Int64(0xFFFFFFFF)), + src_topk=Int32(hi & Int64(0xFFFF)), + ) _MLIR_VALUE_FIELDS = ( @@ -65,17 +201,22 @@ def _store_token_src_metadata_u32x3( "fc1_ready_counter", "token_src_metadata", "combine_output", + "combine_sf", "fc2_output_workspace", + "fc2_output_sf", "fc2_done_counter", + "token_back_schedule_counter", "nvlink_barrier_signal", "nvlink_barrier_counter", "grid_sync_counter", + "local_zero_prefix", + "shared_zero_prefix", "peer_rank_ptr_mapper", + "local_rank", ) _CONST_FIELDS = ( "world_size", - "local_rank", "num_total_experts", "num_experts_per_rank", "num_topk", @@ -110,6 +251,8 @@ def __init__( nvlink_barrier_signal: cute.Tensor, nvlink_barrier_counter: cute.Tensor, grid_sync_counter: cute.Tensor, + local_zero_prefix: cute.Tensor, + shared_zero_prefix: cute.Tensor, peer_rank_ptr_mapper: Any, world_size: int, local_rank: int, @@ -123,6 +266,9 @@ def __init__( sm_count: int, fc2_output_workspace: cute.Tensor = None, fc2_done_counter: cute.Tensor = None, + token_back_schedule_counter: cute.Pointer = None, + combine_sf: cute.Tensor = None, + fc2_output_sf: cute.Tensor = None, ): self.input_token_buffer = input_token_buffer self.input_sf_buffer = input_sf_buffer @@ -138,11 +284,16 @@ def __init__( self.fc1_ready_counter = fc1_ready_counter self.token_src_metadata = token_src_metadata self.combine_output = combine_output + self.combine_sf = combine_sf self.fc2_output_workspace = fc2_output_workspace + self.fc2_output_sf = fc2_output_sf self.fc2_done_counter = fc2_done_counter + self.token_back_schedule_counter = token_back_schedule_counter self.nvlink_barrier_signal = nvlink_barrier_signal self.nvlink_barrier_counter = nvlink_barrier_counter self.grid_sync_counter = grid_sync_counter + self.local_zero_prefix = local_zero_prefix + self.shared_zero_prefix = shared_zero_prefix self.peer_rank_ptr_mapper = peer_rank_ptr_mapper self.world_size = world_size self.local_rank = local_rank @@ -164,6 +315,17 @@ def __extract_mlir_values__(self) -> List[ir.Value]: values.extend(extract_mlir_values(attr)) return values + def __extract_mlir_attributes__(self) -> List[Any]: + # Mirror __extract_mlir_values__ 1:1 so per-arg attrs stay aligned; the + # only non-empty entry is peer_rank_ptr_mapper's byval/grid_constant. + attrs: List[Any] = [] + for name in _MLIR_VALUE_FIELDS: + attr = getattr(self, name) + if attr is None: + continue + attrs.extend(extract_mlir_attributes(attr)) + return attrs + def __new_from_mlir_values__(self, values: List[ir.Value]) -> "TokenCommArgs": idx = 0 @@ -176,9 +338,8 @@ def __new_from_mlir_values__(self, n = len(extract_mlir_values(proto)) rebuilt[name] = new_from_mlir_values(proto, values[idx:idx + n]) idx += n - assert idx == len(values), ( - f"TokenCommArgs serialization mismatch: consumed={idx} provided={len(values)}" - ) + assert idx == len(values), (f"TokenCommArgs serialization mismatch: " + f"consumed={idx} provided={len(values)}") const_kwargs = {name: getattr(self, name) for name in _CONST_FIELDS} return TokenCommArgs(**rebuilt, **const_kwargs) @@ -192,14 +353,13 @@ class TokenInPullTokenBackPush: dispatch_intra_cta_bar_id: int = 10 kernel_tail_named_barrier_id: int = 8 dispatch_to_sched_named_barrier_id: int = 9 - dispatch_to_sched_threads: int = (num_dispatch_warps + 1) * warp_threads + # dispatch_to_sched / kernel_tail thread counts are per-instance (see __init__). experts_per_dispatch_pass: int = num_dispatch_threads def __init__( self, *, world_size: int, - local_rank: int, num_topk: int, num_experts_per_rank: int, num_total_experts: int, @@ -212,11 +372,17 @@ def __init__( cluster_shape_mn, dispatch_warp_start: int, num_other_warps: int, - fc2_output_dtype=None, + combine_format: "CombineFormat" = None, + token_back_by_dispatch: bool = False, fc2_publishes_per_token_cluster_tile: int = 0, + token_back_reduce_topk: bool = False, + token_back_standalone: bool = False, + flag_batch: int = 1, + is_swap_ab: bool = False, + token_back_schedule_mode: Literal["static", + "atomic_counter"] = "static", ) -> None: self.world_size = world_size - self.local_rank = local_rank self.num_topk = num_topk self.num_experts_per_rank = num_experts_per_rank self.num_total_experts = num_total_experts @@ -228,6 +394,19 @@ def __init__( self.sf_padding_block = sf_padding_block self.cluster_tile_tokens = cluster_tile_tokens self.cluster_shape_mn = cluster_shape_mn + if flag_batch < 1 or flag_batch > 32: + raise ValueError( + f"flag_batch must be in [1, 32], got {flag_batch}.") + # Release-flag batch size consumed by dispatch_pull as a Python int. + # One warp lane carries one delayed release target. + self._flag_batch = flag_batch + self.is_swap_ab = is_swap_ab + + if token_back_schedule_mode not in ("static", "atomic_counter"): + raise ValueError( + "token_back_schedule_mode must be 'static' or " + f"'atomic_counter'; got {token_back_schedule_mode!r}.") + self.token_back_schedule_mode = token_back_schedule_mode self.dispatch_warp_start = dispatch_warp_start # Warps that share this CTA with the dispatch group but are not part # of it. They participate in kernel-tail / dispatch-with-other @@ -236,12 +415,53 @@ def __init__( # collapse to dispatch-only). self.num_other_warps = num_other_warps self.num_other_threads = num_other_warps * self.warp_threads - self.num_total_threads = self.num_dispatch_threads + self.num_other_threads + + if combine_format is None: + combine_format = CombineFormat( + act_dtype=cutlass.BFloat16, + scale_dtype=None, + scale_block=None, + ) + self.combine_format = combine_format + self.token_back_by_dispatch = token_back_by_dispatch + self.push_data = token_back_by_dispatch + self.push_sf = combine_format.is_quantized + + # Standalone token-back: a dedicated warpgroup (size == dispatch group) + self.token_back_standalone = token_back_standalone + self.num_token_back_warps = self.num_dispatch_warps if self.token_back_standalone else 0 + self.num_token_back_threads = self.num_token_back_warps * self.warp_threads + self.token_back_warp_start = dispatch_warp_start + self.num_dispatch_warps + # Standalone token-back per-warp pull buffer; token is moved in + # tb_chunk_bytes pieces (last piece carries the remainder), so this is + # independent of hidden. + self.tb_chunk_bytes = 2048 + + self.num_total_threads = (self.num_dispatch_threads + + self.num_other_threads + + self.num_token_back_threads) + self.dispatch_to_sched_threads = ( + self.num_dispatch_warps + 1 + + self.num_token_back_warps) * self.warp_threads self.kernel_tail_threads = self.num_total_threads - self.fc2_output_dtype = fc2_output_dtype - if fc2_output_dtype is not None: - self.fc2_token_bytes = hidden * int(fc2_output_dtype.width) // 8 + # The DATA wire dtype is the combine act dtype (bf16 baseline -> bf16; + # fp4/e4m3 quantized), NOT the kernel's fc2 output dtype: the cross-rank + # payload is what the receiver dequantizes. + self.fc2_output_dtype = combine_format.act_dtype + if token_back_reduce_topk: + if not token_back_by_dispatch: + raise ValueError( + "token_back_reduce_topk=True requires the dispatch " + "token-back DATA path (token_back_by_dispatch).") + if combine_format.act_dtype is not cutlass.BFloat16: + raise NotImplementedError( + "token_back_reduce_topk currently supports a bf16 combine " + f"only, got {combine_format}.") + self.token_back_reduce_topk = token_back_reduce_topk + if self.enable_token_back: + self.fc2_token_bytes = hidden * int( + combine_format.act_dtype.width) // 8 if self.fc2_token_bytes % self.hidden_bytes != 0: raise ValueError( f"fc2_token_bytes={self.fc2_token_bytes} must be a " @@ -251,8 +471,9 @@ def __init__( if fc2_publishes_per_token_cluster_tile <= 0: raise ValueError( "fc2_publishes_per_token_cluster_tile must be > 0 when " - "fc2_output_dtype is set (token_back_by_push enabled).") - self.fc2_publishes_per_token_cluster_tile = fc2_publishes_per_token_cluster_tile + "token-back is enabled (it gates the per-expert push).") + self.fc2_publishes_per_token_cluster_tile = ( + fc2_publishes_per_token_cluster_tile) else: self.fc2_token_bytes = 0 self.fc2_num_chunks = 0 @@ -260,12 +481,29 @@ def __init__( @property def enable_token_back(self) -> bool: - return self.fc2_output_dtype is not None + # token-back warps run if they push the DATA plane, the SF plane, or both. + return self.push_data or self.push_sf def extra_smem_storage_class(self) -> type: hidden_bytes = self.hidden_bytes num_total_experts = self.num_total_experts + if self.token_back_standalone: + + @cute.struct + class TokenCommStorage: + pull_mbar: cute.struct.MemRange[Int64, self.num_dispatch_warps] + smem_expert_count: cute.struct.MemRange[Int32, + num_total_experts] + pull_buffer: cute.struct.Align[cute.struct.MemRange[ + Uint8, self.num_dispatch_warps * hidden_bytes], 16] + tb_pull_mbar: cute.struct.MemRange[Int64, + self.num_token_back_warps] + tb_pull_buffer: cute.struct.Align[cute.struct.MemRange[ + Uint8, self.num_token_back_warps * self.tb_chunk_bytes], 16] + + return TokenCommStorage + @cute.struct class TokenCommStorage: pull_mbar: cute.struct.MemRange[Int64, self.num_dispatch_warps] @@ -288,14 +526,22 @@ def sched_warp_pre_init_wait(self, token_comm_args): @cute.jit def fc1_tma_b_predispatch_spin(self, token_comm_args, work_tile_info): - counter_slot = work_tile_info.cumulative_token_block_count + work_tile_info.tile_n_idx + if cutlass.const_expr(self.is_swap_ab): + counter_slot = work_tile_info.cumulative_token_block_count + work_tile_info.tile_n_idx + peek_threshold = work_tile_info.valid_tokens_in_cta_tile + else: + counter_slot = (work_tile_info.cumulative_token_block_count + + work_tile_info.tile_m_idx // + cutlass.Int32(self.cluster_shape_mn[0])) + peek_threshold = work_tile_info.valid_tokens_in_cluster_tile + counter_ptr = token_comm_args.fc1_ready_counter.iterator + counter_slot if not work_tile_info.peek_ready: _iket.range_push("tma_token_fc1_wait") spin_wait( counter_ptr, - lambda v: v >= work_tile_info.valid_tokens_in_tile, - fail_sleep_cycles=20, + lambda v: v >= peek_threshold, + fail_sleep_cycles=1000, ) _iket.range_pop() @@ -311,6 +557,7 @@ def dispatch_prep( warp_idx, lane_idx, *, + local_rank, num_tokens, num_sms, ): @@ -400,8 +647,9 @@ def dispatch_prep( topk_slot) MAX_SLOT_C: cutlass.Constexpr[ int] = num_tokens * self.num_topk - elem_off = ((local_expert * Int32(self.world_size) + Int32( - self.local_rank)) * Int32(MAX_SLOT_C) + slot) * Int32(4) + elem_off = ((local_expert * Int32(self.world_size) + + Int32(local_rank)) * Int32(MAX_SLOT_C) + + slot) * Int32(4) peer_addr = peer_rank_ptr_mapper.map( src_token_topk_idx.iterator.toint(), dst_rank, @@ -424,8 +672,9 @@ def dispatch_barrier( warp_idx, lane_idx, *, + local_rank, num_sms, - nvlink_barrier_counter=None, + nvlink_barrier_counter, ): # software_grid_sync expects a dispatch-group-relative thread id. tid_in_group = warp_idx * Int32(self.warp_threads) + lane_idx @@ -456,9 +705,9 @@ def dispatch_barrier( ) token_count_u32 = Int32(status_u64 & Int64(0xFFFFFFFF)) erc_local_base = expert_recv_count.iterator.toint() - erc_elem_off = (Int32(self.local_rank) * - Int32(self.num_experts_per_rank) + - dst_local_expert) * Int32(8) + erc_elem_off = ( + Int32(local_rank) * Int32(self.num_experts_per_rank) + + dst_local_expert) * Int32(8) erc_peer_addr = peer_rank_ptr_mapper.map( erc_local_base, dst_rank, @@ -471,7 +720,8 @@ def dispatch_barrier( dst_rank, Int64(dst_local_expert * Int32(8)), ) - red_add_release_sys_u64_raw(ercs_peer_addr, status_u64) + red_add_relaxed_sys_u64_raw(ercs_peer_addr, status_u64) + cute.arch.fence_acq_rel_sys() cute.arch.barrier( barrier_id=self.dispatch_intra_cta_bar_id, number_of_threads=self.num_dispatch_threads, @@ -485,7 +735,6 @@ def dispatch_barrier( sm_idx, warp_idx, lane_idx, - slot=0, num_sms=num_sms, prologue_grid_sync=False, epilogue_grid_sync=True, @@ -530,6 +779,19 @@ def dispatch_pull( # SF rows use their own padding; token and SF pool offsets can diverge. expert_sf_pool_block_offset = Int32(0) + # ── Release-flag batching ──────────────────────────────────────── + # Delay fc1-ready counter publication with the same rotating-lane + # tracker used by the epilogue. Each token's TMA store to the FC1 pool + # is drained CTA-locally by ``cp_async_bulk_wait_group(0)`` before its + # release target is accumulated; the eventual red.release.gpu add + # publishes the corresponding pool data to GPU scope. + flag_tracker = GpuReleaseFlagBatchTracker( + flag_addr=Int64(0), + cumulated_flags=Int32(0), + phase=Int32(0), + tid=lane_idx, + ) + stored_rank_count_lane = Int32(0) NUM_EXPERTS_PER_LANE: cutlass.Constexpr[int] = ( @@ -549,8 +811,8 @@ def dispatch_pull( int] = num_sms * self.num_dispatch_warps token_idx = sm_idx * Int32(self.num_dispatch_warps) + warp_idx - _iket_pull_emit = (sm_idx == Int32(0)) and (warp_idx == Int32(0)) and ( - lane_idx == Int32(0)) + _iket_pull_emit = ((sm_idx == Int32(0)) and (warp_idx == Int32(0)) + and (lane_idx == Int32(0))) while current_expert_idx < Int32(self.num_experts_per_rank): if _iket_pull_emit: @@ -562,17 +824,20 @@ def dispatch_pull( prev_block_count = (prev_valid_count + Int32(self.token_padding_block) - Int32(1)) // Int32(self.token_padding_block) - expert_pool_block_offset = expert_pool_block_offset + prev_block_count + expert_pool_block_offset = (expert_pool_block_offset + + prev_block_count) # Mirror cumul for the release-counter granularity (self.cluster_tile_tokens). prev_task_tile_count = ( prev_valid_count + Int32(self.cluster_tile_tokens) - Int32(1)) // Int32(self.cluster_tile_tokens) - expert_task_tile_offset = expert_task_tile_offset + prev_task_tile_count + expert_task_tile_offset = (expert_task_tile_offset + + prev_task_tile_count) # Mirror cumul for the SF axis granularity (self.sf_padding_block). prev_sf_block_count = (prev_valid_count + Int32(self.sf_padding_block) - Int32(1)) // Int32(self.sf_padding_block) - expert_sf_pool_block_offset = expert_sf_pool_block_offset + prev_sf_block_count + expert_sf_pool_block_offset = (expert_sf_pool_block_offset + + prev_sf_block_count) current_expert_idx = current_expert_idx + Int32(1) if current_expert_idx < Int32(self.num_experts_per_rank): expert_start_idx = expert_end_idx @@ -744,13 +1009,13 @@ def dispatch_pull( ) with cute.arch.elect_one(): - _store_token_src_metadata_u32x3( - token_src_metadata, - pool_token_idx, - Uint32(current_rank_in_expert_idx), - Uint32(src_token), - Uint32(src_topk), - ) + TokenSrcMetadata( + src_rank=current_rank_in_expert_idx, + src_token=src_token, + src_topk=src_topk, + ).store(token_src_metadata.iterator + + Int64(pool_token_idx) * + Int64(TokenSrcMetadata.nbytes)) with cute.arch.elect_one(): cute.arch.cp_async_bulk_commit_group() @@ -760,15 +1025,21 @@ def dispatch_pull( _iket.range_pop() # Pull.TMA_Store _iket.range_push("Pull.Arrival_Atomic") - with cute.arch.elect_one(): - task_tile_idx = expert_task_tile_offset + ( - token_idx_in_expert // Int32(self.cluster_tile_tokens)) - cute.arch.atomic_add( - fc1_ready_counter.iterator + task_tile_idx, - Int32(1), - sem="release", - scope="gpu", - ) + # Accumulate this token's release target into the rotating-lane + # batch tracker. task_tile_idx is warp-uniform (token_idx / + # expert offsets are warp-wide), so every lane runs the same + # state-machine transition while only one lane records the + # current address. + task_tile_idx = expert_task_tile_offset + ( + token_idx_in_expert // Int32(self.cluster_tile_tokens)) + + task_tile_addr = (fc1_ready_counter.iterator + + task_tile_idx).toint() + flag_tracker = flag_tracker.accumulate( + Int32(0), + self._flag_batch, + task_tile_addr, + ) cute.arch.sync_warp() if _iket_pull_emit: @@ -778,16 +1049,54 @@ def dispatch_pull( token_idx = token_idx + Int32(num_global_warps) + # Tail flush: publish any leftover (< self._flag_batch) accumulated release. + flag_tracker.fire() + cute.arch.sync_warp() + return phase_bit, stored_num_tokens_per_expert + @cute.jit + def _adaptive_pace( + self, + avg, + current_window, + *, + lo: cutlass.Constexpr[int], + hi: cutlass.Constexpr[int], + ): + # NVLink pacing: EMA the measured round-trip and nanosleep the deviation + # so outstanding NVLink requests stay bounded and don't head-of-line + # block this SM's non-NVLink (local) load/store traffic. + if current_window > avg: + avg = avg + ((current_window - avg + Int32(3)) // Int32(4)) + sleep_cycle = current_window - avg + if sleep_cycle > Int32(hi): + sleep_cycle = Int32(hi) + if sleep_cycle > Int32(50): + _nanosleep(sleep_cycle) + else: + avg = avg - ((avg - current_window + Int32(3)) // Int32(4)) + sleep_cycle = avg - current_window + if sleep_cycle > Int32(50): + _nanosleep(sleep_cycle) + if avg > Int32(hi): + avg = Int32(hi) + if avg < Int32(lo): + avg = Int32(lo) + return avg + @cute.jit def token_back_by_push( self, - token_comm_storage, + pull_buffer_ptr, + pull_mbar_ptr, fc2_output_workspace, fc2_done_counter, token_src_metadata, combine_output, + combine_sf, + fc2_output_sf, + token_back_schedule_counter, peer_rank_ptr_mapper, phase_bit, stored_num_tokens_per_expert, @@ -795,25 +1104,78 @@ def token_back_by_push( warp_idx, lane_idx, *, + local_rank, num_sms, + chunk_bytes: cutlass.Constexpr[int], ): _iket_emit = (sm_idx == Int32(0)) and (warp_idx == Int32(0)) + avg_token_back_window = Int32(2500) - chunk_bytes: cutlass.Constexpr[int] = self.hidden_bytes - num_chunks: cutlass.Constexpr[int] = self.fc2_num_chunks + # Chunk the fc2 token in ``chunk_bytes`` pieces; the last piece carries + # the remainder so any chunk_bytes works for any fc2_token_bytes. fc2_token_bytes: cutlass.Constexpr[int] = self.fc2_token_bytes - - pull_buffer_ptr = token_comm_storage.pull_buffer.data_ptr() - pull_mbar_ptr = token_comm_storage.pull_mbar.data_ptr() + num_chunks: cutlass.Constexpr[int] = (fc2_token_bytes + chunk_bytes - + 1) // chunk_bytes + last_chunk_bytes: cutlass.Constexpr[int] = ( + fc2_token_bytes - (num_chunks - 1) * chunk_bytes) + + if cutlass.const_expr(self.push_sf): + # (token, topk, hidden):(d_topkxhidden, d_hidden, 1) + combine_sf_u8 = cute.recast_tensor(combine_sf, Uint8) + sf_token_bytes: cutlass.Constexpr[int] = cute.size( + combine_sf_u8[0, None, 0].stride) + num_sf_chunks: cutlass.Constexpr[int] = ( + sf_token_bytes + chunk_bytes - 1) // chunk_bytes + last_sf_chunk_bytes: cutlass.Constexpr[int] = ( + sf_token_bytes - (num_sf_chunks - 1) * chunk_bytes) num_experts_per_lane: cutlass.Constexpr[int] = ( self.num_experts_per_rank + 31) // 32 - num_global_warps: cutlass.Constexpr[ - int] = num_sms * self.num_dispatch_warps + num_global_warps: cutlass.Constexpr[int] = (num_sms * + self.num_dispatch_warps) + schedule_mode = self.token_back_schedule_mode + + # static: stride by the global warp count. atomic_counter: claim the + # next token via one grid-scoped atomicAdd(1) so fast warps keep + # stealing work. cuTeDSL forbids closures over enclosing locals -> + # pass all in. + def update_token_idx( + token_idx, + lane_idx, + schedule_counter, + schedule_mode, + num_global_warps, + ): + if cutlass.const_expr(schedule_mode == "atomic_counter"): + base = Int32(0) + if lane_idx == Int32(0): + base = cute.arch.atomic_add( + schedule_counter, + Int32(1), + sem="relaxed", + scope="gpu", + ) + token_idx = cute.arch.shuffle_sync(base, Int32(0)) + else: + token_idx = token_idx + Int32(num_global_warps) + return token_idx - token_idx = sm_idx * Int32(self.num_dispatch_warps) + warp_idx + if cutlass.const_expr(schedule_mode == "atomic_counter"): + # Claim the initial token. + token_idx = Int32(0) + token_idx = update_token_idx( + token_idx, + lane_idx, + token_back_schedule_counter, + schedule_mode, + num_global_warps, + ) + else: + token_idx = sm_idx * Int32(self.num_dispatch_warps) + warp_idx current_expert_idx = Int32(-1) + confirmed_expert_idx = Int32(-1) + cur_expert_expected = Int32(0) expert_start_idx = Int32(0) expert_end_idx = Int32(0) expert_pool_block_offset = Int32(0) @@ -825,7 +1187,8 @@ def token_back_by_push( prev_block_count = (prev_valid_count + Int32(self.token_padding_block) - Int32(1)) // Int32(self.token_padding_block) - expert_pool_block_offset = expert_pool_block_offset + prev_block_count + expert_pool_block_offset = (expert_pool_block_offset + + prev_block_count) current_expert_idx = current_expert_idx + Int32(1) if current_expert_idx < Int32(self.num_experts_per_rank): @@ -845,81 +1208,164 @@ def token_back_by_push( cluster_tile_cnt = ( total_for_expert + Int32(self.cluster_tile_tokens) - Int32(1)) // Int32(self.cluster_tile_tokens) - expected = cluster_tile_cnt * Int32( + # Stash the threshold; the wait is deferred to the expert we + # actually land on, so stepped-over experts are never waited. + cur_expert_expected = cluster_tile_cnt * Int32( self.fc2_publishes_per_token_cluster_tile) + + if current_expert_idx < Int32(self.num_experts_per_rank): + # Wait once per processed expert (both indices monotonic; fc2 + # completes in expert order so confirming k implies all < k). + if current_expert_idx > confirmed_expert_idx: spin_wait( fc2_done_counter.iterator + current_expert_idx, - lambda v: v >= expected, + lambda v: v >= cur_expert_expected, fail_sleep_cycles=500, ) + confirmed_expert_idx = current_expert_idx - if current_expert_idx < Int32(self.num_experts_per_rank): + remain_experts = Int32( + self.num_experts_per_rank) - current_expert_idx token_idx_in_expert = token_idx - expert_start_idx pool_token_idx = ( expert_pool_block_offset * Int32(self.token_padding_block) + token_idx_in_expert) - md_base = token_src_metadata.iterator + (pool_token_idx * - Int32(12)) - src_rank = Int32( - cute.arch.load(md_base + Int32(0), Int32, scope="gpu")) - src_token = Int32( - cute.arch.load(md_base + Int32(4), Int32, scope="gpu")) - src_topk = Int32( - cute.arch.load(md_base + Int32(8), Int32, scope="gpu")) - - local_token_addr = fc2_output_workspace.iterator.toint( - ) + Int64(pool_token_idx) * Int64(fc2_token_bytes) - peer_combine_ptr = peer_rank_ptr_mapper.ptr_map_to_rank( - combine_output.iterator, - src_rank, - ) - peer_token_ptr = peer_combine_ptr + ( - Int64(src_token * Int32(self.num_topk) + src_topk) * - Int64(fc2_token_bytes)) + md = TokenSrcMetadata.load(token_src_metadata.iterator + + Int64(pool_token_idx) * + Int64(TokenSrcMetadata.nbytes)) + src_rank = md.src_rank + src_token = md.src_token + src_topk = md.src_topk + is_remote_token_back = src_rank != Int32(local_rank) smem_ptr_warp = pull_buffer_ptr + warp_idx * Int32(chunk_bytes) mbar_ptr_warp = pull_mbar_ptr + warp_idx if _iket_emit: _iket.range_push("token_back") + cute.arch.sync_warp() - for chunk in cutlass.range_constexpr(0, num_chunks, 1): - chunk_off = Int64(chunk * chunk_bytes) - # chunk_t0 = read_clock64() - - with cute.arch.elect_one(): - tma_load_1d_raw( - smem_ptr_warp, - local_token_addr + chunk_off, - mbar_ptr_warp, - Int32(chunk_bytes), - ) - cute.arch.mbarrier_arrive_and_expect_tx( - mbar_ptr_warp, - Int32(chunk_bytes), - ) - cute.arch.mbarrier_wait(mbar_ptr_warp, phase_bit) - cute.arch.sync_warp() - - with cute.arch.elect_one(): - tma_store_1d( - peer_token_ptr + chunk_off, - smem_ptr_warp, - Int32(chunk_bytes), - ) + # DATA plane: only the dispatch DATA path pushes here; epi_warps + # has the epilogue STG/UBLK the data straight to the peer. + if cutlass.const_expr(self.push_data): + local_token_addr = ( + fc2_output_workspace.iterator.toint() + + Int64(pool_token_idx) * Int64(fc2_token_bytes)) + peer_combine_ptr = peer_rank_ptr_mapper.ptr_map_to_rank( + combine_output.iterator, + src_rank, + ) + if cutlass.const_expr(self.token_back_reduce_topk): + peer_token_offset = Int64(src_token) * Int64( + fc2_token_bytes) + else: + peer_token_offset = ( + Int64(src_token * Int32(self.num_topk) + src_topk) * + Int64(fc2_token_bytes)) + peer_token_ptr = peer_combine_ptr + peer_token_offset + + for chunk in cutlass.range(num_chunks, unroll=1): + t0 = read_clock64() + chunk_off = Int64(chunk * chunk_bytes) + peer_chunk_ptr = peer_token_ptr + chunk_off + + this_bytes = Int32(chunk_bytes) + if cutlass.const_expr(last_chunk_bytes != chunk_bytes): + if chunk == Int32(num_chunks - 1): + this_bytes = Int32(last_chunk_bytes) + + with cute.arch.elect_one(): + tma_load_1d_raw( + smem_ptr_warp, + local_token_addr + chunk_off, + mbar_ptr_warp, + this_bytes, + ) + cute.arch.mbarrier_arrive_and_expect_tx( + mbar_ptr_warp, + this_bytes, + ) + cute.arch.mbarrier_wait(mbar_ptr_warp, phase_bit) + if cutlass.const_expr(self.token_back_reduce_topk): + cp_reduce_async_bulk_add_noftz_bf16_s2g( + peer_chunk_ptr, + smem_ptr_warp, + this_bytes, + ) + else: + tma_store_1d( + peer_chunk_ptr, + smem_ptr_warp, + this_bytes, + ) + phase_bit = phase_bit ^ Int32(1) cute.arch.cp_async_bulk_commit_group() cute.arch.cp_async_bulk_wait_group(0) - cute.arch.sync_warp() - - # if read_clock64() - chunk_t0 < Int64(600): - # _nanosleep(100) + t1 = read_clock64() + current_window = Int32(t1 - t0) + if is_remote_token_back and remain_experts > Int32(4): + avg_token_back_window = self._adaptive_pace( + avg_token_back_window, + current_window, + lo=1000, + hi=5000, + ) + + if cutlass.const_expr(self.push_sf): + # Int64 like the DATA/metadata paths above: the Int32 + # row index would overflow the multi-GiB SF pool offset + # at large max_tokens_per_rank x EP. + sf_local_addr = fc2_output_sf[Int64(pool_token_idx), 0, + None].iterator.toint() + sf_peer_ptr = peer_rank_ptr_mapper.ptr_map_to_rank( + combine_sf_u8[Int64(src_token), src_topk, + None].iterator, + src_rank, + ) + for chunk in cutlass.range(num_sf_chunks, unroll=1): + t0 = read_clock64() + chunk_off = Int64(chunk * chunk_bytes) + this_bytes = Int32(chunk_bytes) + if cutlass.const_expr( + last_sf_chunk_bytes != chunk_bytes): + if chunk == Int32(num_sf_chunks - 1): + this_bytes = Int32(last_sf_chunk_bytes) + with cute.arch.elect_one(): + tma_load_1d_raw( + smem_ptr_warp, + sf_local_addr + chunk_off, + mbar_ptr_warp, + this_bytes, + ) + cute.arch.mbarrier_arrive_and_expect_tx( + mbar_ptr_warp, + this_bytes, + ) + cute.arch.mbarrier_wait(mbar_ptr_warp, phase_bit) + tma_store_1d( + sf_peer_ptr + chunk_off, + smem_ptr_warp, + this_bytes, + ) + phase_bit = phase_bit ^ Int32(1) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0) + if is_remote_token_back: + round_trip = Int32(read_clock64() - t0) + if round_trip <= this_bytes * Int32(5) // Int32(4): + _nanosleep(this_bytes // Int32(4)) - phase_bit = phase_bit ^ Int32(1) if _iket_emit: _iket.range_pop() - token_idx = token_idx + Int32(num_global_warps) + token_idx = update_token_idx( + token_idx, + lane_idx, + token_back_schedule_counter, + schedule_mode, + num_global_warps, + ) cute.arch.fence_acq_rel_sys() @@ -934,7 +1380,6 @@ def nvlink_barrier( warp_idx, lane_idx, *, - slot: cutlass.Constexpr[int], num_sms, prologue_grid_sync: cutlass.Constexpr[bool], epilogue_grid_sync: cutlass.Constexpr[bool], @@ -943,26 +1388,27 @@ def nvlink_barrier( tid_in_group = warp_idx * Int32(self.warp_threads) + lane_idx if prologue_grid_sync: - software_grid_sync( - grid_sync_counter, - sm_idx, - num_sms, - tid_in_group, - num_threads=self.num_dispatch_threads, - ) + software_grid_sync(grid_sync_counter, + sm_idx, + num_sms, + tid_in_group, + num_threads=self.num_dispatch_threads) if sm_idx == 0: if warp_idx == 0: - signal_phase = Int32(slot) + # Sense-reversing ping-pong barrier. The low 2 bits of the counter + # pick the signal slot (phase 0/1) and the direction (+1 up to + # world_size, then -1 back to 0), so the two slots self-cancel over + # a 4-call cycle and never need an explicit reset of the + # symmetric peer-memory signal. + status = nvlink_barrier_counter[0] & Int32(3) + signal_phase = status & Int32(1) + signal_sign = status >> Int32(1) signal_delta = Int32(1) target = Int32(self.world_size) - if cutlass.const_expr(nvlink_barrier_counter is not None): - status = nvlink_barrier_counter[0] & Int32(3) - signal_phase = status & Int32(1) - signal_sign = status >> Int32(1) - if signal_sign != Int32(0): - signal_delta = Int32(-1) - target = Int32(0) + if signal_sign != Int32(0): + signal_delta = Int32(-1) + target = Int32(0) nbs_local_base = nvlink_barrier_signal.iterator.toint() if lane_idx < Int32(self.world_size): @@ -975,35 +1421,25 @@ def nvlink_barrier( cute.arch.sync_warp() if lane_idx == 0: - if cutlass.const_expr(nvlink_barrier_counter is not None): - cute.arch.atomic_add( - nvlink_barrier_counter.iterator, - Int32(1), - sem="relaxed", - scope="gpu", - ) + cute.arch.atomic_add( + nvlink_barrier_counter.iterator, + Int32(1), + sem="relaxed", + scope="gpu", + ) local_signal_ptr = nvlink_barrier_signal.iterator + signal_phase - if cutlass.const_expr(nvlink_barrier_counter is None): - while (cute.arch.load(local_signal_ptr, - Int32, - sem="acquire", - scope="sys") < target): - pass - else: - while (cute.arch.load(local_signal_ptr, - Int32, - sem="acquire", - scope="sys") != target): - pass + while cute.arch.load(local_signal_ptr, + Int32, + sem="acquire", + scope="sys") != target: + pass if epilogue_grid_sync: - software_grid_sync( - grid_sync_counter, - sm_idx, - num_sms, - tid_in_group, - num_threads=self.num_dispatch_threads, - ) + software_grid_sync(grid_sync_counter, + sm_idx, + num_sms, + tid_in_group, + num_threads=self.num_dispatch_threads) @cute.jit def dispatch_warp_body( @@ -1036,6 +1472,7 @@ def dispatch_warp_body( cta_linear_id, local_warp_idx, lane_idx, + local_rank=token_comm_args.local_rank, num_tokens=token_comm_args.input_token_buffer.shape[0], num_sms=token_comm_args.sm_count, ) @@ -1054,6 +1491,7 @@ def dispatch_warp_body( cta_linear_id, local_warp_idx, lane_idx, + local_rank=token_comm_args.local_rank, num_sms=token_comm_args.sm_count, nvlink_barrier_counter=token_comm_args.nvlink_barrier_counter, ) @@ -1091,61 +1529,146 @@ def dispatch_warp_body( if iket_active: _iket.range_pop() - if cutlass.const_expr(self.enable_token_back): + if cutlass.const_expr(self.enable_token_back + and not self.token_back_standalone): if iket_active: _iket.range_push("Token_Back_By_Push") self.token_back_by_push( - token_comm_storage, + token_comm_storage.pull_buffer.data_ptr(), + token_comm_storage.pull_mbar.data_ptr(), token_comm_args.fc2_output_workspace, token_comm_args.fc2_done_counter, token_comm_args.token_src_metadata, token_comm_args.combine_output, + token_comm_args.combine_sf, + token_comm_args.fc2_output_sf, + token_comm_args.token_back_schedule_counter, token_comm_args.peer_rank_ptr_mapper, phase_bit, stored_num_tokens_per_expert, cta_linear_id, local_warp_idx, lane_idx, + local_rank=token_comm_args.local_rank, num_sms=token_comm_args.sm_count, + chunk_bytes=self.hidden_bytes, ) if iket_active: _iket.range_pop() @cute.jit - def tail_reset_shared_counters( + def token_back_warp_body( + self, + token_comm_args, + token_comm_storage, + *, + warp_idx, + lane_idx, + tidx, + ): + bidx, bidy, bidz = cute.arch.block_idx() + cta_linear_id = ( + Int32(bidx) + Int32(self.cluster_shape_mn[1]) * Int32(bidy) + + Int32(self.cluster_shape_mn[1] * self.cluster_shape_mn[0]) * + Int32(bidz)) + local_warp_idx = Int32(warp_idx) - Int32(self.token_back_warp_start) + + # Handshake: dispatch_barrier done => expert_recv_count_sum populated. + nb_dispatch_to_sched = pipeline.NamedBarrier( + barrier_id=self.dispatch_to_sched_named_barrier_id, + num_threads=self.dispatch_to_sched_threads, + ) + nb_dispatch_to_sched.arrive_and_wait() + + tb_pull_mbar_ptr = token_comm_storage.tb_pull_mbar.data_ptr() + tb_pull_buffer_ptr = token_comm_storage.tb_pull_buffer.data_ptr() + if lane_idx == Int32(0): + cute.arch.mbarrier_init(tb_pull_mbar_ptr + local_warp_idx, 1) + cute.arch.sync_warp() + + NUM_EXPERTS_PER_LANE: cutlass.Constexpr[int] = ( + self.num_experts_per_rank + 31) // 32 + stored_num_tokens_per_expert = [] + for _ in cutlass.range_constexpr(0, NUM_EXPERTS_PER_LANE, 1): + stored_num_tokens_per_expert.append(Int32(0)) + for i in cutlass.range_constexpr(0, NUM_EXPERTS_PER_LANE, 1): + e_idx_for_lane = Int32(i * self.warp_threads) + lane_idx + if e_idx_for_lane < Int32(self.num_experts_per_rank): + sum_packed_init = token_comm_args.expert_recv_count_sum[ + e_idx_for_lane] + stored_num_tokens_per_expert[i] = Int32( + Int64(sum_packed_init) & Int64(0xFFFFFFFF)) + cute.arch.sync_warp() + + iket_active = (cta_linear_id == Int32(0)) and (local_warp_idx + == Int32(0)) + if iket_active: + _iket.range_push("Token_Back_By_Push_Standalone") + + self.token_back_by_push( + tb_pull_buffer_ptr, + tb_pull_mbar_ptr, + token_comm_args.fc2_output_workspace, + token_comm_args.fc2_done_counter, + token_comm_args.token_src_metadata, + token_comm_args.combine_output, + token_comm_args.combine_sf, + token_comm_args.fc2_output_sf, + token_comm_args.token_back_schedule_counter, + token_comm_args.peer_rank_ptr_mapper, + Int32(0), + stored_num_tokens_per_expert, + cta_linear_id, + local_warp_idx, + lane_idx, + local_rank=token_comm_args.local_rank, + num_sms=token_comm_args.sm_count, + chunk_bytes=self.tb_chunk_bytes, + ) + + if iket_active: + _iket.range_pop() + + @cute.jit + def tail_reset_counters( self, token_comm_args, + target_zero_tensor, *, cta_linear_id, local_warp_idx, lane_idx, ): - thread_linear = (cta_linear_id * Int32(self.num_dispatch_warps) + - local_warp_idx) * Int32(self.warp_threads) + lane_idx + """Per-lane 4B (Int32) bulk-zero of one accumulating-counter prefix. + + ``target_zero_tensor`` is an Int32 view over a workspace's front counter + region (megamoe_kernel front-places every counter that must restart at 0 + each launch; data buffers and the phase-flip ``nvlink_barrier_signal`` sit + after the prefix and are untouched). The zeroing is spread across all + dispatch threads grid-wide. kernel_tail calls this twice: + * SHARED prefix (expert_recv_count / _sum) BETWEEN the two nvlink + barriers, so the final barrier publishes the zeros cross-rank -- needed + when the next launch is another MegaMoE reusing the shared workspace + with no intervening rank sync; + * LOCAL prefix (l1_arrival / expert_send / grid_sync / nvlink_barrier / + fc1_done [+ fc2_done / token_back_schedule / load_balance]) AFTER the + last barrier -- rank-local (next kernel sees it via stream order) and + grid_sync/nvlink barrier counters stay live until that last barrier. + Only the FIRST launch relies on a caller-zeroed workspace. + """ + thread_linear = ( + (cta_linear_id * Int32(self.num_dispatch_warps) + local_warp_idx) * + Int32(self.warp_threads) + lane_idx) stride = Int32(token_comm_args.sm_count * self.num_dispatch_threads) - recv_total: cutlass.Constexpr[ - int] = self.world_size * self.num_experts_per_rank - i = thread_linear - while i < Int32(recv_total): - rank_idx = i // Int32(self.num_experts_per_rank) - expert_idx = i % Int32(self.num_experts_per_rank) - token_comm_args.expert_recv_count[rank_idx, expert_idx] = Int64(0) - i = i + stride - + count = cute.size(target_zero_tensor) i = thread_linear - while i < Int32(self.num_experts_per_rank): - token_comm_args.expert_recv_count_sum[i] = Int64(0) + while i < Int32(count): + target_zero_tensor[i] = Int32(0) i = i + stride - if cutlass.const_expr(self.enable_token_back): - i = thread_linear - while i < Int32(self.num_experts_per_rank): - token_comm_args.fc2_done_counter[i] = Int32(0) - i = i + stride - @cute.jit def kernel_tail( self, @@ -1161,13 +1684,21 @@ def kernel_tail( ) nb_kernel_tail.arrive_and_wait() - if warp_idx >= self.dispatch_warp_start: + # Only the dispatch warps run NVLink cleanup; standalone token-back + # warps (>= token_back_warp_start) just join the rendezvous above. + if (warp_idx >= self.dispatch_warp_start) and ( + warp_idx < self.dispatch_warp_start + self.num_dispatch_warps): bidx, bidy, bidz = cute.arch.block_idx() cta_linear_id = ( Int32(bidx) + Int32(self.cluster_shape_mn[1]) * Int32(bidy) + Int32(self.cluster_shape_mn[1] * self.cluster_shape_mn[0]) * Int32(bidz)) local_warp_idx = Int32(warp_idx) - Int32(self.dispatch_warp_start) + # Per-launch nvlink barrier count is 3: 1 (dispatch_barrier) + 2 + # below (drain + publish, around the shared reset). The + # sense-reversing signal rides the phase counter across launch + # boundaries, so the count does not need to self-cancel within a + # single launch. self.nvlink_barrier( token_comm_args.nvlink_barrier_signal, token_comm_args.nvlink_barrier_counter, @@ -1176,26 +1707,15 @@ def kernel_tail( cta_linear_id, local_warp_idx, lane_idx, - slot=1, - num_sms=token_comm_args.sm_count, - prologue_grid_sync=True, - epilogue_grid_sync=True, - ) - self.nvlink_barrier( - token_comm_args.nvlink_barrier_signal, - token_comm_args.nvlink_barrier_counter, - token_comm_args.grid_sync_counter, - token_comm_args.peer_rank_ptr_mapper, - cta_linear_id, - local_warp_idx, - lane_idx, - slot=1, num_sms=token_comm_args.sm_count, prologue_grid_sync=True, epilogue_grid_sync=True, ) - self.tail_reset_shared_counters( + # Shared counters between the barriers: the slot=0 barrier below + # publishes these zeros cross-rank for a back-to-back MegaMoE relaunch. + self.tail_reset_counters( token_comm_args, + token_comm_args.shared_zero_prefix, cta_linear_id=cta_linear_id, local_warp_idx=local_warp_idx, lane_idx=lane_idx, @@ -1208,8 +1728,16 @@ def kernel_tail( cta_linear_id, local_warp_idx, lane_idx, - slot=0, num_sms=token_comm_args.sm_count, prologue_grid_sync=True, epilogue_grid_sync=True, ) + # Local counters last: rank-local, and grid_sync/nvlink_barrier + # counters above stay live until this final barrier completes. + self.tail_reset_counters( + token_comm_args, + token_comm_args.local_zero_prefix, + cta_linear_id=cta_linear_id, + local_warp_idx=local_warp_idx, + lane_idx=lane_idx, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.py index b99f61b04de0..9976de3d2dab 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/topk_reduce.py @@ -1,1651 +1,509 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Standalone CuTeDSL topk reduce kernel. - -Form A writes one BF16 fc2 output row per ``(token, topk)`` cell into -``combine_output`` with logical shape ``(T, K, H)``. This module provides the -device-side final reduce used by the default form-A path: - - BF16 (T, K, H) -> FP32 accumulate over K -> FP32/BF16 (T, H) - -It also supports an explicit MXFP8 input mode: - - FP8_E4M3 (T, K, H) + UE8M0 scale -> FP32 dequant/reduce -> BF16 (T, H) - -and an explicit NVFP4 input mode: - - FP4_E2M1 (T, K, H) + per-16 FP8 scale + per-128 FP32 scale - -> FP32 dequant/reduce -> FP32/BF16 (T, H) - -It intentionally does not touch dispatch metadata, peer pointer mapping, or -the fc2 epilogue STG path. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause +"""Device combine reduce: collapse per-(token, topk) fc2 cells into one row. + +The combine step writes one fc2 output per ``(token, topk)`` cell; this reduces +over the topk axis into the token-centric ``(token, hidden)`` output. The wire +format is a :class:`~src.token_comm.CombineFormat`: + + bf16 -- no staging: bf16 terms reduced directly. + 32e4m3xe8m0 -- MXFP8: fp8 e4m3 data + per-32 e8m0 (power-of-2) scale. + 16e2m1xbf16 -- fp4 e2m1 data + per-16 bf16 amax (one level, no global); + dequant per element x = fp4 * (amax * (1 / 6)). + +Task partition: each worker owns one ``(token, hidden_tile)`` and loops topk; the +flat worker index decodes into ``(token_idx, hidden_tile_idx)`` via a constant +divide by ``hidden_tiles``. The per-block scale is broadcast to a logical +per-hidden view (stride 0) so it tiles by the same worker index as the data. The +activation load stays in the topk loop (too large to hoist); the small scale and +score loads are hoisted ahead of the loop when topk is small. """ from __future__ import annotations -import argparse -from typing import Optional, Tuple +from typing import ClassVar, Dict, Optional import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute -import cutlass.torch as cutlass_torch -import torch -from cutlass.cute.typing import AddressSpace -from cutlass.cutlass_dsl import Float32, Int32 - -DEFAULT_THREADS = 256 - -BF16_VECTOR_THREADS = 512 -BF16_HIDDEN_PER_THREAD = 8 -BF16_STORE_ELEMENTS_PER_256B = 16 - -MXFP8_VECTOR_THREADS = 128 -MXFP8_HIDDEN_PER_THREAD = 16 -MXFP8_SCALE_BLOCK_SIZE = 32 - -NVFP4_VECTOR_THREADS = 128 -NVFP4_HIDDEN_PER_THREAD = 32 -NVFP4_SFC_SCALE_BLOCK_SIZE = 16 -NVFP4_SFC_PACKED_BYTES = NVFP4_SFC_SCALE_BLOCK_SIZE // 2 -NVFP4_SFC_INPUT_BITS_PER_COPY = NVFP4_SFC_PACKED_BYTES * 8 -NVFP4_GLOBAL_SCALE_BLOCK_SIZE = 128 - -NVFP4_E2M1_MAX = 6.0 -FP8_E4M3FN_MAX = 448.0 - -_Fp4DecodeTable: torch.Tensor = torch.tensor( - [ - 0.0, - 0.5, - 1.0, - 1.5, - 2.0, - 3.0, - 4.0, - 6.0, - -0.0, - -0.5, - -1.0, - -1.5, - -2.0, - -3.0, - -4.0, - -6.0, - ], - dtype=torch.float32, -) - -_Fp4ValuesEvenFirst: torch.Tensor = torch.tensor( - [ - 0.0, - 1.0, - 2.0, - 4.0, - -0.0, - -1.0, - -2.0, - -4.0, - 0.5, - 1.5, - 3.0, - 6.0, - -0.5, - -1.5, - -3.0, - -6.0, - ], - dtype=torch.float32, -) +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import Float32, Int32, T -_ReorderToNibble: torch.Tensor = torch.tensor( - [ - 0x0, - 0x2, - 0x4, - 0x6, - 0x8, - 0xA, - 0xC, - 0xE, - 0x1, - 0x3, - 0x5, - 0x7, - 0x9, - 0xB, - 0xD, - 0xF, - ], - dtype=torch.uint8, -) +from .megamoe_constants import Nvfp4E2M1RcpLimit +from .token_comm import CombineFormat +# --------------------------------------------------------------------------- +# fp4 (e2m1) -> fp32 register decode. +# +# Blackwell has no e2m1->f32 upconvert: the framework's ``term.load().to(f32)`` +# lowers to an ALU subnormal-normalization path (~60% DRAM SOL). The helper +# below forces a table-driven decode instead; ``e2m1_reg`` (N e2m1 codes, N % 8 +# == 0) is read as packed b32 words and N fp32 values are written into +# ``fp32_reg`` in code order. The 16 e2m1 values are exact in fp32. +# --------------------------------------------------------------------------- -def logical_io_bytes( - combine_output: torch.Tensor, - reduced_output: torch.Tensor, - topk_score: Optional[torch.Tensor] = None, - mxfp8_scale: Optional[torch.Tensor] = None, - nvfp4_sfc_scale: Optional[torch.Tensor] = None, - nvfp4_global_scale: Optional[torch.Tensor] = None, -) -> Tuple[int, int, int]: - """Return logical read, write and total bytes for one topk reduce pass.""" - read_bytes = combine_output.numel() * combine_output.element_size() - if topk_score is not None: - read_bytes += topk_score.numel() * topk_score.element_size() - if mxfp8_scale is not None: - read_bytes += mxfp8_scale.numel() * mxfp8_scale.element_size() - if nvfp4_sfc_scale is not None: - read_bytes += nvfp4_sfc_scale.numel() * nvfp4_sfc_scale.element_size() - if nvfp4_global_scale is not None: - read_bytes += nvfp4_global_scale.numel( - ) * nvfp4_global_scale.element_size() - write_bytes = reduced_output.numel() * reduced_output.element_size() - return int(read_bytes), int(write_bytes), int(read_bytes + write_bytes) +@cute.jit +def cvt_e2m1_to_fp32_cvt_ptx(e2m1_reg: cute.Tensor, + fp32_reg: cute.Tensor) -> None: + """Decode via the e2m1->f16 HW cvt (``cvt.rn.f16x2.e2m1x2``) then widen f16->f32. -def bandwidth_gbps(num_bytes: int, elapsed_ms: float) -> float: - if elapsed_ms <= 0.0: - return float("inf") - return float(num_bytes) / (elapsed_ms * 1.0e6) - - -def make_mxfp8_input( - src: torch.Tensor, - *, - scale_rank: int = 3, -) -> tuple[torch.Tensor, torch.Tensor]: - """Quantize FP32 ``src`` to MXFP8 data plus UE8M0 dequant scale.""" - if src.dim() != 3: - raise ValueError( - f"src must have shape (T, K, H), got {tuple(src.shape)}.") - if src.dtype != torch.float32: - raise TypeError(f"src must be torch.float32, got {src.dtype}.") - if not src.is_cuda: - raise ValueError("src must be a CUDA tensor.") - - T, K, H = src.shape - block = MXFP8_SCALE_BLOCK_SIZE - scale_cols = (H + block - 1) // block - padded_abs = torch.zeros( - (T, K, scale_cols * block), - device=src.device, - dtype=torch.float32, - ) - padded_abs[:, :, :H] = src.abs() - amax = padded_abs.reshape(T, K, scale_cols, block).amax(dim=-1) - if scale_rank == 2: - scale_f32 = amax.amax(dim=1) / 448.0 - scale_for_q = scale_f32[:, None, :] - elif scale_rank == 3: - scale_f32 = amax / 448.0 - scale_for_q = scale_f32 - else: - raise ValueError(f"scale_rank must be 2 or 3, got {scale_rank}.") - - def _round_up_to_power_of_two(scale: torch.Tensor) -> torch.Tensor: - return torch.pow( - torch.full_like(scale, 2.0), - torch.ceil(torch.log2(torch.clamp(scale, min=2.0**-30))), - ) - - scale_f32 = _round_up_to_power_of_two(scale_f32) - scale_for_q = _round_up_to_power_of_two(scale_for_q) - expanded_scale = scale_for_q.repeat_interleave(block, dim=-1)[:, :, :H] - q = (src / expanded_scale).to(torch.float8_e4m3fn) - return q, scale_f32.to(torch.float8_e8m0fnu) - - -def _pack_f32_to_fp4(fp32: torch.Tensor) -> torch.Tensor: - """Round FP32 to FP4 E2M1 and pack pairs along the last dimension.""" - if fp32.dim() == 0 or fp32.shape[-1] % 2 != 0: - raise ValueError( - f"FP4 packing requires an even non-empty last dim, got {tuple(fp32.shape)}." - ) - device = fp32.device - boundaries = torch.tensor( - [ - -5.0, - -3.5, - -2.5, - -1.75, - -1.25, - -0.75, - -0.25, - 0.25, - 0.75, - 1.25, - 1.75, - 2.5, - 3.5, - 5.0, - ], - device=device, - dtype=fp32.dtype, - ) - bucket_to_nibble = torch.tensor( - [ - 0xF, - 0xE, - 0xD, - 0xC, - 0xB, - 0xA, - 0x9, - 0x0, - 0x1, - 0x2, - 0x3, - 0x4, - 0x5, - 0x6, - 0x7, - ], - device=device, - dtype=torch.uint8, - ) - bucket = torch.bucketize(fp32.contiguous(), boundaries) - indices = bucket_to_nibble[bucket] - lo = indices[..., 0::2] - hi = indices[..., 1::2] - return ((hi << 4) | lo).contiguous() - - -def unpack_fp4_to_f32(packed: torch.Tensor) -> torch.Tensor: - """Unpack a last-dim-packed FP4 tensor or uint8 byte tensor to FP32.""" - if packed.dtype == torch.uint8: - raw = packed - elif hasattr(torch, - "float4_e2m1fn_x2") and packed.dtype == torch.float4_e2m1fn_x2: - raw = packed.view(torch.uint8) - else: - raise TypeError( - f"packed must be torch.uint8 or torch.float4_e2m1fn_x2, got {packed.dtype}." - ) - lo = (raw & 0x0F).to(torch.int64) - hi = (raw >> 4).to(torch.int64) - lut = _Fp4DecodeTable.to(raw.device) - unpacked_shape = list(raw.shape) - unpacked_shape[-1] *= 2 - unpacked = torch.empty(unpacked_shape, - dtype=torch.float32, - device=raw.device) - unpacked[..., 0::2] = lut[lo] - unpacked[..., 1::2] = lut[hi] - return unpacked - - -def make_nvfp4_input( - src: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Quantize FP32 ``src`` to NVFP4 plus per-16 FP8 and per-128 FP32 scales. - - The returned scales are dequant scales along hidden: - ``x_hat = fp4 * sfc_fp8 * global_fp32``. + Safe baseline: the per-pair ``cvt`` instruction is itself a HW PRMT+F2FP, so + the e2m1->f16 step already avoids ALU normalization; f16->f32 is one cheap + ``cvt.f32.f16`` per element. """ - if not hasattr(torch, "float8_e4m3fn"): - raise TypeError("NVFP4 mode requires torch float8_e4m3fn.") - if src.dim() != 3: - raise ValueError( - f"src must have shape (T, K, H), got {tuple(src.shape)}.") - if src.dtype != torch.float32: - raise TypeError(f"src must be torch.float32, got {src.dtype}.") - if not src.is_cuda: - raise ValueError("src must be a CUDA tensor.") - if src.shape[-1] % 2 != 0: - raise ValueError( - f"NVFP4 input hidden must be even for fp4x2 packing, got {src.shape[-1]}." + src_words = cute.recast_tensor(e2m1_reg, Int32) # (N,) e2m1 -> (N/8,) b32 + for w in cutlass.range_constexpr(cute.size(src_words)): + res = llvm.inline_asm( + llvm.StructType.get_literal([T.f32()] * 8), + [src_words[w].ir_value()], + "{\n" + " .reg .b8 b0, b1, b2, b3;\n" + " .reg .b32 p0, p1, p2, p3;\n" + " .reg .b16 c0, d0, c1, d1, c2, d2, c3, d3;\n" + " mov.b32 {b0, b1, b2, b3}, $8;\n" + " cvt.rn.f16x2.e2m1x2 p0, b0;\n" + " cvt.rn.f16x2.e2m1x2 p1, b1;\n" + " cvt.rn.f16x2.e2m1x2 p2, b2;\n" + " cvt.rn.f16x2.e2m1x2 p3, b3;\n" + " mov.b32 {c0, d0}, p0;\n" + " mov.b32 {c1, d1}, p1;\n" + " mov.b32 {c2, d2}, p2;\n" + " mov.b32 {c3, d3}, p3;\n" + " cvt.f32.f16 $0, c0;\n" + " cvt.f32.f16 $1, d0;\n" + " cvt.f32.f16 $2, c1;\n" + " cvt.f32.f16 $3, d1;\n" + " cvt.f32.f16 $4, c2;\n" + " cvt.f32.f16 $5, d2;\n" + " cvt.f32.f16 $6, c3;\n" + " cvt.f32.f16 $7, d3;\n" + "}", + "=f,=f,=f,=f,=f,=f,=f,=f,r", + has_side_effects=False, ) + for i in cutlass.range_constexpr(8): + fp32_reg[w * 8 + i] = Float32(llvm.extractvalue(T.f32(), res, [i])) - T, K, H = src.shape - sfc_block = NVFP4_SFC_SCALE_BLOCK_SIZE - global_block = NVFP4_GLOBAL_SCALE_BLOCK_SIZE - sfc_cols = (H + sfc_block - 1) // sfc_block - global_cols = (H + global_block - 1) // global_block - - padded_abs_sfc = torch.zeros( - (T, K, sfc_cols * sfc_block), - device=src.device, - dtype=torch.float32, - ) - padded_abs_sfc[:, :, :H] = src.abs() - amax16 = padded_abs_sfc.reshape(T, K, sfc_cols, sfc_block).amax(dim=-1) - padded_abs_global = torch.zeros( - (T, K, global_cols * global_block), - device=src.device, - dtype=torch.float32, - ) - padded_abs_global[:, :, :H] = src.abs() - amax128 = padded_abs_global.reshape(T, K, global_cols, - global_block).amax(dim=-1) +class TopkReduce: + """Combine reduce for a fixed ``(hidden, num_topk, combine_format)``. - global_scale = torch.clamp( - amax128 / (NVFP4_E2M1_MAX * FP8_E4M3FN_MAX), - min=2.0**-16, - ) - global_for_sfc = global_scale.repeat_interleave( - global_block // sfc_block, - dim=-1, - )[:, :, :sfc_cols] - sfc_fp32 = amax16 / (NVFP4_E2M1_MAX * global_for_sfc) - sfc_fp32 = torch.clamp(sfc_fp32, min=2.0**-16, max=FP8_E4M3FN_MAX) - sfc_fp8 = sfc_fp32.to(torch.float8_e4m3fn) - sfc_rt = sfc_fp8.to(torch.float32) + ``__init__`` pins the static shape and format (and the derived launch + geometry); ``__call__`` (a ``@cute.jit`` launcher) sizes a 1D grid from the + runtime token count and dispatches the format's kernel. The caller owns the + torch->cute conversion and the ``cute.compile``. + """ - expanded_sfc = sfc_rt.repeat_interleave(sfc_block, dim=-1)[:, :, :H] - expanded_global = global_scale.repeat_interleave(global_block, - dim=-1)[:, :, :H] - q = _pack_f32_to_fp4(src / (expanded_sfc * expanded_global)) - return q, sfc_fp8, global_scale + _threads: ClassVar[int] = 128 + # combine_format.name -> hidden elements per worker (one LDG of data: + # bf16 8*2B=16B, e4m3 16*1B=16B, e2m1 16*0.5B=8B). For quantized formats this + # stays <= the scale block, so each worker reads exactly one scale entry. + _hidden_per_thread: ClassVar[Dict[str, int]] = { + "bf16": 8, + "32e4m3xe8m0": 16, + "16e2m1xbf16": 16, + } + # topk count at/below which the scale + score loads are hoisted ahead of the + # topk loop (small enough to not bloat registers; a CTA-broadcast read). + _prefetch_limit: ClassVar[int] = 16 + + def __init__(self, hidden: int, num_topk: int, + combine_format: CombineFormat) -> None: + self.hidden = int(hidden) + self.num_topk = int(num_topk) + self.combine_format = combine_format + self.hidden_per_thread = self._hidden_per_thread[combine_format.name] + # hidden must tile cleanly both into worker slices and into scale blocks. + align = max(combine_format.scale_block or self.hidden_per_thread, + self.hidden_per_thread) + if self.hidden % align != 0: + raise ValueError( + f"hidden ({self.hidden}) must be divisible by max(scale_block, " + f"hidden_per_thread) = {align} for combine_format {combine_format}." + ) + self.hidden_tiles = self.hidden // self.hidden_per_thread + # tail guard only needed when the worker count per token is not a whole + # number of CTAs; prefetch only when topk is small enough to hoist. + self.require_predicate = self.hidden_tiles % self._threads != 0 + self.prefetch = self.num_topk <= self._prefetch_limit + # -- launcher ------------------------------------------------------------- -def mxfp8_reference_sum( - q: torch.Tensor, - scale: torch.Tensor, - topk_score: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """Return K-ordered FP32 reduce of MXFP8 input after dequantization.""" - T, K, H = q.shape - block = MXFP8_SCALE_BLOCK_SIZE - if scale.dim() == 2: - scale_for_q = scale.to(torch.float32)[:, None, :] - else: - scale_for_q = scale.to(torch.float32) - expanded_scale = scale_for_q.repeat_interleave(block, dim=-1)[:, :, :H] - dequant = q.to(torch.float32) * expanded_scale - acc = torch.zeros((T, H), device=q.device, dtype=torch.float32) - for k in range(K): - contrib = dequant[:, k, :] - if topk_score is not None: - acc = torch.addcmul(acc, contrib, topk_score[:, k, None]) + @cute.jit + def __call__( + self, + combine_quant: cute.Tensor, # (token, topk, hidden) + combine_sf: Optional[cute.Tensor], # (token, topk, hidden) + reduced_output: cute.Tensor, # (token, hidden) + topk_score: Optional[cute.Tensor], # (token, topk) + stream: cuda.CUstream, + ): + threads = self._threads + total_workers = reduced_output.shape[0] * self.hidden_tiles + grid = [(total_workers + threads - 1) // threads, 1, 1] + block = [threads, 1, 1] + + combine_quant = cute.make_tensor( + combine_quant.iterator, + cute.make_layout( + (combine_quant.shape[0], self.num_topk, self.hidden), + stride=combine_quant.stride)) + reduced_output = cute.make_tensor( + reduced_output.iterator, + cute.make_layout((reduced_output.shape[0], self.hidden), + stride=reduced_output.stride)) + if cutlass.const_expr(topk_score is not None): + topk_score = cute.make_tensor( + topk_score.iterator, + cute.make_layout((topk_score.shape[0], self.num_topk), + stride=topk_score.stride)) + + if cutlass.const_expr(not self.combine_format.is_quantized): + self._reduce_bf16(combine_quant, topk_score, reduced_output).launch( + grid=grid, + block=block, + stream=stream, + ) + return + + # The mega kernel hands sf in already as the depth-2 broadcast layout; a + # plain (torch) sf is depth-1 and gets its hidden mode split into + # (sf_vec, hidden/sf_vec):(0, s_h) so logical hidden h reads block h//sf_vec. + sf_vec = self.combine_format.scale_block + if cutlass.const_expr(cute.depth(combine_sf.layout) >= 2): + sf = cute.make_tensor( + combine_sf.iterator, + cute.make_layout((combine_sf.shape[0], self.num_topk, + (sf_vec, self.hidden // sf_vec)), + stride=combine_sf.stride)) else: - acc = acc + contrib - return acc - - -def nvfp4_reference_sum( - q: torch.Tensor, - sfc_scale: torch.Tensor, - global_scale: torch.Tensor, - topk_score: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """Return K-ordered FP32 reduce of hierarchical NVFP4 input.""" - unpacked = unpack_fp4_to_f32(q) - T, K, H = unpacked.shape - expanded_sfc = sfc_scale.to(torch.float32).repeat_interleave( - NVFP4_SFC_SCALE_BLOCK_SIZE, - dim=-1, - )[:, :, :H] - expanded_global = global_scale.to(torch.float32).repeat_interleave( - NVFP4_GLOBAL_SCALE_BLOCK_SIZE, - dim=-1, - )[:, :, :H] - acc = torch.zeros((T, H), device=q.device, dtype=torch.float32) - for k in range(K): - contrib = unpacked[:, k, :] * expanded_sfc[:, k, :] - if topk_score is not None: - contrib = contrib * expanded_global[:, k, :] - acc = torch.addcmul(acc, contrib, topk_score[:, k, None]) + sf = cute.make_tensor( + combine_sf.iterator, + cute.make_layout( + (combine_sf.shape[0], self.num_topk, + (sf_vec, self.hidden // sf_vec)), + stride=(combine_sf.stride[0], combine_sf.stride[1], + (0, combine_sf.stride[2])), + ), + ) + + if cutlass.const_expr( + self.combine_format.act_dtype is cutlass.Float8E4M3FN): + self._reduce_mxfp8(combine_quant, sf, topk_score, + reduced_output).launch( + grid=grid, + block=block, + stream=stream, + ) else: - acc = torch.addcmul(acc, contrib, expanded_global[:, k, :]) - return acc - - -def weighted_reference_sum( - src: torch.Tensor, - topk_score: torch.Tensor, -) -> torch.Tensor: - """Return K-ordered FP32 weighted reduce using FMA/addcmul semantics.""" - src_f32 = src.to(torch.float32) - acc = torch.zeros( - (src.shape[0], src.shape[2]), - device=src.device, - dtype=torch.float32, - ) - for k in range(src.shape[1]): - acc = torch.addcmul(acc, src_f32[:, k, :], topk_score[:, k, None]) - return acc - - -def ordered_reference_sum(src: torch.Tensor) -> torch.Tensor: - """Return K-ordered FP32 reduce of BF16 input.""" - src_f32 = src.to(torch.float32) - acc = torch.zeros( - (src.shape[0], src.shape[2]), - device=src.device, - dtype=torch.float32, - ) - for k in range(src.shape[1]): - acc = acc + src_f32[:, k, :] - return acc + self._reduce_fp4(combine_quant, sf, topk_score, + reduced_output).launch( + grid=grid, + block=block, + stream=stream, + ) + @cute.jit + def _mark_alignment(self, tensor: cute.Tensor, + align_bytes: int) -> cute.Tensor: + p = tensor.iterator + return cute.make_tensor( + cute.make_ptr(p.dtype, + p.toint(), + p.memspace, + assumed_align=align_bytes), + tensor.layout, + ) -@cute.jit -def _fp4_e2m1_nibble_to_f32(nibble: Int32) -> Float32: - value = Float32(0.0) - if nibble == Int32(1): - value = Float32(0.5) - elif nibble == Int32(2): - value = Float32(1.0) - elif nibble == Int32(3): - value = Float32(1.5) - elif nibble == Int32(4): - value = Float32(2.0) - elif nibble == Int32(5): - value = Float32(3.0) - elif nibble == Int32(6): - value = Float32(4.0) - elif nibble == Int32(7): - value = Float32(6.0) - elif nibble == Int32(9): - value = Float32(-0.5) - elif nibble == Int32(10): - value = Float32(-1.0) - elif nibble == Int32(11): - value = Float32(-1.5) - elif nibble == Int32(12): - value = Float32(-2.0) - elif nibble == Int32(13): - value = Float32(-3.0) - elif nibble == Int32(14): - value = Float32(-4.0) - elif nibble == Int32(15): - value = Float32(-6.0) - return value - - -@cute.kernel -def topk_reduce_bf16_vec_kernel( - combine_output: cute.Tensor, - topk_score: Optional[cute.Tensor], - reduced_output: cute.Tensor, - num_topk: cutlass.Constexpr[int], - hidden: cutlass.Constexpr[int], - store_dtype: cutlass.Constexpr[str], -): - """BF16 reduce with one thread handling one 8-hidden vector.""" - - hidden_vec_block_idx, token_idx, _ = cute.arch.block_idx() - tid = cute.arch.thread_idx()[0] - block_dim = cute.arch.block_dim()[0] - vec_idx = hidden_vec_block_idx * block_dim + tid - base_h = vec_idx * Int32(BF16_HIDDEN_PER_THREAD) - - if base_h < Int32(hidden): - acc = cute.make_rmem_tensor((BF16_HIDDEN_PER_THREAD, ), cutlass.Float32) - for i in cutlass.range_constexpr(0, BF16_HIDDEN_PER_THREAD, 1): - acc[i] = Float32(0.0) + # -- kernels -------------------------------------------------------------- - copy_atom_bf16_vec = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - cutlass.BFloat16, - num_bits_per_copy=128, - ) + @cute.kernel + def _reduce_bf16( + self, + combine_output: cute.Tensor, + topk_score: Optional[cute.Tensor], + reduced_output: cute.Tensor, + ): + threads = self._threads + hidden_per_thread = self.hidden_per_thread + hidden_tiles = self.hidden_tiles + num_topk: cutlass.Constexpr[int] = self.num_topk + needs_guard = self.require_predicate + prefetch = self.prefetch + out_dtype = reduced_output.element_type + + worker_idx = cute.arch.block_idx()[0] * Int32( + threads) + cute.arch.thread_idx()[0] + token_idx = worker_idx // hidden_tiles + hidden_tile_idx = worker_idx % hidden_tiles + + score_dtype = topk_score.dtype if cutlass.const_expr( + topk_score is not None) else cutlass.Float32 + score_reg = cute.make_rmem_tensor((num_topk, ), score_dtype) + + if (not needs_guard) or token_idx < reduced_output.shape[0]: + # (token, topk, hidden) -> (topk, hidden_per_thread) + terms = cute.zipped_divide( + combine_output[token_idx, None, None], + (num_topk, hidden_per_thread), + )[(None, None), (0, hidden_tile_idx)] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide( + reduced_output[token_idx, None], + (hidden_per_thread, ), + )[(None, ), (hidden_tile_idx, )] - for k in cutlass.range_constexpr(0, num_topk, 1): - score_value = Float32(1.0) if cutlass.const_expr(topk_score is not None): - score_value = Float32(topk_score[token_idx, Int32(k)]) - score_pair = (score_value, score_value) + if cutlass.const_expr(prefetch): + cute.autovec_copy(topk_score[token_idx, None], score_reg) + else: + for k in cutlass.range_constexpr(num_topk): + score_reg[k] = score_dtype(1) - in_regs = cute.make_rmem_tensor( - (BF16_HIDDEN_PER_THREAD, ), + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, + num_bits_per_copy=128, ) - in_row = combine_output[token_idx, Int32(k), None] - in_tile = cute.local_tile( - in_row, - (BF16_HIDDEN_PER_THREAD, ), - (base_h // Int32(BF16_HIDDEN_PER_THREAD), ), - ) - in_aligned_iter = cute.make_ptr( - in_tile.element_type, - in_tile.iterator.toint(), - AddressSpace.gmem, - assumed_align=16, - ) - in_tile = cute.make_tensor(in_aligned_iter, in_tile.layout) - cute.copy( - copy_atom_bf16_vec, - cute.coalesce(in_tile), - cute.coalesce(in_regs), - ) - - for pair_i in cutlass.range_constexpr( - 0, - BF16_HIDDEN_PER_THREAD // 2, - 1, - ): - val_pair = ( - Float32(in_regs[2 * pair_i]), - Float32(in_regs[2 * pair_i + 1]), - ) - old_acc_pair = (acc[2 * pair_i], acc[2 * pair_i + 1]) - if cutlass.const_expr(topk_score is not None): - acc_pair = cute.arch.fma_packed_f32x2( - val_pair, - score_pair, - old_acc_pair, - ) - else: - acc_pair = cute.arch.add_packed_f32x2( - old_acc_pair, - val_pair, - ) - acc[2 * pair_i] = acc_pair[0] - acc[2 * pair_i + 1] = acc_pair[1] + acc = cute.make_rmem_tensor((hidden_per_thread, ), cutlass.Float32) - out_row = reduced_output[token_idx, None] - out_tile = cute.local_tile( - out_row, - (BF16_HIDDEN_PER_THREAD, ), - (base_h // Int32(BF16_HIDDEN_PER_THREAD), ), - ) - if cutlass.const_expr(store_dtype == "bf16"): - out_regs = cute.make_rmem_tensor( - (BF16_HIDDEN_PER_THREAD, ), - cutlass.BFloat16, - ) - out_regs.store(acc.load().to(cutlass.BFloat16)) - out_aligned_iter = cute.make_ptr( - out_tile.element_type, - out_tile.iterator.toint(), - AddressSpace.gmem, - assumed_align=16, - ) - out_tile = cute.make_tensor(out_aligned_iter, out_tile.layout) + for k in cutlass.range_constexpr(0, num_topk, 1): + term = cute.make_rmem_tensor((hidden_per_thread, ), + cutlass.BFloat16) + cute.copy(load_atom, terms[k, None], term) + if cutlass.const_expr(topk_score is not None and not prefetch): + score_reg[k] = topk_score[token_idx, Int32(k)] + score_pair = (Float32(score_reg[k]), Float32(score_reg[k])) + + for i in cutlass.range_constexpr(0, hidden_per_thread, 2): + value_pair = (Float32(term[i]), Float32(term[i + 1])) + if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2( + value_pair, score_pair, (acc[i], acc[i + 1])) + else: + if cutlass.const_expr(topk_score is not None): + acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2( + value_pair, score_pair) + else: + acc[i] = value_pair[0] + acc[i + 1] = value_pair[1] + + out = cute.make_rmem_tensor((hidden_per_thread, ), out_dtype) + out.store(acc.load().to(out_dtype)) cute.copy( - copy_atom_bf16_vec, - cute.coalesce(out_regs), - cute.coalesce(out_tile), - ) - else: - for i in cutlass.range_constexpr(0, BF16_HIDDEN_PER_THREAD, 1): - out_tile[i] = acc[i] - - -@cute.kernel -def topk_reduce_mxfp8_vec_kernel( - combine_output: cute.Tensor, - topk_score: Optional[cute.Tensor], - mxfp8_scale: cute.Tensor, - reduced_output: cute.Tensor, - num_topk: cutlass.Constexpr[int], - hidden: cutlass.Constexpr[int], - mxfp8_scale_rank: cutlass.Constexpr[int], -): - """MXFP8 reduce with one thread handling one 16-hidden vector.""" - - hidden_vec_block_idx, token_idx, _ = cute.arch.block_idx() - tid = cute.arch.thread_idx()[0] - block_dim = cute.arch.block_dim()[0] - vec_idx = hidden_vec_block_idx * block_dim + tid - base_h = vec_idx * Int32(MXFP8_HIDDEN_PER_THREAD) - - if base_h < Int32(hidden): - acc = cute.make_rmem_tensor((MXFP8_HIDDEN_PER_THREAD, ), - cutlass.Float32) - for i in cutlass.range_constexpr(0, MXFP8_HIDDEN_PER_THREAD, 1): - acc[i] = Float32(0.0) - - copy_atom_ldg_128b = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - cutlass.Float8E4M3FN, - num_bits_per_copy=128, - ) - copy_atom_stg_256b = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - cutlass.BFloat16, - num_bits_per_copy=256, - ) - scale_col = base_h // Int32(MXFP8_SCALE_BLOCK_SIZE) + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), + out_dtype, + num_bits_per_copy=128), + out, + self._mark_alignment(dst, + hidden_per_thread * out_dtype.width // 8), + ) + + @cute.kernel + def _reduce_mxfp8( + self, + combine_quant: cute.Tensor, + combine_sf: cute. + Tensor, # depth-2 broadcast view: logical (token, topk, hidden) e8m0 + topk_score: Optional[cute.Tensor], + reduced_output: cute.Tensor, + ): + threads = self._threads + hidden_per_thread = self.hidden_per_thread + hidden_tiles = self.hidden_tiles + num_topk: cutlass.Constexpr[int] = self.num_topk + needs_guard = self.require_predicate + prefetch = self.prefetch + out_dtype = reduced_output.element_type + + worker_idx = cute.arch.block_idx()[0] * Int32( + threads) + cute.arch.thread_idx()[0] + token_idx = worker_idx // hidden_tiles + hidden_tile_idx = worker_idx % hidden_tiles + + score_dtype = topk_score.dtype if cutlass.const_expr( + topk_score is not None) else cutlass.Float32 + score_reg = cute.make_rmem_tensor((num_topk, ), score_dtype) + scale_reg = cute.make_rmem_tensor((num_topk, ), cutlass.Float8E8M0FNU) + + if (not needs_guard) or token_idx < reduced_output.shape[0]: + # (token, topk, hidden) -> (topk, hidden_per_thread) + codes = cute.zipped_divide( + combine_quant[token_idx, None, None], + (num_topk, hidden_per_thread), + )[(None, None), (0, hidden_tile_idx)] + # (token, topk, hidden) -> (topk, hidden_per_thread) + sf = cute.zipped_divide( + combine_sf[token_idx, None, None], + (num_topk, hidden_per_thread), + )[(None, None), (0, hidden_tile_idx)] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide( + reduced_output[token_idx, None], + (hidden_per_thread, ), + )[(None, ), (hidden_tile_idx, )] - for k in cutlass.range_constexpr(0, num_topk, 1): - if cutlass.const_expr(mxfp8_scale_rank == 3): - scale = Float32(mxfp8_scale[token_idx, Int32(k), scale_col]) - else: - scale = Float32(mxfp8_scale[token_idx, scale_col]) - scale_pair = (scale, scale) - score_value = Float32(1.0) if cutlass.const_expr(topk_score is not None): - score_value = Float32(topk_score[token_idx, Int32(k)]) - score_pair = (score_value, score_value) - - in_regs = cute.make_rmem_tensor( - (MXFP8_HIDDEN_PER_THREAD, ), + if cutlass.const_expr(prefetch): + cute.autovec_copy(topk_score[token_idx, None], score_reg) + else: + for k in cutlass.range_constexpr(num_topk): + score_reg[k] = score_dtype(1) + if cutlass.const_expr(prefetch): + cute.autovec_copy( + sf[None, 0], + scale_reg) # one scale per topk slot (stride-0 broadcast) + + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float8E4M3FN, + num_bits_per_copy=128, ) - in_row = combine_output[token_idx, Int32(k), None] - in_tile = cute.local_tile( - in_row, - (MXFP8_HIDDEN_PER_THREAD, ), - (base_h // Int32(MXFP8_HIDDEN_PER_THREAD), ), - ) - in_aligned_iter = cute.make_ptr( - in_tile.element_type, - in_tile.iterator.toint(), - AddressSpace.gmem, - assumed_align=16, - ) - in_tile = cute.make_tensor(in_aligned_iter, in_tile.layout) - cute.copy( - copy_atom_ldg_128b, - cute.coalesce(in_tile), - cute.coalesce(in_regs), - ) - in_vals = cute.make_rmem_tensor( - (MXFP8_HIDDEN_PER_THREAD, ), - cutlass.Float32, - ) - in_vals.store(in_regs.load().to(cutlass.Float32)) - - for pair_i in cutlass.range_constexpr( - 0, - MXFP8_HIDDEN_PER_THREAD // 2, - 1, - ): - val_pair = ( - in_vals[2 * pair_i], - in_vals[2 * pair_i + 1], - ) - old_acc_pair = (acc[2 * pair_i], acc[2 * pair_i + 1]) - if cutlass.const_expr(topk_score is not None): - contrib_pair = cute.arch.mul_packed_f32x2( - val_pair, - scale_pair, - ) - acc_pair = cute.arch.fma_packed_f32x2( - contrib_pair, - score_pair, - old_acc_pair, - ) - else: - acc_pair = cute.arch.fma_packed_f32x2( - val_pair, - scale_pair, - old_acc_pair, - ) - acc[2 * pair_i] = acc_pair[0] - acc[2 * pair_i + 1] = acc_pair[1] + acc = cute.make_rmem_tensor((hidden_per_thread, ), cutlass.Float32) - out_row = reduced_output[token_idx, None] - for chunk in cutlass.range_constexpr( - 0, - MXFP8_HIDDEN_PER_THREAD // BF16_STORE_ELEMENTS_PER_256B, - 1, - ): - out_regs = cute.make_rmem_tensor( - (BF16_STORE_ELEMENTS_PER_256B, ), - cutlass.BFloat16, - ) - for i in cutlass.range_constexpr(0, BF16_STORE_ELEMENTS_PER_256B, - 1): - out_regs[i] = acc[chunk * BF16_STORE_ELEMENTS_PER_256B + i].to( - cutlass.BFloat16) - out_h = base_h + Int32(chunk * BF16_STORE_ELEMENTS_PER_256B) - out_tile = cute.local_tile( - out_row, - (BF16_STORE_ELEMENTS_PER_256B, ), - (out_h // Int32(BF16_STORE_ELEMENTS_PER_256B), ), - ) - out_aligned_iter = cute.make_ptr( - out_tile.element_type, - out_tile.iterator.toint(), - AddressSpace.gmem, - assumed_align=32, - ) - out_tile = cute.make_tensor(out_aligned_iter, out_tile.layout) + for k in cutlass.range_constexpr(0, num_topk, 1): + term = cute.make_rmem_tensor((hidden_per_thread, ), + cutlass.Float8E4M3FN) + cute.copy(load_atom, codes[k, None], term) + value = cute.make_rmem_tensor((hidden_per_thread, ), + cutlass.Float32) + value.store(term.load().to(cutlass.Float32)) + + if cutlass.const_expr(not prefetch): + scale_reg[k] = sf[k, 0] + if cutlass.const_expr(topk_score is not None): + score_reg[k] = topk_score[token_idx, Int32(k)] + + scale = Float32(scale_reg[k]) # e8m0 -> f32 + scale_pair = (scale, scale) + score_pair = (Float32(score_reg[k]), Float32(score_reg[k])) + + for i in cutlass.range_constexpr(0, hidden_per_thread, 2): + dequant_pair = cute.arch.mul_packed_f32x2( + (value[i], value[i + 1]), scale_pair) + if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2( + dequant_pair, score_pair, (acc[i], acc[i + 1])) + else: + if cutlass.const_expr(topk_score is not None): + acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2( + dequant_pair, score_pair) + else: + acc[i] = dequant_pair[0] + acc[i + 1] = dequant_pair[1] + + out = cute.make_rmem_tensor((hidden_per_thread, ), out_dtype) + out.store(acc.load().to(out_dtype)) cute.copy( - copy_atom_stg_256b, - cute.coalesce(out_regs), - cute.coalesce(out_tile), - ) - - -@cute.kernel -def topk_reduce_kernel( - combine_output: cute.Tensor, - topk_score: Optional[cute.Tensor], - mxfp8_scale: Optional[cute.Tensor], - nvfp4_sfc_scale: Optional[cute.Tensor], - nvfp4_global_scale: Optional[cute.Tensor], - reduced_output: cute.Tensor, - num_topk: cutlass.Constexpr[int], - hidden: cutlass.Constexpr[int], - store_dtype: cutlass.Constexpr[str], - mxfp8_scale_rank: cutlass.Constexpr[int], -): - """Reduce ``combine_output[t, :, h]`` into ``reduced_output[t, h]``. - - In the default path, ``combine_output`` is BF16. In MXFP8 mode, - ``combine_output`` is FP8 E4M3 and ``mxfp8_scale`` is UE8M0 with either - shape ``(T, ceil_div(H, 32))`` or ``(T, K, ceil_div(H, 32))``. Optional - ``topk_score`` is FP32 with shape ``(T, K)`` and scales each K - contribution before accumulation. Shapes and store dtype are supplied as - constexprs by the launcher so the K loop is fully unrolled and matches the - host reference order exactly. - """ - - hidden_block_idx, token_idx, _ = cute.arch.block_idx() - - h = hidden_block_idx * cute.arch.block_dim()[0] + cute.arch.thread_idx()[0] + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), + out_dtype, + num_bits_per_copy=256), + out, + self._mark_alignment(dst, + hidden_per_thread * out_dtype.width // 8), + ) + + @cute.kernel + def _reduce_fp4( + self, + combine_quant: cute.Tensor, # (token, topk, hidden) e2m1 (logical) + combine_sf: cute. + Tensor, # depth-2 broadcast view: logical (token, topk, hidden) bf16 amax + topk_score: Optional[cute.Tensor], + reduced_output: cute.Tensor, + ): + threads = self._threads + hidden_per_thread = self.hidden_per_thread + hidden_tiles = self.hidden_tiles + num_topk: cutlass.Constexpr[int] = self.num_topk + needs_guard = self.require_predicate + prefetch = self.prefetch + out_dtype = reduced_output.element_type + + worker_idx = cute.arch.block_idx()[0] * Int32( + threads) + cute.arch.thread_idx()[0] + token_idx = worker_idx // hidden_tiles + hidden_tile_idx = worker_idx % hidden_tiles + + score_dtype = topk_score.dtype if cutlass.const_expr( + topk_score is not None) else cutlass.Float32 + score_reg = cute.make_rmem_tensor((num_topk, ), score_dtype) + scale_reg = cute.make_rmem_tensor((num_topk, ), cutlass.BFloat16) + + if (not needs_guard) or token_idx < reduced_output.shape[0]: + # (token, topk, hidden) -> (topk, hidden_per_thread) + codes = cute.zipped_divide( + combine_quant[token_idx, None, None], + (num_topk, hidden_per_thread), + )[(None, None), (0, hidden_tile_idx)] + # (token, topk, hidden) -> (topk, hidden_per_thread) + sf = cute.zipped_divide( + combine_sf[token_idx, None, None], + (num_topk, hidden_per_thread), + )[(None, None), (0, hidden_tile_idx)] + # (token, hidden) -> (hidden_per_thread) + dst = cute.zipped_divide( + reduced_output[token_idx, None], + (hidden_per_thread, ), + )[(None, ), (hidden_tile_idx, )] - if h < Int32(hidden): - acc = Float32(0.0) - for k in cutlass.range_constexpr(0, num_topk, 1): - if cutlass.const_expr(nvfp4_sfc_scale is not None): - byte_col = h // Int32(2) - shift = (h - byte_col * Int32(2)) * Int32(4) - packed = Int32(combine_output[token_idx, Int32(k), byte_col]) - nibble = (packed >> shift) & Int32(0x0F) - contrib = _fp4_e2m1_nibble_to_f32(nibble) - sfc_col = h // Int32(NVFP4_SFC_SCALE_BLOCK_SIZE) - global_col = h // Int32(NVFP4_GLOBAL_SCALE_BLOCK_SIZE) - sfc = Float32(nvfp4_sfc_scale[token_idx, Int32(k), sfc_col]) - global_sf = Float32(nvfp4_global_scale[token_idx, - Int32(k), global_col]) - contrib = contrib * sfc * global_sf - else: - contrib = Float32(combine_output[token_idx, Int32(k), h]) - if cutlass.const_expr(mxfp8_scale is not None): - scale_col = h // Int32(MXFP8_SCALE_BLOCK_SIZE) - if cutlass.const_expr(mxfp8_scale_rank == 3): - scale = Float32(mxfp8_scale[token_idx, - Int32(k), scale_col]) - else: - scale = Float32(mxfp8_scale[token_idx, scale_col]) - contrib = contrib * scale if cutlass.const_expr(topk_score is not None): - contrib = contrib * Float32(topk_score[token_idx, Int32(k)]) - acc = acc + contrib + if cutlass.const_expr(prefetch): + cute.autovec_copy(topk_score[token_idx, None], score_reg) else: - acc = acc + contrib - if cutlass.const_expr(store_dtype == "bf16"): - reduced_output[token_idx, h] = acc.to(cutlass.BFloat16) - else: - reduced_output[token_idx, h] = acc - + for k in cutlass.range_constexpr(num_topk): + score_reg[k] = score_dtype(1) + if cutlass.const_expr(prefetch): + cute.autovec_copy(sf[None, 0], scale_reg) -@cute.kernel -def topk_reduce_nvfp4_vec_kernel( - combine_output: cute.Tensor, - topk_score: Optional[cute.Tensor], - nvfp4_sfc_scale: cute.Tensor, - nvfp4_global_scale: cute.Tensor, - reduced_output: cute.Tensor, - num_topk: cutlass.Constexpr[int], - hidden: cutlass.Constexpr[int], - store_dtype: cutlass.Constexpr[str], -): - """NVFP4 reduce with one thread handling two per-16 hidden blocks.""" - - hidden_vec_block_idx, token_idx, _ = cute.arch.block_idx() - tid = cute.arch.thread_idx()[0] - block_dim = cute.arch.block_dim()[0] - vec_idx = hidden_vec_block_idx * block_dim + tid - base_h = vec_idx * Int32(NVFP4_HIDDEN_PER_THREAD) - - if base_h < Int32(hidden): - sfc_col_base = base_h // Int32(NVFP4_SFC_SCALE_BLOCK_SIZE) - global_col = base_h // Int32(NVFP4_GLOBAL_SCALE_BLOCK_SIZE) - - copy_atom_ldg_sfc = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - cutlass.Uint8, - num_bits_per_copy=NVFP4_SFC_INPUT_BITS_PER_COPY, - ) - copy_atom_stg_256b = cute.make_copy_atom( - cute.nvgpu.CopyUniversalOp(), - cutlass.BFloat16, - num_bits_per_copy=256, - ) - - global_regs = cute.make_rmem_tensor((num_topk, ), cutlass.Float32) - for k in cutlass.range_constexpr(0, num_topk, 1): - global_regs[k] = Float32(nvfp4_global_scale[token_idx, - Int32(k), global_col]) - if cutlass.const_expr(topk_score is not None): - score_regs = cute.make_rmem_tensor((num_topk, ), cutlass.Float32) - for k in cutlass.range_constexpr(0, num_topk, 1): - score_regs[k] = Float32(topk_score[token_idx, Int32(k)]) - - out_row = reduced_output[token_idx, None] - for sfc_block_i in cutlass.range_constexpr( - 0, - NVFP4_HIDDEN_PER_THREAD // NVFP4_SFC_SCALE_BLOCK_SIZE, - 1, - ): - acc = cute.make_rmem_tensor( - (NVFP4_SFC_SCALE_BLOCK_SIZE, ), - cutlass.Float32, + load_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Float4E2M1FN, + num_bits_per_copy=64, ) - for i in cutlass.range_constexpr(0, NVFP4_SFC_SCALE_BLOCK_SIZE, 1): - acc[i] = Float32(0.0) + acc = cute.make_rmem_tensor((hidden_per_thread, ), cutlass.Float32) - sfc_base_h = base_h + Int32( - sfc_block_i * NVFP4_SFC_SCALE_BLOCK_SIZE) for k in cutlass.range_constexpr(0, num_topk, 1): - global_sf = global_regs[k] - global_pair = (global_sf, global_sf) - score_value = Float32(1.0) - if cutlass.const_expr(topk_score is not None): - score_value = score_regs[k] - score_pair = (score_value, score_value) - - q_bytes = cute.make_rmem_tensor( - (NVFP4_SFC_PACKED_BYTES, ), - cutlass.Uint8, - ) - q_row = combine_output[token_idx, Int32(k), None] - q_tile = cute.local_tile( - q_row, - (NVFP4_SFC_PACKED_BYTES, ), - (sfc_base_h // Int32(NVFP4_SFC_SCALE_BLOCK_SIZE), ), - ) - q_aligned_iter = cute.make_ptr( - q_tile.element_type, - q_tile.iterator.toint(), - AddressSpace.gmem, - assumed_align=NVFP4_SFC_PACKED_BYTES, - ) - q_tile = cute.make_tensor(q_aligned_iter, q_tile.layout) - cute.copy( - copy_atom_ldg_sfc, - cute.coalesce(q_tile), - cute.coalesce(q_bytes), - ) - q_fp4 = cute.recast_tensor(q_bytes, cutlass.Float4E2M1FN) - q_vals = q_fp4.load().to(cutlass.Float32) - sfc = Float32(nvfp4_sfc_scale[ - token_idx, - Int32(k), - sfc_col_base + Int32(sfc_block_i), - ]) - sfc_pair = (sfc, sfc) - for byte_offset in cutlass.range_constexpr( - 0, - NVFP4_SFC_SCALE_BLOCK_SIZE // 2, - 1, - ): - val_pair = ( - q_vals[2 * byte_offset], - q_vals[2 * byte_offset + 1], - ) - contrib_pair = cute.arch.mul_packed_f32x2( - val_pair, sfc_pair) - old_acc_pair = ( - acc[2 * byte_offset], - acc[2 * byte_offset + 1], - ) + term = cute.make_rmem_tensor((hidden_per_thread, ), + cutlass.Float4E2M1FN) + cute.copy(load_atom, codes[k, None], term) + value = cute.make_rmem_tensor((hidden_per_thread, ), + cutlass.Float32) + cvt_e2m1_to_fp32_cvt_ptx(term, value) + + if cutlass.const_expr(not prefetch): + scale_reg[k] = sf[k, 0] if cutlass.const_expr(topk_score is not None): - contrib_pair = cute.arch.mul_packed_f32x2( - contrib_pair, - global_pair, - ) - acc_pair = cute.arch.fma_packed_f32x2( - contrib_pair, - score_pair, - old_acc_pair, - ) + score_reg[k] = topk_score[token_idx, Int32(k)] + + # amax (bf16) -> per-element scale; (1/6) folds the fp4 grid max. + scale = Float32(scale_reg[k]) * Float32(Nvfp4E2M1RcpLimit) + scale_pair = (scale, scale) + score_pair = (Float32(score_reg[k]), Float32(score_reg[k])) + + for i in cutlass.range_constexpr(0, hidden_per_thread, 2): + dequant_pair = cute.arch.mul_packed_f32x2( + (value[i], value[i + 1]), scale_pair) + if cutlass.const_expr(k != 0): + acc[i], acc[i + 1] = cute.arch.fma_packed_f32x2( + dequant_pair, score_pair, (acc[i], acc[i + 1])) + elif cutlass.const_expr(topk_score is not None): + acc[i], acc[i + 1] = cute.arch.mul_packed_f32x2( + dequant_pair, score_pair) else: - acc_pair = cute.arch.fma_packed_f32x2( - contrib_pair, - global_pair, - old_acc_pair, - ) - acc[2 * byte_offset] = acc_pair[0] - acc[2 * byte_offset + 1] = acc_pair[1] - - if cutlass.const_expr(store_dtype == "bf16"): - out_regs = cute.make_rmem_tensor( - (BF16_STORE_ELEMENTS_PER_256B, ), - cutlass.BFloat16, - ) - for i in cutlass.range_constexpr(0, - BF16_STORE_ELEMENTS_PER_256B, - 1): - out_regs[i] = acc[i].to(cutlass.BFloat16) - out_tile = cute.local_tile( - out_row, - (BF16_STORE_ELEMENTS_PER_256B, ), - (sfc_base_h // Int32(BF16_STORE_ELEMENTS_PER_256B), ), - ) - out_aligned_iter = cute.make_ptr( - out_tile.element_type, - out_tile.iterator.toint(), - AddressSpace.gmem, - assumed_align=32, - ) - out_tile = cute.make_tensor(out_aligned_iter, out_tile.layout) - cute.copy( - copy_atom_stg_256b, - cute.coalesce(out_regs), - cute.coalesce(out_tile), - ) - else: - out_tile = cute.local_tile( - out_row, - (NVFP4_SFC_SCALE_BLOCK_SIZE, ), - (sfc_base_h // Int32(NVFP4_SFC_SCALE_BLOCK_SIZE), ), - ) - for i in cutlass.range_constexpr(0, NVFP4_SFC_SCALE_BLOCK_SIZE, - 1): - out_tile[i] = acc[i] - - -def _validate_tensors( - combine_output: torch.Tensor, - reduced_output: torch.Tensor, - topk_score: Optional[torch.Tensor] = None, - mxfp8_scale: Optional[torch.Tensor] = None, - nvfp4_sfc_scale: Optional[torch.Tensor] = None, - nvfp4_global_scale: Optional[torch.Tensor] = None, -) -> Tuple[int, int, int, int]: - if combine_output.dim() != 3: - raise ValueError( - f"combine_output must have shape (T, K, H), got {tuple(combine_output.shape)}." - ) - if reduced_output.dim() != 2: - raise ValueError( - f"reduced_output must have shape (T, H), got {tuple(reduced_output.shape)}." - ) - if reduced_output.dtype not in (torch.float32, torch.bfloat16): - raise TypeError( - f"reduced_output must be torch.float32 or torch.bfloat16, got {reduced_output.dtype}." - ) - if not combine_output.is_cuda or not reduced_output.is_cuda: - raise ValueError( - "combine_output and reduced_output must both be CUDA tensors.") - if combine_output.device != reduced_output.device: - raise ValueError( - f"combine_output and reduced_output must be on the same device, got " - f"{combine_output.device} and {reduced_output.device}.") - - if mxfp8_scale is not None and (nvfp4_sfc_scale is not None - or nvfp4_global_scale is not None): - raise ValueError("MXFP8 and NVFP4 modes are mutually exclusive.") - if (nvfp4_sfc_scale is None) != (nvfp4_global_scale is None): - raise ValueError( - "nvfp4_sfc_scale and nvfp4_global_scale must be provided together.") - - T, K, H_storage = combine_output.shape - if T <= 0 or K <= 0 or H_storage <= 0: - raise ValueError( - f"combine_output shape must have positive dimensions, got " - f"{tuple(combine_output.shape)}.") - - nvfp4_mode = nvfp4_sfc_scale is not None - H = int(H_storage) * 2 if nvfp4_mode else int(H_storage) - if reduced_output.shape != (T, H): - raise ValueError( - f"reduced_output shape must be {(T, H)}, got {tuple(reduced_output.shape)}." - ) - - mxfp8_scale_rank = 0 - if mxfp8_scale is None and not nvfp4_mode: - if combine_output.dtype != torch.bfloat16: - raise TypeError( - f"combine_output must be torch.bfloat16 unless mxfp8_scale is " - f"or NVFP4 scales are provided, got {combine_output.dtype}.") - elif mxfp8_scale is not None: - if not hasattr(torch, "float8_e4m3fn") or not hasattr( - torch, "float8_e8m0fnu"): - raise TypeError( - "MXFP8 mode requires torch float8_e4m3fn and float8_e8m0fnu.") - if combine_output.dtype != torch.float8_e4m3fn: - raise TypeError( - f"MXFP8 combine_output must be torch.float8_e4m3fn, got {combine_output.dtype}." - ) - if mxfp8_scale.dtype != torch.float8_e8m0fnu: - raise TypeError( - f"mxfp8_scale must be torch.float8_e8m0fnu, got {mxfp8_scale.dtype}." - ) - if reduced_output.dtype != torch.bfloat16: - raise TypeError( - f"MXFP8 reduced_output must be torch.bfloat16, got {reduced_output.dtype}." - ) - if not mxfp8_scale.is_cuda: - raise ValueError("mxfp8_scale must be a CUDA tensor.") - if mxfp8_scale.device != combine_output.device: - raise ValueError( - f"mxfp8_scale must be on {combine_output.device}, got {mxfp8_scale.device}." - ) - scale_cols = (H + MXFP8_SCALE_BLOCK_SIZE - 1) // MXFP8_SCALE_BLOCK_SIZE - if mxfp8_scale.dim() == 2: - expected_scale_shape = (T, scale_cols) - elif mxfp8_scale.dim() == 3: - expected_scale_shape = (T, K, scale_cols) - else: - raise ValueError( - "mxfp8_scale must have shape (T, ceil_div(H, 32)) or " - f"(T, K, ceil_div(H, 32)), got {tuple(mxfp8_scale.shape)}.") - if mxfp8_scale.shape != expected_scale_shape: - raise ValueError( - f"mxfp8_scale shape must be {expected_scale_shape}, got {tuple(mxfp8_scale.shape)}." - ) - mxfp8_scale_rank = mxfp8_scale.dim() - else: - if not hasattr(torch, "float8_e4m3fn"): - raise TypeError("NVFP4 mode requires torch float8_e4m3fn.") - if combine_output.dtype != torch.uint8: - raise TypeError( - f"NVFP4 combine_output must be packed torch.uint8, got {combine_output.dtype}." - ) - if nvfp4_sfc_scale.dtype != torch.float8_e4m3fn: - raise TypeError( - f"nvfp4_sfc_scale must be torch.float8_e4m3fn, got {nvfp4_sfc_scale.dtype}." - ) - if nvfp4_global_scale.dtype != torch.float32: - raise TypeError( - f"nvfp4_global_scale must be torch.float32, got {nvfp4_global_scale.dtype}." - ) - if not nvfp4_sfc_scale.is_cuda or not nvfp4_global_scale.is_cuda: - raise ValueError("NVFP4 scales must be CUDA tensors.") - if nvfp4_sfc_scale.device != combine_output.device: - raise ValueError( - f"nvfp4_sfc_scale must be on {combine_output.device}, got {nvfp4_sfc_scale.device}." - ) - if nvfp4_global_scale.device != combine_output.device: - raise ValueError( - f"nvfp4_global_scale must be on {combine_output.device}, got " - f"{nvfp4_global_scale.device}.") - sfc_cols = (H + NVFP4_SFC_SCALE_BLOCK_SIZE - - 1) // NVFP4_SFC_SCALE_BLOCK_SIZE - global_cols = (H + NVFP4_GLOBAL_SCALE_BLOCK_SIZE - - 1) // NVFP4_GLOBAL_SCALE_BLOCK_SIZE - expected_sfc_shape = (T, K, sfc_cols) - expected_global_shape = (T, K, global_cols) - if nvfp4_sfc_scale.dim( - ) != 3 or nvfp4_sfc_scale.shape != expected_sfc_shape: - raise ValueError( - f"nvfp4_sfc_scale shape must be {expected_sfc_shape}, got " - f"{tuple(nvfp4_sfc_scale.shape)}.") - if nvfp4_global_scale.dim( - ) != 3 or nvfp4_global_scale.shape != expected_global_shape: - raise ValueError( - f"nvfp4_global_scale shape must be {expected_global_shape}, got " - f"{tuple(nvfp4_global_scale.shape)}.") - if topk_score is not None: - if topk_score.dim() != 2: - raise ValueError( - f"topk_score must have shape (T, K), got {tuple(topk_score.shape)}." - ) - if topk_score.dtype != torch.float32: - raise TypeError( - f"topk_score must be torch.float32, got {topk_score.dtype}.") - if not topk_score.is_cuda: - raise ValueError("topk_score must be a CUDA tensor.") - if topk_score.device != combine_output.device: - raise ValueError( - f"topk_score must be on {combine_output.device}, got {topk_score.device}." - ) - if topk_score.shape != (T, K): - raise ValueError( - f"topk_score shape must be {(T, K)}, got {tuple(topk_score.shape)}." - ) - return int(T), int(K), int(H), int(mxfp8_scale_rank) - - -def _infer_assumed_align(tensor: torch.Tensor, max_align: int = 16) -> int: - ptr = int(tensor.data_ptr()) - for align in (16, 8, 4, 2, 1): - if align <= max_align and ptr % align == 0: - return align - return 1 - - -def _to_cute_tensor(tensor: torch.Tensor) -> cute.Tensor: - assumed_align = _infer_assumed_align(tensor) - cute_tensor = cutlass_torch.from_dlpack(tensor, assumed_align=assumed_align) - leading_dim = cutlass_torch.get_leading_dim(tensor) - return cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) - - -def compile_topk_reduce( - combine_output: torch.Tensor, - reduced_output: torch.Tensor, - topk_score: Optional[torch.Tensor] = None, - *, - mxfp8_scale: Optional[torch.Tensor] = None, - nvfp4_sfc_scale: Optional[torch.Tensor] = None, - nvfp4_global_scale: Optional[torch.Tensor] = None, - threads: Optional[int] = None, - stream: Optional[cuda.CUstream] = None, -): - """Compile a shape-specialized topk reduce launcher. - - The returned tuple is always ``(compiled, combine_cute, reduced_cute, - topk_score_cute, mxfp8_scale_cute, nvfp4_sfc_scale_cute, - nvfp4_global_scale_cute, stream)``. Missing optional inputs are represented - by ``None``. Callers that only need a one-shot reduce should use - :func:`run_topk_reduce`. - """ - T, K, H, mxfp8_scale_rank = _validate_tensors( - combine_output, - reduced_output, - topk_score, - mxfp8_scale, - nvfp4_sfc_scale, - nvfp4_global_scale, - ) - if threads is not None and threads <= 0: - raise ValueError(f"threads must be positive, got {threads}.") - if stream is None: - stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) - store_dtype = "bf16" if reduced_output.dtype == torch.bfloat16 else "fp32" + acc[i] = dequant_pair[0] + acc[i + 1] = dequant_pair[1] - combine_cute = _to_cute_tensor(combine_output) - reduced_cute = _to_cute_tensor(reduced_output) - topk_score_cute = _to_cute_tensor( - topk_score) if topk_score is not None else None - mxfp8_scale_cute = _to_cute_tensor( - mxfp8_scale) if mxfp8_scale is not None else None - nvfp4_sfc_scale_cute = _to_cute_tensor( - nvfp4_sfc_scale) if nvfp4_sfc_scale is not None else None - nvfp4_global_scale_cute = (_to_cute_tensor(nvfp4_global_scale) - if nvfp4_global_scale is not None else None) - nvfp4_mode = nvfp4_sfc_scale is not None - bf16_vectorized = ( - not nvfp4_mode and mxfp8_scale is None - and combine_output.dtype == torch.bfloat16 - and H % BF16_HIDDEN_PER_THREAD == 0 and combine_output.stride(-1) == 1 - and combine_output.stride(-2) % BF16_HIDDEN_PER_THREAD == 0 - and reduced_output.stride(-1) == 1 - and (reduced_output.dtype != torch.bfloat16 - or reduced_output.stride(0) % BF16_HIDDEN_PER_THREAD == 0)) - mxfp8_vectorized = ( - mxfp8_scale is not None and not nvfp4_mode - and H % MXFP8_HIDDEN_PER_THREAD == 0 and combine_output.stride(-1) == 1 - and combine_output.stride(-2) % MXFP8_HIDDEN_PER_THREAD == 0 - and reduced_output.dtype == torch.bfloat16 - and reduced_output.stride(-1) == 1 - and reduced_output.stride(0) % MXFP8_HIDDEN_PER_THREAD == 0) - nvfp4_vectorized = ( - nvfp4_mode and H % NVFP4_HIDDEN_PER_THREAD == 0 - and combine_output.stride(-1) == 1 - and combine_output.stride(-2) % (NVFP4_HIDDEN_PER_THREAD // 2) == 0 - and reduced_output.stride(-1) == 1 - and (reduced_output.dtype != torch.bfloat16 - or reduced_output.stride(0) % NVFP4_HIDDEN_PER_THREAD == 0)) - if bf16_vectorized: - hidden_per_thread = BF16_HIDDEN_PER_THREAD - elif mxfp8_vectorized: - hidden_per_thread = MXFP8_HIDDEN_PER_THREAD - elif nvfp4_vectorized: - hidden_per_thread = NVFP4_HIDDEN_PER_THREAD - else: - hidden_per_thread = 1 - if threads is None: - if bf16_vectorized: - launch_threads = BF16_VECTOR_THREADS - elif mxfp8_vectorized: - launch_threads = MXFP8_VECTOR_THREADS - elif nvfp4_vectorized: - launch_threads = NVFP4_VECTOR_THREADS - else: - launch_threads = DEFAULT_THREADS - else: - launch_threads = threads - hidden_blocks = (H + launch_threads * hidden_per_thread - - 1) // (launch_threads * hidden_per_thread) - launch_grid = [hidden_blocks, T, 1] - - @cute.jit - def _launcher( - combine_cute: cute.Tensor, - reduced_cute: cute.Tensor, - topk_score_cute: Optional[cute.Tensor], - mxfp8_scale_cute: Optional[cute.Tensor], - nvfp4_sfc_scale_cute: Optional[cute.Tensor], - nvfp4_global_scale_cute: Optional[cute.Tensor], - stream: cuda.CUstream, - ): - if cutlass.const_expr(bf16_vectorized): - topk_reduce_bf16_vec_kernel( - combine_cute, - topk_score_cute, - reduced_cute, - num_topk=K, - hidden=H, - store_dtype=store_dtype, - ).launch( - grid=launch_grid, - block=[launch_threads, 1, 1], - stream=stream, - ) - elif cutlass.const_expr(mxfp8_vectorized): - topk_reduce_mxfp8_vec_kernel( - combine_cute, - topk_score_cute, - mxfp8_scale_cute, - reduced_cute, - num_topk=K, - hidden=H, - mxfp8_scale_rank=mxfp8_scale_rank, - ).launch( - grid=launch_grid, - block=[launch_threads, 1, 1], - stream=stream, - ) - elif cutlass.const_expr(nvfp4_vectorized): - topk_reduce_nvfp4_vec_kernel( - combine_cute, - topk_score_cute, - nvfp4_sfc_scale_cute, - nvfp4_global_scale_cute, - reduced_cute, - num_topk=K, - hidden=H, - store_dtype=store_dtype, - ).launch( - grid=launch_grid, - block=[launch_threads, 1, 1], - stream=stream, - ) - else: - topk_reduce_kernel( - combine_cute, - topk_score_cute, - mxfp8_scale_cute, - nvfp4_sfc_scale_cute, - nvfp4_global_scale_cute, - reduced_cute, - num_topk=K, - hidden=H, - store_dtype=store_dtype, - mxfp8_scale_rank=mxfp8_scale_rank, - ).launch( - grid=launch_grid, - block=[launch_threads, 1, 1], - stream=stream, + out = cute.make_rmem_tensor((hidden_per_thread, ), out_dtype) + out.store(acc.load().to(out_dtype)) + cute.copy( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), + out_dtype, + num_bits_per_copy=256), + out, + self._mark_alignment(dst, + hidden_per_thread * out_dtype.width // 8), ) - - compiled = cute.compile( - _launcher, - combine_cute, - reduced_cute, - topk_score_cute, - mxfp8_scale_cute, - nvfp4_sfc_scale_cute, - nvfp4_global_scale_cute, - stream, - ) - return ( - compiled, - combine_cute, - reduced_cute, - topk_score_cute, - mxfp8_scale_cute, - nvfp4_sfc_scale_cute, - nvfp4_global_scale_cute, - stream, - ) - - -def launch_compiled_topk_reduce( - compiled, - combine_cute: cute.Tensor, - reduced_cute: cute.Tensor, - topk_score_cute: Optional[cute.Tensor], - mxfp8_scale_cute: Optional[cute.Tensor], - nvfp4_sfc_scale_cute: Optional[cute.Tensor], - nvfp4_global_scale_cute: Optional[cute.Tensor], - stream: cuda.CUstream, - *, - synchronize: bool = False, - return_elapsed_ms: bool = False, -) -> Optional[float]: - """Launch a topk reduce plan returned by :func:`compile_topk_reduce`.""" - - def _launch() -> None: - compiled( - combine_cute, - reduced_cute, - topk_score_cute, - mxfp8_scale_cute, - nvfp4_sfc_scale_cute, - nvfp4_global_scale_cute, - stream, - ) - - if return_elapsed_ms: - start = torch.cuda.Event(enable_timing=True) - stop = torch.cuda.Event(enable_timing=True) - start.record() - _launch() - stop.record() - stop.synchronize() - elapsed_ms = float(start.elapsed_time(stop)) - else: - _launch() - elapsed_ms = None - - if synchronize: - torch.cuda.synchronize() - return elapsed_ms - - -def run_topk_reduce( - combine_output: torch.Tensor, - reduced_output: torch.Tensor, - topk_score: Optional[torch.Tensor] = None, - *, - mxfp8_scale: Optional[torch.Tensor] = None, - nvfp4_sfc_scale: Optional[torch.Tensor] = None, - nvfp4_global_scale: Optional[torch.Tensor] = None, - threads: Optional[int] = None, - stream: Optional[cuda.CUstream] = None, - synchronize: bool = False, - return_elapsed_ms: bool = False, -) -> Optional[float]: - """Compile and launch the topk reduce kernel. - - Returns the measured kernel elapsed time in milliseconds when - ``return_elapsed_ms`` is True, otherwise returns ``None``. - """ - plan = compile_topk_reduce( - combine_output, - reduced_output, - topk_score, - mxfp8_scale=mxfp8_scale, - nvfp4_sfc_scale=nvfp4_sfc_scale, - nvfp4_global_scale=nvfp4_global_scale, - threads=threads, - stream=stream, - ) - ( - compiled, - combine_cute, - reduced_cute, - topk_score_cute, - mxfp8_scale_cute, - nvfp4_sfc_scale_cute, - nvfp4_global_scale_cute, - stream, - ) = plan - return launch_compiled_topk_reduce( - compiled, - combine_cute, - reduced_cute, - topk_score_cute, - mxfp8_scale_cute, - nvfp4_sfc_scale_cute, - nvfp4_global_scale_cute, - stream, - synchronize=synchronize, - return_elapsed_ms=return_elapsed_ms, - ) - - -def benchmark_topk_reduce_vs_torch_sum( - *, - tokens: int, - topk: int, - hidden: int, - warmup: int = 5, - iters: int = 50, - output_dtype: torch.dtype = torch.float32, - seed: int = 20260531, - use_topk_score: bool = False, - use_mxfp8: bool = False, - use_nvfp4: bool = False, - mxfp8_scale_rank: int = 3, - threads: Optional[int] = None, - print_result: bool = True, -) -> dict[str, float]: - """Compare CuTeDSL topk_reduce against torch K-axis sum. - - The torch baseline intentionally uses the runner/reference expression: - ``combine_output_ref.to(torch.float32).sum(dim=1)`` when ``topk_score`` - is absent, or the weighted equivalent when present. MXFP8 and NVFP4 - inputs are converted from FP32 into their quantized data plus scale tensors - before both benchmark paths. CuTeDSL compile time is excluded from the - measured kernel time. - """ - if not torch.cuda.is_available(): - raise RuntimeError("CUDA GPU is required for topk_reduce benchmark.") - if output_dtype not in (torch.float32, torch.bfloat16): - raise ValueError( - f"output_dtype must be FP32 or BF16, got {output_dtype}.") - if use_mxfp8 and use_nvfp4: - raise ValueError( - "MXFP8 and NVFP4 benchmark modes are mutually exclusive.") - if use_mxfp8 and output_dtype != torch.bfloat16: - raise ValueError("MXFP8 benchmark requires BF16 output.") - if threads is None: - if use_nvfp4: - threads = NVFP4_VECTOR_THREADS - elif use_mxfp8: - threads = MXFP8_VECTOR_THREADS - else: - threads = BF16_VECTOR_THREADS - - torch.manual_seed(seed) - combine_output_fp32 = torch.randn( - (tokens, topk, hidden), - device="cuda", - dtype=torch.float32, - ) - if use_mxfp8: - combine_output_ref, mxfp8_scale = make_mxfp8_input( - combine_output_fp32, - scale_rank=mxfp8_scale_rank, - ) - nvfp4_sfc_scale = None - nvfp4_global_scale = None - input_dtype_name = "mxfp8" - elif use_nvfp4: - combine_output_ref, nvfp4_sfc_scale, nvfp4_global_scale = make_nvfp4_input( - combine_output_fp32, ) - mxfp8_scale = None - input_dtype_name = "nvfp4" - else: - combine_output_ref = combine_output_fp32.to(torch.bfloat16) - mxfp8_scale = None - nvfp4_sfc_scale = None - nvfp4_global_scale = None - input_dtype_name = "bf16" - topk_output = torch.empty( - (tokens, hidden), - device="cuda", - dtype=output_dtype, - ) - topk_score = None - if use_topk_score: - topk_score = torch.rand((tokens, topk), - device="cuda", - dtype=torch.float32) - - if print_result: - print( - "compiling topk_reduce " - f"shape={(tokens, topk, hidden)} output_dtype={output_dtype} " - f"input_dtype={input_dtype_name} " - f"mxfp8_scale_rank={mxfp8_scale_rank if use_mxfp8 else 'none'} " - f"topk_score={'on' if topk_score is not None else 'off'} " - f"threads={threads}", - flush=True, - ) - ( - compiled, - combine_cute, - reduced_cute, - topk_score_cute, - mxfp8_scale_cute, - nvfp4_sfc_scale_cute, - nvfp4_global_scale_cute, - stream, - ) = compile_topk_reduce( - combine_output_ref, - topk_output, - topk_score, - mxfp8_scale=mxfp8_scale, - nvfp4_sfc_scale=nvfp4_sfc_scale, - nvfp4_global_scale=nvfp4_global_scale, - threads=threads, - ) - - compiled( - combine_cute, - reduced_cute, - topk_score_cute, - mxfp8_scale_cute, - nvfp4_sfc_scale_cute, - nvfp4_global_scale_cute, - stream, - ) - torch.cuda.synchronize() - - def reference_result(*, timed_baseline: bool) -> torch.Tensor: - if use_mxfp8: - return mxfp8_reference_sum( - combine_output_ref, - mxfp8_scale, - topk_score, - ).to(output_dtype) - if use_nvfp4: - return nvfp4_reference_sum( - combine_output_ref, - nvfp4_sfc_scale, - nvfp4_global_scale, - topk_score, - ).to(output_dtype) - if topk_score is None: - if output_dtype == torch.bfloat16 and not timed_baseline: - return ordered_reference_sum(combine_output_ref).to( - output_dtype) - return combine_output_ref.to( - torch.float32).sum(dim=1).to(output_dtype) - return weighted_reference_sum(combine_output_ref, - topk_score).to(output_dtype) - - expected_result = reference_result(timed_baseline=False) - torch.testing.assert_close(topk_output, - expected_result, - atol=1e-5, - rtol=1e-5) - - def measure_cuda_ms(fn) -> float: - for _ in range(warmup): - fn() - torch.cuda.synchronize() - - start = torch.cuda.Event(enable_timing=True) - stop = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(iters): - fn() - stop.record() - stop.synchronize() - return float(start.elapsed_time(stop)) / float(iters) - - def run_compiled_topk_reduce() -> None: - compiled( - combine_cute, - reduced_cute, - topk_score_cute, - mxfp8_scale_cute, - nvfp4_sfc_scale_cute, - nvfp4_global_scale_cute, - stream, - ) - - torch_result = None - - def run_torch_sum() -> None: - nonlocal torch_result - torch_result = reference_result(timed_baseline=True) - - topk_ms = measure_cuda_ms(run_compiled_topk_reduce) - torch_ms = measure_cuda_ms(run_torch_sum) - assert torch_result is not None - speedup = torch_ms / topk_ms - read_bytes, write_bytes, total_bytes = logical_io_bytes( - combine_output_ref, - topk_output, - topk_score, - mxfp8_scale, - nvfp4_sfc_scale, - nvfp4_global_scale, - ) - topk_bw = bandwidth_gbps(total_bytes, topk_ms) - torch_bw = bandwidth_gbps(total_bytes, torch_ms) - - if print_result: - print("topk_reduce_vs_torch_sum " - f"shape={(tokens, topk, hidden)} output_dtype={output_dtype} " - f"input_dtype={input_dtype_name} " - f"mxfp8_scale_rank={mxfp8_scale_rank if use_mxfp8 else 'none'} " - f"topk_score={'on' if topk_score is not None else 'off'} " - f"threads={threads} " - f"warmup={warmup} iters={iters} " - f"topk_reduce_ms={topk_ms:.6f} " - f"torch_sum_ms={torch_ms:.6f} " - f"speedup_vs_torch={speedup:.3f}x " - f"read_gb={read_bytes / 1.0e9:.6f} " - f"write_gb={write_bytes / 1.0e9:.6f} " - f"topk_reduce_bw_gbps={topk_bw:.3f} " - f"torch_sum_bw_gbps={torch_bw:.3f}") - - return { - "topk_reduce_ms": topk_ms, - "torch_sum_ms": torch_ms, - "speedup_vs_torch": speedup, - "read_bytes": float(read_bytes), - "write_bytes": float(write_bytes), - "total_bytes": float(total_bytes), - "topk_reduce_bw_gbps": topk_bw, - "torch_sum_bw_gbps": torch_bw, - "use_mxfp8": float(use_mxfp8), - "use_nvfp4": float(use_nvfp4), - "threads": float(threads), - } - - -def _parse_bench_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=( - "Benchmark CuTeDSL topk_reduce against combine_output_ref.to(torch.float32).sum(dim=1)." - )) - parser.add_argument("--tokens", type=int, default=192) - parser.add_argument("--topk", type=int, default=8) - parser.add_argument("--hidden", type=int, default=7168) - parser.add_argument("--warmup", type=int, default=5) - parser.add_argument("--iters", type=int, default=50) - parser.add_argument("--output_dtype", - choices=["fp32", "bf16"], - default="bf16") - parser.add_argument("--use_topk_score", action="store_true") - parser.add_argument("--use_mxfp8", action="store_true") - parser.add_argument("--use_nvfp4", action="store_true") - parser.add_argument("--mxfp8_scale_rank", - type=int, - choices=[2, 3], - default=3) - parser.add_argument("--threads", type=int, default=None) - parser.add_argument("--seed", type=int, default=20260531) - return parser.parse_args() - - -def main() -> int: - args = _parse_bench_args() - output_dtype = torch.bfloat16 if args.output_dtype == "bf16" else torch.float32 - benchmark_topk_reduce_vs_torch_sum( - tokens=args.tokens, - topk=args.topk, - hidden=args.hidden, - warmup=args.warmup, - iters=args.iters, - output_dtype=output_dtype, - use_topk_score=args.use_topk_score, - use_mxfp8=args.use_mxfp8, - use_nvfp4=args.use_nvfp4, - mxfp8_scale_rank=args.mxfp8_scale_rank, - threads=args.threads, - seed=args.seed, - ) - print("DONE") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 3d65b1ee7595..5724c1052d8c 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + # -------------------------------------------------- # Portions of this code were derived from DeepSeek‑V3: # https://github.com/deepseek-ai/DeepSeek-V3 @@ -296,6 +299,43 @@ def _resolve_enable_fused_hc(config: PretrainedConfig) -> bool: return bool(getattr(config, "enable_fused_hc", True)) +def _normalize_deepseek_v4_nvfp4_mixed_precision_config( + model_config: ModelConfig[PretrainedConfig], +) -> ModelConfig[PretrainedConfig]: + """Resolve FP8 base layers in DeepSeek-V4 NVFP4 checkpoints.""" + quant_config = model_config.quant_config + hf_quant_config = getattr(model_config.pretrained_config, "quantization_config", None) + layer_quant_configs = model_config.quant_config_dict or {} + has_nvfp4_experts = any( + name.endswith(".mlp.experts") and config.quant_algo == QuantAlgo.NVFP4 + for name, config in layer_quant_configs.items() + ) + if ( + quant_config.quant_algo != QuantAlgo.MIXED_PRECISION + or not has_nvfp4_experts + or not isinstance(hf_quant_config, dict) + or hf_quant_config.get("quant_method") != "fp8" + or tuple(hf_quant_config.get("weight_block_size", ())) != (128, 128) + ): + return model_config + + default_exclude = ["*kv_b_proj*", "*k_b_proj*", "*eh_proj*"] + hf_exclude_modules = hf_quant_config.get("modules_to_not_convert") or [] + exclude_modules = list(dict.fromkeys(list(hf_exclude_modules) + default_exclude)) + fp8_quant_config = quant_config.model_copy( + deep=True, + update={ + "quant_algo": QuantAlgo.FP8_BLOCK_SCALES, + "group_size": 128, + "exclude_modules": exclude_modules, + }, + ) + fp8_quant_config.__dict__.pop("quant_mode", None) + fp8_quant_config.__dict__.pop("layer_quant_mode", None) + model_config.quant_config = fp8_quant_config + return model_config + + def _copy_deepseek_v4_fused_a_weight_scale( module: Linear, fused_a: torch.Tensor, fused_a_scale: torch.Tensor ) -> None: @@ -2483,6 +2523,7 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: } def __init__(self, model_config: ModelConfig[PretrainedConfig]): + model_config = _normalize_deepseek_v4_nvfp4_mixed_precision_config(model_config) self.mapping_with_cp = None # Note: Currently the usage of mapping is all over the place making its usage brittle # in this file. As a temporary WAR, we hold on to an original copy of mapping when CP diff --git a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md index b84f0e5c4577..97751dd47e20 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -251,5 +251,5 @@ When adding new components, use these reference implementations: - **Schedulers MUST NOT write `moe.repeat_idx`** — `repeat_idx` is wrapper state advanced once per `forward_impl` regardless of chunk count - **Do NOT allocate symmetric memory from `run_moe` in `FUSED_COMM` backends** — Symmetric-memory rendezvous is a build-time collective and is unsafe under PP / layer-skip or CUDA graph capture; allocate from `create_weights()` after `ConfigurableMoE` has synchronized EPLB-derived attributes. See `mega_moe/mega_moe_deepgemm.py` for the DG pattern and `mega_moe/mega_moe_cute_dsl.py:_alloc_symm_provider` for the NVSHMEM-equivalent provider. - **Do NOT add a new `FUSED_COMM` backend without a zero-token `quantize_input` regression test** — `FusedCommMoEScheduler` calls `quantize_input` for every chunk (including zero-token chunks) so each backend must return its own empty-tensor layout. See `tests/unittest/_torch/modules/moe/test_moe_backend.py::test_megamoe_deepgemm_quantize_input_zero_tokens` and `test_megamoe_cutedsl_quantize_input_zero_tokens` for the pattern. -- **Do NOT use a dataclass for an autotuner tactic without a tested `__repr__` round-trip** — `AutoTuner` serializes tactic values through `json.dumps`/`json.loads` and `eval(repr(tactic))`; a plain dataclass fails the `eval(repr(...))` check. Prefer a JSON-friendly **tuple of primitives or lists of primitives** (lists are JSON-friendly; tuples round-trip via `eval(repr(...))`). See `Sm100MegaMoENvfp4Runner` in `tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py` for the 6-tuple tactic pattern (mma_tiler/cluster_shape as `list[int]`, the rest as `bool`/`int`/`str`). The fallback tactic is built inline in `Sm100MegaMoENvfp4Runner.forward(tactic=-1)` from `DEFAULT_MEGAMOE_TACTIC`, not via a separate `fallback_tactic()` method. -- **Do NOT forget `distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL` on a multi-rank `FUSED_COMM` backend's `TuningConfig`** — Every EP rank must converge on the same compiled tactic for every chunk, otherwise the in-kernel NVLink dispatch barrier deadlocks. Reference: `Sm100MegaMoENvfp4Runner.get_tuning_config` and every multi-rank op in `cute_dsl_custom_ops.py`. +- **Do NOT use a dataclass for an autotuner tactic without a tested `__repr__` round-trip** — `AutoTuner` serializes tactic values through `json.dumps`/`json.loads` and `eval(repr(tactic))`; a plain dataclass fails the `eval(repr(...))` check. Prefer a JSON-friendly **tuple of primitives or lists of primitives** (lists are JSON-friendly; tuples round-trip via `eval(repr(...))`). See the tactic-representation comment block in `tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py` for the 8-tuple tactic pattern (mma_tiler/cluster_shape as `list[int]`, `epi_flag_batch` as a nested `(int, int)` tuple, the rest as `bool`/`int`/`str`; `_unpack_tactic` is the single source of truth for the field order). The fallback tactic is the token-aware `default_megamoe_tactic(num_tokens)` helper, selected by `Sm100MegaMoENvfp4Runner.forward(tactic=-1)`, not a separate `fallback_tactic()` method. +- **Use `distributed_tuning_strategy=DistributedTuningStrategy.MERGE` on a multi-rank `FUSED_COMM` backend's `TuningConfig`** — Every EP rank must converge on the same compiled tactic for every chunk, otherwise the in-kernel NVLink dispatch barrier deadlocks. `PARALLEL` can profile different tactics on different ranks and is unsafe for fused collectives. Reference: `Sm100MegaMoENvfp4Runner.get_tuning_config`. diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index 1f3e9b112dd0..bb24f7ba45ec 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -338,9 +338,11 @@ def create_moe_backend( ], f"swiglu_limit is not supported in {moe_cls.__name__}." if swiglu_limit_scalar is not None: + # MegaMoECuteDsl uses the scalar only as a fallback when no per-expert + # tensor limit is given (see the MegaMoE branch below). assert moe_cls in [ CutlassFusedMoE, TRTLLMGenFusedMoE, WideEPMoE, DeepGemmFusedMoE, - MegaMoEDeepGemm, CuteDslFusedMoE + MegaMoEDeepGemm, CuteDslFusedMoE, MegaMoECuteDsl ], f"swiglu_limit_scalar is not supported in {moe_cls.__name__}." if moe_cls == TRTLLMGenFusedMoE: @@ -519,7 +521,11 @@ def create_moe_backend( activation_type=activation_type, ) if moe_cls is MegaMoECuteDsl: - megamoe_kwargs["swiglu_limit"] = swiglu_limit + # ``_resolve_gate_up_clamp`` accepts tensor or scalar; fall back + # to the scalar form when only that was wired. + megamoe_kwargs["swiglu_limit"] = (swiglu_limit + if swiglu_limit is not None else + swiglu_limit_scalar) else: megamoe_kwargs["swiglu_limit_scalar"] = swiglu_limit_scalar return moe_cls(**megamoe_kwargs) diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py index 2d57eb24354a..4f01ec5ee562 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -31,8 +31,8 @@ * BF16 -> NVFP4 activation quantization (``quantize_input``) * ``run_moe`` boundary: stage activation + topk into the kernel ABI, build the ``MegaMoECuteDslWeightView`` from the quant method, call - ``torch.ops.trtllm.cute_dsl_megamoe_nvfp4_blackwell``, sum the - per-topk axis (form A), return ``(T, hidden)`` output. + ``torch.ops.trtllm.cute_dsl_megamoe_nvfp4_blackwell``, return the + ``(T, hidden)`` output (the kernel collapses the top-k axis). ``run_moe`` is a single unified path for both topologies. Only the SOURCE of the kernel's input/output buffers branches on ``ep_size``: @@ -49,7 +49,7 @@ ``peer_rank_ptr_mapper.map``. ``_acquire_buffers`` is the only branch point; staging, kernel launch, -and the host-side top-k reduction are identical across topologies. +and output finalization are identical across topologies. Remaining hard gate: @@ -68,6 +68,9 @@ from __future__ import annotations +import os +import socket +import weakref from dataclasses import dataclass from typing import Dict, List, Optional, Tuple, Union @@ -76,8 +79,11 @@ from tensorrt_llm._utils import get_sm_version, is_sm_100f from tensorrt_llm.logger import logger +from tensorrt_llm.math_utils import ceil_div from tensorrt_llm.models.modeling_utils import QuantAlgo +from ....autotuner import AutoTuner + # ``megamoe_activation_sf_bytes_per_row`` lives at module top of the # custom-op file (NOT inside its ``IS_MEGAMOE_OP_AVAILABLE`` gate), so # it is always importable. The provider / shared-workspace helpers used @@ -329,6 +335,10 @@ class MegaMoECuteDsl(MoE): _SUPPORTED_ACTIVATION_DTYPES = frozenset({torch.bfloat16}) + # Legal combine wire formats; must stay in sync with + # ``CombineFormat.parse`` in the kernel package's token_comm.py. + _SUPPORTED_COMBINE_FORMATS = frozenset({"bf16", "32e4m3xe8m0", "16e2m1xbf16"}) + # Kernel owns dispatch + GEMM1 + SwiGLU + GEMM2 + combine via the # CuteDSL three-stage dispatch primitives + NVLink barrier; the # scheduler must skip host-side comm and lockstep every chunk. @@ -486,29 +496,28 @@ def __init__( # topk-score application point. v2 default is the deepgemm graph # (apply_topk_in_fc1=True): the fused kernel folds the topk score into - # the SwiGLU output before the fc1-out NVFP4 quant and the host reduces - # combine_output.sum(dim=1). Kept as an internal backend constant until - # the transformers route is GPU-validated and promoted to MoeConfig. + # the SwiGLU output before the fc1-out NVFP4 quant. Kept as an internal + # backend constant until the transformers route is GPU-validated and + # promoted to MoeConfig. self.apply_topk_in_fc1 = True - # Cross-rank combine path. ``False`` (default): the FC2 epilogue writes - # ``combine_output`` directly (scattered symmetric writes back to the - # source rank). ``True``: the kernel stages FC2 output in a local - # ``fc2_output_workspace`` and a fused in-kernel NVLink ``token_back_by_push`` - # bulk-returns it to the source rank's ``combine_output`` -- faster for - # multi-rank EP at the cost of the extra (local) fc2_output_workspace + - # fc2_done_counter budget (auto-sized by ``get_workspace_sizes``). The - # ``combine_output`` shape / host ``.sum(dim=1)`` reduce are unchanged - # (those depend on ``in_kernel_fc2_reduce``, not this knob). Internal - # backend constant for now; flip to opt into the fused-combine path. - self.token_back_by_dispatch = False - - # FC2 output store path (codegen-time). ``True`` (default): non-bulk - # TMA store (upstream default). ``False``: bulk store path. Kept as an - # internal backend attribute so different shapes/cases can pick the - # cheaper store; it changes the generated kernel, so it is part of the - # runner ``unique_id`` / compile-cache key (never a per-call runtime kwarg). - self.non_ubulk_fc2_store = True + # Cross-rank combine wire format is selected with + # MEGAMOE_COMBINE_FORMAT (default bf16). It MUST be rank-identical: + # it sizes the symmetric workspace and picks the compiled kernel, so + # divergence desyncs the rendezvous / NVLink barrier. + combine_format = os.environ.get("MEGAMOE_COMBINE_FORMAT", "bf16") + if combine_format not in self._SUPPORTED_COMBINE_FORMATS: + raise ValueError( + f"MEGAMOE_COMBINE_FORMAT must be one of " + f"{sorted(self._SUPPORTED_COMBINE_FORMATS)}; got {combine_format!r}." + ) + self.combine_format = combine_format + + # AutoTuner tactic-sweep opt-in via MEGAMOE_TACTIC_AUTOTUNE=1 + # (bench-only knob; default OFF so serving warmup never pays the + # MERGE lockstep sweep). Must be set identically on ALL ranks (it + # gates collective tuning behavior). + self.tactic_autotune = os.environ.get("MEGAMOE_TACTIC_AUTOTUNE", "0") == "1" # SwiGLU clamp: map the model-provided per-layer ``swiglu_limit`` tensor # to the kernel's codegen-time scalar ``gate_up_clamp``. The MegaMoE @@ -526,6 +535,24 @@ def __init__( or 4096 ) + # Adaptive ``max_tokens_per_rank`` bucketing: one symmetric provider + + # kernel per ladder bucket; per launch pick the smallest bucket that + # fits the lockstep chunk max (see ``_select_launch_max_tokens``). The + # ladder cap is min(moe_max_num_tokens, engine max_num_tokens) -- both + # rank-identical, keeping the ladder rank-identical (the + # NVLink-barrier invariant). + per_rank_cap = self.max_num_tokens + _engine_max = int(getattr(model_config, "max_num_tokens", 0) or 0) + if _engine_max > 0: + per_rank_cap = min(per_rank_cap, _engine_max) + self._maxt_buckets = self._resolve_maxt_buckets(per_rank_cap) + # Per-bucket symmetric providers (bucket -> MegaMoeSymmMemProvider); + # filled in ``create_weights`` on the multi-rank EP path. + self._symm_providers = {} + # Lockstep per-chunk cross-rank max for the NEXT run_moe (set via + # ``set_adaptive_launch_tokens``); ``None`` -> full bucket. + self._active_launch_max_tokens: Optional[int] = None + # Resolve EP ProcessGroup at construction. Resolving at forward # time would be collective on a non-synchronous call stack and # deadlock under PP / layer-skip. Construction is globally @@ -540,24 +567,37 @@ def __init__( f"to single-rank degenerate mode at run_moe time." ) self._ep_pg = None + # Pure-DEP invariant checked at (rank-synchronous) construction: the + # op's barrier-reset fence runs on the WORLD group, so WORLD != EP + # would build fine and then fail or hang mid-serving. ValueError on + # purpose -- the ``except RuntimeError`` above is the single-rank + # fallback and must not swallow this. + if dist.is_available() and dist.is_initialized() and dist.get_world_size() != self.ep_size: + raise ValueError( + f"MegaMoECuteDsl requires pure DEP (torch.distributed WORLD " + f"== EP): the forward-time barrier-reset fence runs on the " + f"WORLD process group. Got WORLD={dist.get_world_size()}, " + f"EP={self.ep_size}." + ) - # Weight tensors are owned by the quant method. ``_symm_provider`` - # is the symmetric-memory provider for multi-rank EP execution; - # allocated build-time in ``create_weights`` (collective - # rendezvous), shared across MoE layers via the module-scope - # cache in ``cute_dsl_megamoe_custom_op.py``. ``None`` for the - # single-rank degenerate path. - self._symm_provider = None + # WEAK ref to the AutoTuner profiling scratch provider: the custom + # op's process-global cache is the SOLE strong owner, so + # ``release_megamoe_profiling_scratch()`` reclaims it and this ref + # goes dead (no per-module teardown). Materialized lazily by the + # deferred factory in ``_launch_megamoe_kernel``. + self._profiling_scratch_provider_ref = None + self._profiling_scratch_kwargs = None self._weights_loaded = False self._weights_created = False - self._post_load_done = False self.quant_method = None # Per-instance staging cache (key -> {tensor name: tensor}) and # last-staged-T tracker; together they implement the always-pad- # to-max_T launch contract by refreshing only the rows that - # changed between calls. + # changed between calls. ``_last_staged_T`` is keyed by bucket max_T: + # each bucket's staging buffer has its own row-reset history (a shared + # scalar would reset the wrong buffer when buckets alternate). self._local_staging_cache: Dict[Tuple, Dict[str, torch.Tensor]] = {} - self._last_staged_T: Optional[int] = None + self._last_staged_T: Dict[int, int] = {} if not model_config.skip_create_weights_in_init: self.create_weights() @@ -595,6 +635,62 @@ def _resolve_gate_up_clamp( ) return float(first.item()) + @staticmethod + def _resolve_maxt_buckets(max_num_tokens: int) -> List[int]: + """Resolve the adaptive ``max_tokens_per_rank`` bucket ladder. + + MUST be a pure function of ``max_num_tokens`` so every EP rank derives + the identical ladder (divergence breaks the cross-rank NVLink barrier): + {256, 1024, 4096} rungs below the per-rank cap plus the cap itself. + Kept small -- each rung multiplies the compile/capture/memory surface. + """ + full = int(max_num_tokens) + if full <= 0: + full = 4096 + cand = {b for b in (256, 1024, 4096) if b < full} + cand.add(full) + return sorted(cand) + + def _select_launch_max_tokens(self, chunk_max_tokens: Optional[int]) -> int: + """Pick the smallest ladder bucket that fits ``chunk_max_tokens``. + + ``chunk_max_tokens`` is the lockstep ``max(all_rank_num_tokens)`` for + the chunk (rank-identical) or ``None`` (-> full bucket); input and + ladder rank-identical => bucket rank-identical, so all ranks launch + the same compiled kernel against same-sized symmetric buffers. + + Tuning mode with the ``tactic_autotune`` opt-in ALWAYS takes the top + bucket (profiling scratch and tuned-cache keys are sized/keyed for + it); both gates are rank-identical so the bucket stays lockstep. + """ + buckets = self._maxt_buckets + if self.tactic_autotune and AutoTuner.get().is_tuning_mode: + return buckets[-1] + if chunk_max_tokens is None or chunk_max_tokens <= 0: + return buckets[-1] + for b in buckets: + if chunk_max_tokens <= b: + return b + # RAISE (rank-identical input -> lockstep failure). Clamping instead + # would raise only on ranks whose local count exceeds the clamp, + # desyncing the NVLink barrier and hanging the survivors. + raise RuntimeError( + f"MegaMoE-CuteDSL: chunk max_tokens_per_rank={chunk_max_tokens} exceeds " + f"the top adaptive bucket {buckets[-1]} (per-rank cap = " + f"min(moe_max_num_tokens, max_num_tokens)). Raise moe_max_num_tokens or " + f"reduce the per-rank chunk." + ) + + def set_adaptive_launch_tokens(self, chunk_max_tokens: Optional[int]) -> None: + """Scheduler hook: record the lockstep per-chunk cross-rank max token + count for the next ``run_moe``. + + MUST be ``max(all_rank_num_tokens)`` for the chunk (rank-identical) so + all ranks select the same bucket; ``None`` -> full bucket. Consumed + and reset by the next ``_run_moe``. + """ + self._active_launch_max_tokens = chunk_max_tokens + def _supports_load_balancer(self) -> bool: # Both static and dynamic EPLB are supported: the four MegaMoE- # format derived parameters (``mega_fc{1,2}_weight{,_sf}``) are @@ -663,12 +759,80 @@ def validate_configurable_moe(self, moe) -> None: # ------------------------------------------------------------------ # EP process-group resolution (no collective at forward time) # ------------------------------------------------------------------ + def _maybe_init_torch_dist_under_mpi(self): + """Bootstrap a torch.distributed NCCL WORLD group under the MPI orchestrator. + + The MPI orchestrator (trtllm-bench / mpirun) never calls + ``init_process_group``, but the symmetric-memory rendezvous needs a + ProcessGroup; for pure DEP a node-local NCCL WORLD suffices. No-op + under Ray or single-rank. + """ + if not dist.is_available() or dist.is_initialized(): + return + from tensorrt_llm._utils import mpi_comm, mpi_rank, mpi_world_size + + try: + world = mpi_world_size() + rank = mpi_rank() + except Exception as e: # not under MPI either -> leave uninitialized + logger.debug( + f"[MegaMoECuteDsl] MPI rank query failed ({e!r}); " + "skipping torch.distributed bootstrap." + ) + return + if world <= 1: + return + # Check pure DEP BEFORE any collective: under PP/CP hybrids some ranks + # never construct a MegaMoE module, so the bcast/init below would hang + # instead of failing loud. Rank-identical inputs -> lockstep raise. + if world != self.ep_size: + raise ValueError( + f"MegaMoECuteDsl requires pure DEP (MPI WORLD == EP) to " + f"bootstrap torch.distributed: got MPI world={world}, " + f"EP={self.ep_size}." + ) + + # Rank 0 draws a free port (a fixed one collides across same-node + # jobs). The bcast is UNCONDITIONAL and rank 0 alone decides + # (honoring ITS pre-set MASTER_* env): gating the collective on + # per-rank env presence would desync the communicator when the vars + # are exported on only a subset of ranks. No bind-race retry: a sound + # one needs cross-rank failure agreement, and the close()->bind() + # window is negligible next to the fixed-port collision this replaces. + def _pick_rendezvous(): + addr = os.environ.get("MASTER_ADDR") + port = os.environ.get("MASTER_PORT") + if addr and port: + return (addr, port) + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("", 0)) + port = sock.getsockname()[1] + try: + addr = socket.gethostbyname(socket.gethostname()) + except OSError: + # Unresolvable hostname (minimal containers / broken DNS): + # this bootstrap is documented single-node, loopback works. + addr = "127.0.0.1" + return (addr, str(port)) + + host, port = mpi_comm().bcast(_pick_rendezvous() if rank == 0 else None, root=0) + os.environ["MASTER_ADDR"] = str(host) + os.environ["MASTER_PORT"] = str(port) + logger.info( + f"[MegaMoECuteDsl] torch.distributed not initialized under MPI; " + f"bootstrapping NCCL WORLD group (rank={rank}/{world}, " + f"{os.environ['MASTER_ADDR']}:{os.environ['MASTER_PORT']}) for the " + f"EP rendezvous." + ) + dist.init_process_group(backend="cuda:nccl,cpu:gloo", rank=rank, world_size=world) + def _resolve_ep_pg(self): """Return the torch.distributed ProcessGroup for the EP sub-world. Mirrors :meth:`MegaMoEDeepGemm._resolve_ep_pg` so the two MegaMoE backends share the same fallback chain. """ + self._maybe_init_torch_dist_under_mpi() if not dist.is_available() or not dist.is_initialized(): raise RuntimeError( "MegaMoECuteDsl requires torch.distributed to be initialized " @@ -730,22 +894,30 @@ def create_weights(self): return # Step 1: build-time symmetric memory allocation (multi-rank only). # Single-rank degenerate uses local CUDA tensors and skips here. - self._symm_provider = None if self.ep_size > 1: - self._symm_provider = self._alloc_symm_provider() + self._alloc_symm_provider() # Step 2-3: quant method registers all NVFP4 + MegaMoE-format params. self.quant_method = self._get_quant_method() self.quant_method.create_weights(self) # Step 4. self._weights_created = True + @property + def _symm_provider(self): + """Full-bucket symmetric provider, or ``None`` (pre-``create_weights`` + or single-rank) -- callers use this as the provider-available guard. + """ + return self._symm_providers.get(self._maxt_buckets[-1]) + def _alloc_symm_provider(self): """Build-time symmetric provider allocation. See ``create_weights``. - Returns a :class:`MegaMoeSymmMemProvider` from the module-scope - cache. Raises :class:`MegaMoeCuteDslUnavailable` with an - actionable message when no ProcessGroup is available -- that - would block the rendezvous and is a hard error for multi-rank. + One provider per adaptive bucket. Each ``get_megamoe_symm_provider`` + call is a COLLECTIVE rendezvous, so the loop must iterate identical + buckets in identical order on every EP rank (guaranteed: the ladder is + a pure function of rank-identical inputs); the module-scope cache + makes only the first layer pay. Raises MegaMoeCuteDslUnavailable when + no ProcessGroup is available (hard error for multi-rank). """ from ....custom_ops.cute_dsl_megamoe_custom_op import ( get_megamoe_symm_provider, @@ -759,26 +931,58 @@ def _alloc_symm_provider(self): "or initialize torch.distributed before model build." ) top_k = self.routing_method.experts_per_token - shared_workspace_bytes = query_megamoe_shared_workspace_bytes( - world_size=self.ep_size, - local_rank=self.ep_rank, - num_topk=top_k, - num_experts_per_rank=int(self.expert_size_per_partition), - hidden_size=self.hidden_size, - intermediate_size_per_partition=int(self.intermediate_size_per_partition), - expand_intermediate_size_per_partition=int(self.expand_intermediate_size_per_partition), - max_tokens_per_rank=int(self.max_num_tokens), - ) - return get_megamoe_symm_provider( - process_group=self._ep_pg, - world_size=self.ep_size, - rank=self.ep_rank, - hidden_size=self.hidden_size, - max_tokens_per_rank=int(self.max_num_tokens), - num_topk=top_k, - output_dtype=self.dtype or torch.bfloat16, - shared_workspace_bytes=shared_workspace_bytes, - ) + full_T = self._maxt_buckets[-1] + self._profiling_scratch_provider_ref = None + for max_T in self._maxt_buckets: + shared_workspace_bytes = query_megamoe_shared_workspace_bytes( + world_size=self.ep_size, + num_topk=top_k, + num_experts_per_rank=int(self.expert_size_per_partition), + hidden_size=self.hidden_size, + intermediate_size_per_partition=int(self.intermediate_size_per_partition), + expand_intermediate_size_per_partition=int( + self.expand_intermediate_size_per_partition + ), + max_tokens_per_rank=max_T, + in_kernel_fc2_reduce=False, + combine_format=self.combine_format, + ) + provider = get_megamoe_symm_provider( + process_group=self._ep_pg, + world_size=self.ep_size, + rank=self.ep_rank, + hidden_size=self.hidden_size, + max_tokens_per_rank=max_T, + num_topk=top_k, + output_dtype=self.dtype or torch.bfloat16, + combine_format=self.combine_format, + shared_workspace_bytes=shared_workspace_bytes, + ) + self._symm_providers[max_T] = provider + # AutoTuner profiling scratch: record only the request kwargs + # (full bucket only, matching tuning mode's forced top bucket); + # the multi-GiB allocation itself is deferred to the op's + # profiling pre-hook via the factory in ``_launch_megamoe_kernel``, + # so a process that never tunes never pays it. + if max_T == full_T and self.tactic_autotune: + self._profiling_scratch_kwargs = dict( + process_group=self._ep_pg, + world_size=self.ep_size, + rank=self.ep_rank, + hidden_size=self.hidden_size, + max_tokens_per_rank=max_T, + num_topk=top_k, + output_dtype=self.dtype or torch.bfloat16, + combine_format=self.combine_format, + shared_workspace_bytes=shared_workspace_bytes, + ) + if len(self._maxt_buckets) > 1 and self.ep_rank == 0: + logger.info( + "[MegaMoECuteDsl] adaptive max_tokens_per_rank buckets=%s " + "(one symmetric provider + kernel per bucket; per-launch bucket " + "= smallest >= max(all_rank_num_tokens))", + self._maxt_buckets, + ) def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False) -> None: if self.quant_method is None: @@ -804,14 +1008,6 @@ def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False) self.quant_method.load_weights( self, weights, self.weight_loading_mode, allow_partial_loading=allow_partial_loading ) - # Eager loading path: ``FusedMoEMethodBase.load_weights`` already - # ran ``quant_method.process_weights_after_loading(self)`` at its - # tail. Mark the sentinel so a subsequent - # ``backend.process_weights_after_loading()`` becomes a no-op - # instead of re-stacking ``mega_fc*_weight*`` from - # already-finalised parent buffers. - if not allow_partial_loading: - self._post_load_done = True def post_load_weights(self) -> None: if self.quant_method is None: @@ -819,45 +1015,24 @@ def post_load_weights(self) -> None: self.transform_weights() self.cache_derived_state() - def process_weights_after_loading(self) -> None: - """Run quant-method weight transforms; idempotent across calls. - - The real MegaMoE-format build (``[w3|w1]`` cat, 16-atom gate/up - interleave, ``to_blocked`` swizzle, and ``fc1_norm_const`` setup) lives in - :meth:`NVFP4MegaMoECuteDslMethod.process_weights_after_loading`. - This hook must dispatch to that method directly so two paths - both reach it: - - * Eager loading (``allow_partial_loading=False``) -- fired by - ``FusedMoEMethodBase.load_weights`` itself. - * Partial loading (RLHF reload, etc.) -- ``load_weights`` - skips its tail call, so the caller invokes this hook on - ``ConfigurableMoE`` -> backend to finalise. - - ``_post_load_done`` keeps the call idempotent: a second - invocation after eager finalisation must not re-run the - transforms (``_build_mega_format_weights`` would re-stack - ``mega_fc*_weight*`` from already-finalised parent buffers). - """ - if getattr(self, "_post_load_done", False): - return + # Staged-hook guards: ConfigurableMoE delegates these straight to the + # backend, and the base MoE implementations dereference + # ``self.quant_method`` unguarded. + def transform_weights(self) -> None: if self.quant_method is None: self.create_weights() - self.quant_method.process_weights_after_loading(self) - self._post_load_done = True + super().transform_weights() - def pre_reload_weights(self) -> None: - """Reset cached state before a hot weight reload. + def cache_derived_state(self) -> None: + if self.quant_method is None: + self.create_weights() + super().cache_derived_state() - ``_post_load_done`` is cleared so the next ``process_weights_after_loading`` - re-runs the MegaMoE-format weight transforms over the new - checkpoint bytes. The symmetric-memory provider is forward-time - scratch that does not need to be re-rendezvoused on weight - reload; we keep it as-is to avoid an unnecessary collective. - """ - self._post_load_done = False - if self.quant_method is not None and hasattr(self.quant_method, "pre_reload_weights"): - self.quant_method.pre_reload_weights(self) + def pre_reload_weights(self) -> None: + raise NotImplementedError( + "MegaMoE-CuteDSL does not support hot weight reloading; its " + "source weights are released after initial packing." + ) def _build_weight_view(self) -> MegaMoECuteDslWeightView: """Bundle the MegaMoE-format weight tensors registered by the @@ -923,10 +1098,11 @@ def quantize_input( ) # ``fp4_quantize(is_sf_swizzled=False)`` returns LINEAR layout # ``(rows, ceil(hidden/16))`` with no column pad. The kernel TMA - # load needs ``round_up(ceil(hidden/16), 4)`` bytes per row, so + # load needs ``pad_up(ceil_div(hidden, 16), 4)`` bytes per row + # (== ``megamoe_activation_sf_bytes_per_row``), so # 32-aligned-but-not-64-aligned hidden sizes (1568, 1632, 2080) # come back 2 bytes short; pad the tail before returning. - raw_cols = (hidden + 15) // 16 + raw_cols = ceil_div(hidden, 16) x_sf_raw = x_sf.view(x_bf16.shape[0], raw_cols) if sf_cols == raw_cols: return x_fp4, x_sf_raw @@ -949,8 +1125,8 @@ def run_moe( Casts ``token_selected_experts`` to ``int64`` (the scheduler keeps ``int32`` for the EPLB stats kernel; the MegaMoE kernel reads ``topk_idx`` as Int64) and delegates the staging + kernel launch - to :meth:`_run_moe`. The host then sums the form-A - ``(T, top_k, hidden)`` combine output along the top-k axis. + to :meth:`_run_moe`, which returns the reduced ``(T, hidden)`` + output. """ del unused_kwargs if output_dtype is None: @@ -1000,7 +1176,7 @@ def run_moe( output_dtype=output_dtype, ) - def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype): + def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype, max_T: int): """Allocate (and cache) the per-instance local staging tensors. Always allocates ``topk_idx`` (the kernel reads it as a local-only @@ -1011,16 +1187,17 @@ def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype for ``ep_size == 1``; multi-rank pulls them from the symmetric provider's regions instead. - All staging tensors are sized to ``max_num_tokens`` along dim 0 - so the kernel's constexpr ``num_tokens`` matches the buffer-time - ``max_tokens_per_rank``. Diverging the two would make + All staging tensors are sized to ``max_T`` (the per-launch adaptive + bucket) along dim 0 so the kernel's constexpr ``num_tokens`` matches the + buffer-time ``max_tokens_per_rank``. Diverging the two would make ``_dispatch_prep`` round 3 (``MAX_SLOT_C = num_tokens * num_topk`` in dispatch_kernel.py) write per-(expert, rank) advertise cards at the wrong stride relative to the symm allocation (``max_tokens_per_rank * num_topk`` in megamoe_kernel.py), - silently corrupting multi-rank metadata. + silently corrupting multi-rank metadata. The cache is keyed by + ``max_T`` so each bucket owns a distinct staging set. """ - max_T = int(self.max_num_tokens) + max_T = int(max_T) cache_key = (max_T, top_k, hidden, str(device), output_dtype) cached = self._local_staging_cache if cache_key in cached: @@ -1045,8 +1222,11 @@ def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype staging["activation_sf"] = torch.empty( (max_T, sf_bytes_per_row), dtype=torch.uint8, device=device ) + # Always 1: the kernel collapses the top-k axis in-op (TopkReduce + # form-A / REDG form-B); the op aliases this as the 2D output. + combine_k = 1 staging["combine_output"] = torch.empty( - (max_T, top_k, hidden), + (max_T, combine_k, hidden), dtype=torch.bfloat16, device=device, ) @@ -1058,7 +1238,6 @@ def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype shared_bytes = query_megamoe_shared_workspace_bytes( world_size=1, - local_rank=0, num_topk=top_k, num_experts_per_rank=int(self.expert_size_per_partition), hidden_size=hidden, @@ -1067,24 +1246,34 @@ def _ensure_local_staging(self, *, top_k: int, hidden: int, device, output_dtype self.expand_intermediate_size_per_partition ), max_tokens_per_rank=max_T, + in_kernel_fc2_reduce=False, + combine_format=self.combine_format, ) - staging["shared_workspace"] = torch.empty( + # zeros (not empty): the leading counter prefix is an atomic-add + # target the kernel only tail-resets for the NEXT launch, so the + # FIRST launch needs it pre-zeroed (one-time; cached per bucket). + staging["shared_workspace"] = torch.zeros( shared_bytes, dtype=torch.uint8, device=device ) cached[cache_key] = staging return staging - def _acquire_buffers(self, *, top_k: int, hidden: int, device, output_dtype) -> _MegaMoeBuffers: - """Resolve the kernel's input/output buffers. + def _acquire_buffers( + self, *, top_k: int, hidden: int, device, output_dtype, launch_max_T: int + ) -> _MegaMoeBuffers: + """Resolve the kernel's input/output buffers for the ``launch_max_T`` + adaptive bucket. This is the ONLY structural branch between single-rank and multi- rank execution; the source of activation / activation_sf / topk_weights / combine_output / shared_workspace differs per the :class:`_MegaMoeBuffers` contract. ``topk_idx_local`` always lives - in plain CUDA memory. + in plain CUDA memory. ``launch_max_T`` selects both the local staging + set and (multi-rank) the symmetric provider, which MUST share the same + ``max_T`` as the kernel's compile-time ``max_tokens_per_rank``. """ staging = self._ensure_local_staging( - top_k=top_k, hidden=hidden, device=device, output_dtype=output_dtype + top_k=top_k, hidden=hidden, device=device, output_dtype=output_dtype, max_T=launch_max_T ) if self.ep_size == 1: return _MegaMoeBuffers( @@ -1111,13 +1300,16 @@ def _acquire_buffers(self, *, top_k: int, hidden: int, device, output_dtype) -> f"model_config.skip_create_weights_in_init was not set " f"without a follow-up create_weights() call." ) - if self._symm_provider.num_topk != top_k: + # ``launch_max_T`` comes from ``_select_launch_max_tokens``, so it is + # always a ladder member. + provider = self._symm_providers[launch_max_T] + if provider.num_topk != top_k: raise MegaMoeCuteDslUnavailable( f"MegaMoECuteDsl symm provider was built for top_k=" - f"{self._symm_provider.num_topk} but run_moe called with " + f"{provider.num_topk} but run_moe called with " f"top_k={top_k}; recreate the backend." ) - regions = self._symm_provider.get_regions() + regions = provider.get_regions() return _MegaMoeBuffers( activation=regions.activation, activation_sf=regions.activation_sf, @@ -1163,7 +1355,7 @@ def _stage_inputs( it. """ max_T = bufs.topk_idx_local.shape[0] - last_T = getattr(self, "_last_staged_T", None) + last_T = self._last_staged_T.get(max_T) if last_T is not None and last_T > num_tokens: bufs.topk_idx_local[num_tokens:last_T].fill_(-1) if num_tokens > 0: @@ -1173,7 +1365,7 @@ def _stage_inputs( bufs.topk_weights[:num_tokens, :top_k].copy_(topk_weights, non_blocking=True) if num_tokens < max_T: bufs.topk_weights[num_tokens:max_T, :top_k].zero_() - self._last_staged_T = num_tokens + self._last_staged_T[max_T] = num_tokens def _launch_megamoe_kernel( self, @@ -1192,53 +1384,104 @@ def _launch_megamoe_kernel( peer_offsets: List[int], num_tokens: int, output_dtype: torch.dtype, + launch_max_T: int, ) -> torch.Tensor: - """Launch the fused MegaMoE CuteDSL kernel and reduce form-A output. + """Launch the fused MegaMoE CuteDSL kernel and return its reduced output. Single-rank and multi-rank reach this point with identical kernel inputs; only the source of the staged buffers differs (decided - upstream by :meth:`_acquire_buffers`). The host-side top-k reduction - is the same across topologies. NVFP4 / FP8-SF dtype views happen - through the module-level :func:`_as_nvfp4` / :func:`_as_fp8_sf` - helpers (the kernel rejects raw uint8 byte tensors). + upstream by :meth:`_acquire_buffers`). NVFP4 / FP8-SF dtype views + happen through the module-level :func:`_as_nvfp4` / + :func:`_as_fp8_sf` helpers (the kernel rejects raw uint8 byte + tensors). """ - torch.ops.trtllm.cute_dsl_megamoe_nvfp4_blackwell( - activation=_as_nvfp4(activation), - activation_sf=_as_fp8_sf(activation_sf), - topk_idx=topk_idx, - topk_weights=topk_weights, - fc1_weight=_as_nvfp4(weight_view.fc1_weight), - fc1_weight_sf=_as_fp8_sf(weight_view.fc1_weight_sf), - fc2_weight=_as_nvfp4(weight_view.fc2_weight), - fc2_weight_sf=_as_fp8_sf(weight_view.fc2_weight_sf), - fc1_alpha=weight_view.fc31_alpha, - fc2_alpha=weight_view.fc2_alpha, - fc1_norm_const=weight_view.fc1_norm_const, - combine_output=combine_output, - shared_workspace=shared_workspace, - world_size=world_size, - local_rank=local_rank, - num_topk=top_k, - num_experts_per_rank=int(self.expert_size_per_partition), - hidden_size=hidden, - intermediate_size_per_partition=int(self.intermediate_size_per_partition), - expand_intermediate_size_per_partition=int(self.expand_intermediate_size_per_partition), - max_tokens_per_rank=int(self.max_num_tokens), - peer_offsets=peer_offsets, - apply_topk_in_fc1=bool(self.apply_topk_in_fc1), - gate_up_clamp=self.gate_up_clamp, - token_back_by_dispatch=bool(self.token_back_by_dispatch), - non_ubulk_fc2_store=bool(self.non_ubulk_fc2_store), + # Hand the symmetric profiling scratch (or a deferred factory) to the + # op for tuning-mode profiling launches; cleared in the ``finally`` so + # a later same-process call for a different shape never sees a stale + # scratch (ownership: see ``_profiling_scratch_provider_ref``). + from ....custom_ops.cute_dsl_megamoe_custom_op import ( + set_active_megamoe_profiling_scratch, + set_active_megamoe_profiling_scratch_factory, + ) + + scratch_provider = ( + self._profiling_scratch_provider_ref() + if self._profiling_scratch_provider_ref is not None + else None + ) + scratch_factory = None + if ( + scratch_provider is None + and self._profiling_scratch_kwargs is not None + and world_size > 1 + ): + if AutoTuner.get().is_tuning_mode: + # Tuning but no live provider (lazy-only, or an earlier + # engine's release freed it): hand a DEFERRED factory -- a + # cache-hitting tuning forward (e.g. the spec-decode draft) + # must never allocate the multi-GiB scratch. It runs in the + # op's profiling pre-hook (where the collective rendezvous is + # lockstep-safe) and refreshes this module's weakref. + from ....custom_ops.cute_dsl_megamoe_custom_op import get_megamoe_profiling_scratch + + def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): + provider = _get(**_self._profiling_scratch_kwargs) + _self._profiling_scratch_provider_ref = weakref.ref(provider) + return provider.get_regions() + + scratch_factory = _deferred_scratch_factory + set_active_megamoe_profiling_scratch( + scratch_provider.get_regions() + if (scratch_provider is not None and world_size > 1) + else None ) + set_active_megamoe_profiling_scratch_factory(scratch_factory) + try: + torch.ops.trtllm.cute_dsl_megamoe_nvfp4_blackwell( + activation=_as_nvfp4(activation), + activation_sf=_as_fp8_sf(activation_sf), + topk_idx=topk_idx, + topk_weights=topk_weights, + fc1_weight=_as_nvfp4(weight_view.fc1_weight), + fc1_weight_sf=_as_fp8_sf(weight_view.fc1_weight_sf), + fc2_weight=_as_nvfp4(weight_view.fc2_weight), + fc2_weight_sf=_as_fp8_sf(weight_view.fc2_weight_sf), + fc1_alpha=weight_view.fc31_alpha, + fc2_alpha=weight_view.fc2_alpha, + fc1_norm_const=weight_view.fc1_norm_const, + combine_output=combine_output, + shared_workspace=shared_workspace, + world_size=world_size, + local_rank=local_rank, + num_topk=top_k, + num_experts_per_rank=int(self.expert_size_per_partition), + hidden_size=hidden, + intermediate_size_per_partition=int(self.intermediate_size_per_partition), + expand_intermediate_size_per_partition=int( + self.expand_intermediate_size_per_partition + ), + # Bucket max_T == the staging/symm buffer leading dim above. + max_tokens_per_rank=int(launch_max_T), + peer_offsets=peer_offsets, + apply_topk_in_fc1=bool(self.apply_topk_in_fc1), + gate_up_clamp=self.gate_up_clamp, + # Keep the deterministic standalone TopkReduce until form-B + # has dedicated GPU correctness and performance coverage. + in_kernel_fc2_reduce=False, + combine_format=self.combine_format, + tactic_autotune=bool(self.tactic_autotune), + num_tokens=num_tokens, + ) + finally: + # Always clear: a raising op call must not leave the factory + # global alive (its closure strongly references this module). + set_active_megamoe_profiling_scratch(None) + set_active_megamoe_profiling_scratch_factory(None) if num_tokens == 0: return torch.empty((0, hidden), dtype=output_dtype, device=combine_output.device) - # Deepgemm graph (apply_topk_in_fc1=True): the kernel already folded - # the topk score into the per-route BF16 terms, so the host reduce is - # a plain sum over the top-k axis. Accumulate in fp32 explicitly to - # match the design reference ``bf16(sum_fp32(term))`` and to be robust - # against any future change to the bf16 reduction accumulator type. - out = combine_output[:num_tokens].to(torch.float32).sum(dim=1).to(output_dtype) - return out + # The kernel already collapsed the top-k axis into the (T, 1, hidden) + # output (TopkReduce form-A / REDG form-B); just drop the singleton. + return combine_output[:num_tokens].squeeze(1).to(output_dtype) def _run_moe( self, @@ -1256,29 +1499,41 @@ def _run_moe( ) -> torch.Tensor: """Unified MegaMoE CuteDSL forward: acquire -> stage -> launch. - The kernel is always launched with ``T = max_num_tokens`` (its - compile-time constexpr); live tokens fill the first - ``num_tokens`` rows and the tail is masked via ``topk_idx == -1`` - (skipped by dispatch_kernel) and zero ``topk_weights`` (combine - stale-data guard). + The kernel is launched with ``T = launch_max_T`` -- the smallest + adaptive bucket >= the lockstep cross-rank chunk token count (set by + ``set_adaptive_launch_tokens``), capped at ``max_num_tokens``. Live + tokens fill the first ``num_tokens`` rows and the tail is masked via + ``topk_idx == -1`` (skipped by dispatch_kernel) and zero + ``topk_weights`` (combine stale-data guard). ``FusedCommMoEScheduler`` invariant 7 forces every EP rank to cross the NVLink barrier even with zero local tokens; only single-rank short-circuits ``num_tokens == 0`` because no peer is waiting. """ - if num_tokens > self.max_num_tokens: + # Pick and consume the adaptive bucket; un-hinted callers default to + # the full bucket. + launch_max_T = self._select_launch_max_tokens(self._active_launch_max_tokens) + self._active_launch_max_tokens = None + if num_tokens > launch_max_T: + # By construction num_tokens <= max(all_rank_num_tokens) <= + # launch_max_T; fail loudly rather than desync the barrier. raise RuntimeError( f"MegaMoECuteDsl run_moe got {num_tokens} tokens but the " - f"staging buffer is sized for {self.max_num_tokens}. Raise " - f"model_config.moe_max_num_tokens so peers do not read " - f"invalid rows." + f"selected adaptive bucket is sized for {launch_max_T} " + f"(buckets={self._maxt_buckets}, max_num_tokens=" + f"{self.max_num_tokens}). This indicates the scheduler's " + f"lockstep chunk-max hint diverged from the staged tokens." ) if num_tokens == 0 and self.ep_size == 1: return torch.empty((0, hidden), dtype=output_dtype, device=device) bufs = self._acquire_buffers( - top_k=top_k, hidden=hidden, device=device, output_dtype=output_dtype + top_k=top_k, + hidden=hidden, + device=device, + output_dtype=output_dtype, + launch_max_T=launch_max_T, ) self._stage_inputs( bufs=bufs, @@ -1304,4 +1559,5 @@ def _run_moe( peer_offsets=bufs.peer_offsets, num_tokens=num_tokens, output_dtype=output_dtype, + launch_max_T=launch_max_T, ) diff --git a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py index fe0153386763..05b4124ab5d9 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py +++ b/tensorrt_llm/_torch/modules/fused_moe/moe_scheduler.py @@ -1183,6 +1183,13 @@ def _forward_chunk( # quantize_input contracts. moe_input, x_sf = moe.backend.quantize_input(x_chunk_real) + # CuteDSL needs the scheduler's rank-identical chunk maximum to select + # one adaptive bucket on every EP rank; using a local token count could + # diverge and deadlock its in-kernel NVLink barrier. + set_adaptive = getattr(moe.backend, "set_adaptive_launch_tokens", None) + if set_adaptive is not None: + set_adaptive(max(all_rank_num_tokens) if all_rank_num_tokens else None) + # ----- MoE compute ----- # ``token_selected_slots`` is in [0, num_slots), matching the kernel's # ``num_experts`` template parameter (SymmBuffer / weights sized to diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index a39b53dcdbdb..2364aaf99eed 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -15,6 +15,7 @@ import inspect import math +import threading from abc import ABC, abstractmethod from enum import Enum, auto from typing import Dict, List, NamedTuple, Optional, Tuple, Union @@ -2094,6 +2095,11 @@ class NVFP4FusedMoEMethod(FusedMoEMethodBase): """ eplb_support_status = EplbSupportStatus.SUPPORTED + # Whether raw per-expert block-scale staging is an EPLB migration + # target. Children that migrate derived formats and free the raw + # sources (MegaMoE-CuteDSL) set this False. + _eplb_migrate_raw_block_scales = True + def get_weights_shapes(self, module: torch.nn.Module, weight_vec_size: int, block_scales_vec_size: int): # Divide by 16 because we use int64 to pack 16 fp4 values @@ -2647,11 +2653,14 @@ def process_weights_after_loading(self, module: torch.nn.Module): # before register_all_parameter_slot_and_to_fix_weight_fns below. self._prepare_shared_weight_scales_for_finalization(module) weight_fns = { - 'w3_w1_weight_scale': module.local_shared_w3_w1_scale_tensors, - 'w2_weight_scale': module.local_shared_w2_scale_tensors, 'fc31_alpha': shared_fc31_alpha, 'fc2_alpha': shared_fc2_alpha, } + if self._eplb_migrate_raw_block_scales: + weight_fns['w3_w1_weight_scale'] = ( + module.local_shared_w3_w1_scale_tensors) + weight_fns['w2_weight_scale'] = ( + module.local_shared_w2_scale_tensors) if shared_fc31_weight_scale_2 is not None: weight_fns['fc31_weight_scale_2'] = shared_fc31_weight_scale_2 if shared_fc2_weight_scale_2 is not None: @@ -3446,21 +3455,22 @@ class NVFP4MegaMoECuteDslMethod(NVFP4FusedMoEMethod): MegaMoE-format derived tensors, and fills the per-slot ``fc1_norm_const`` tensor from each expert's raw ``w2.input_scale``. - EPLB support is ``SUPPORTED``: dynamic EPLB migrates the four - ``mega_fc*_weight*`` derived parameters and per-expert - ``fc1_norm_const`` via CPU shared-staging buffers built in - :meth:`_build_mega_shared_staging` / :meth:`_build_fc1_norm_const` and - registered through :meth:`register_all_parameter_slot_and_to_fix_weight_fns`, - in addition to the standard NVFP4 family (``w3_w1_weight`` / - ``w2_weight`` / ``w*_weight_scale`` / ``fc*_alpha``) handled by the - base / grandparent classes. Slot migration replaces all raw + - MegaMoE-derived parameters atomically with byte-consistent values from - the source rank (the source built mega = transform(raw) once at - load time, so the migrated raw and mega bytes stay paired). + EPLB support is ``SUPPORTED`` with DERIVED-ONLY migration (the + MegaMoE-DeepGemm pattern): dynamic EPLB migrates the ``mega_fc*`` + derived parameters and ``fc1_norm_const`` via CPU shared staging, plus + the per-expert ``fc*_alpha`` / ``fc*_weight_scale_2`` handled by the + NVFP4 parent. The RAW source params are NOT migration targets (see + :meth:`_finalize_shared_weights`, ``_eplb_migrate_raw_block_scales``) + and are freed unconditionally after packing (``run_moe`` reads only + the mega buffers). Dynamic EPLB + hot weight RELOAD is unsupported. """ eplb_support_status = EplbSupportStatus.SUPPORTED + # Raw block scales are freed after packing; migrating them would slice + # 0-element placeholders (see class docstring). + _eplb_migrate_raw_block_scales = False + # On-device NVFP4 byte formats. Same constants the Cutlass child # uses; they describe the NVFP4 weight / FP8 block-scale packing, # not anything Cutlass-kernel-specific. @@ -3508,7 +3518,9 @@ def _round_up_int(a: int, b: int) -> int: @classmethod def fc1_sf_flat_size(cls, intermediate: int, hidden: int) -> int: """``round_up(expand_intermediate, SfPaddingBlock=128) * - round_up(ceil(hidden / 16), 4)`` -- matches kernel_fc12.py:880-890. + round_up(ceil(hidden / 16), 4)`` -- matches the FC1 weight-SF view + in ``kernel_fc12.py`` (``intermediate_gateup_padded`` / + ``expected_fc1_weight_sf_cols``). ``expand_intermediate = 2 * intermediate``. """ expand_intermediate = intermediate * 2 @@ -3518,11 +3530,25 @@ def fc1_sf_flat_size(cls, intermediate: int, hidden: int) -> int: @classmethod def fc2_sf_flat_size(cls, hidden: int, intermediate: int) -> int: """``round_up(hidden, SfPaddingBlock=128) * - round_up(ceil(intermediate / 16), 4)`` -- matches runner_fc12.py:1305. + round_up(ceil(intermediate / 16), 4)`` -- matches the FC2 weight-SF + view in ``kernel_fc12.py`` (``hidden_padded_fc2`` / + ``expected_fc2_weight_sf_cols``). """ return (cls._round_up_int(hidden, 128) * cls._round_up_int(cls._ceil_div_int(intermediate, 16), 4)) + # Source-checkpoint tensors kept as 0-element placeholders outside the + # load window so the full source set and the mega buffers never coexist + # (streaming load: peak = steady set + ONE layer of sources). + _STREAMED_SOURCE_PARAMS = ("w3_w1_weight", "w3_w1_weight_scale", + "w2_weight", "w2_weight_scale") + + # Serializes the transient source-set window (materialize -> load -> + # eager finalize) across MODULES: the generic loader runs module loads + # on a ThreadPoolExecutor, and concurrent windows would break the + # "peak = steady set + ONE layer" bound. Class-level on purpose. + _streamed_transient_lock = threading.Lock() + # ----------------------------------------------------------------- # create_weights: register MegaMoE-format parameters in addition to # the grandparent's standard NVFP4 parameters. @@ -3545,6 +3571,12 @@ def create_weights(self, module: torch.nn.Module): f"got expand_intermediate=" f"{module.expand_intermediate_size_per_partition}, " f"intermediate={module.intermediate_size_per_partition}.") + # Fail fast: the derived-only EPLB migration path registers no bias + # staging. + if module.bias and self.need_load_shared_weights(module): + raise NotImplementedError( + "NVFP4MegaMoECuteDslMethod does not support expert bias " + "together with dynamic-EPLB shared weight loading.") weight_vec_size = torch.iinfo(self.weight_dtype).bits // 4 self.block_scales_vec_size = torch.iinfo( @@ -3560,6 +3592,21 @@ def create_weights(self, module: torch.nn.Module): self.block_scales_dtype, self.block_scales_vec_size) + # Streaming load: shrink the source params to 0-element placeholders + # (full shapes preserved in rebuild_tensor_metadata); load_weights + # rematerializes one module at a time. Otherwise init materializes + # the full source set AND the mega buffers (~229 GB/rank). + for _name in self._STREAMED_SOURCE_PARAMS: + _p = getattr(module, _name) + replace_parameter_and_save_metadata( + module, _name, + nn.Parameter(torch.empty(0, dtype=_p.dtype), + requires_grad=False), + module.rebuild_tensor_metadata) + # Rebind quant_scales to the placeholders so the full-size CPU + # init tensors are actually released. + self.setup_quant_scales(module) + num_local_slots = module.expert_size_per_partition hidden = module.hidden_size intermediate = module.intermediate_size_per_partition @@ -3624,6 +3671,55 @@ def create_weights(self, module: torch.nn.Module): ) module.register_parameter("fc1_norm_const", fc1_norm_const) + def _materialize_source_params(self, module: torch.nn.Module): + """Rematerialize this module's streamed source params (full shape) + so the loader can fill them; no-op when already materialized. + """ + for name in self._STREAMED_SOURCE_PARAMS: + p = getattr(module, name, None) + if (p is not None and p.data.numel() == 0 + and name in module.rebuild_tensor_metadata): + meta = module.rebuild_tensor_metadata[name]['meta'] + module.register_parameter( + name, + nn.Parameter(torch.empty_like(meta, device="cuda"), + requires_grad=False)) + + def _finalize_shared_weights(self, module: torch.nn.Module): + # Derived-only EPLB migration: skip the base raw-weight + # registration -- the balancer would IndexError slicing the freed + # 0-element raw placeholders (see class docstring). + if not self.need_load_shared_weights(module): + return + if not hasattr(module, 'local_shared_w3_w1_tensors'): + # Already finalized (idempotent, mirroring the base contract). + return + if module.bias: + raise ValueError( + "MegaMoE-CuteDSL EPLB shared loading does not support " + "expert bias.") + delattr(module, 'local_shared_w3_w1_tensors') + delattr(module, 'local_shared_w2_tensors') + module.layer_load_balancer.host_tensor_sharer.finalize_layer_weights() + + def load_weights(self, + module: torch.nn.Module, + weights: List[Dict], + weight_loading_mode: MoEWeightLoadingMode, + allow_partial_loading: bool = False): + if allow_partial_loading: + raise NotImplementedError( + "MegaMoE-CuteDSL only supports full initial weight loading.") + + # The transient source-set window is serialized across modules + # (_streamed_transient_lock). + with NVFP4MegaMoECuteDslMethod._streamed_transient_lock: + self._materialize_source_params(module) + super().load_weights(module, + weights, + weight_loading_mode, + allow_partial_loading=allow_partial_loading) + # ----------------------------------------------------------------- # Loader overrides (4x @abstractmethod hooks on the grandparent). # Each one stashes the raw checkpoint shard in a tmp dict keyed by @@ -3638,23 +3734,18 @@ def load_expert_w3_w1_weight(self, dst_w3_w1_weight: torch.Tensor, allow_partial_loading: bool = False, expert_idx: int = -1): - if not allow_partial_loading: - assert w1_weight is not None and w3_weight is not None - if w1_weight is None and w3_weight is None: - return + assert w1_weight is not None and w3_weight is not None device = dst_w3_w1_weight.device - w1_weight_shard = load_weight_shard( - w1_weight, - module.tp_size, - module.tp_rank, - TensorParallelMode.COLUMN, - device=device) if w1_weight is not None else None - w3_weight_shard = load_weight_shard( - w3_weight, - module.tp_size, - module.tp_rank, - TensorParallelMode.COLUMN, - device=device) if w3_weight is not None else None + w1_weight_shard = load_weight_shard(w1_weight, + module.tp_size, + module.tp_rank, + TensorParallelMode.COLUMN, + device=device) + w3_weight_shard = load_weight_shard(w3_weight, + module.tp_size, + module.tp_rank, + TensorParallelMode.COLUMN, + device=device) if not hasattr(module, 'tmp_cutlass_w3_w1_weights'): module.tmp_cutlass_w3_w1_weights = {} @@ -3663,22 +3754,17 @@ def load_expert_w3_w1_weight(self, dict_key = (dst_base, expert_idx) expert_entry = module.tmp_cutlass_w3_w1_weights.setdefault(dict_key, {}) expert_entry['dst'] = dst_w3_w1_weight - if w1_weight_shard is not None: - expert_entry['w1'] = w1_weight_shard.contiguous().view( - dst_w3_w1_weight.dtype) - if w3_weight_shard is not None: - expert_entry['w3'] = w3_weight_shard.contiguous().view( - dst_w3_w1_weight.dtype) + expert_entry['w1'] = w1_weight_shard.contiguous().view( + dst_w3_w1_weight.dtype) + expert_entry['w3'] = w3_weight_shard.contiguous().view( + dst_w3_w1_weight.dtype) def load_expert_w2_weight(self, module: torch.nn.Module, w2_weight: torch.Tensor, dst_w2_weight: torch.Tensor, allow_partial_loading: bool = False): - if not allow_partial_loading: - assert w2_weight is not None - if w2_weight is None: - return + assert w2_weight is not None device = dst_w2_weight.device w2_weight_shard = load_weight_shard(w2_weight, module.tp_size, @@ -3690,6 +3776,9 @@ def load_expert_w2_weight(self, cast_w2_weight_shard = self._maybe_padding_shape( cast_w2_weight_shard, dst_w2_weight) dst_w2_weight.copy_(cast_w2_weight_shard, non_blocking=True) + if not hasattr(module, '_streamed_w2_covered'): + module._streamed_w2_covered = set() + module._streamed_w2_covered.add(dst_w2_weight.data_ptr()) def load_expert_w3_w1_weight_scale_nvfp4( self, @@ -3699,18 +3788,16 @@ def load_expert_w3_w1_weight_scale_nvfp4( dst_w3_w1_weight_scale: torch.Tensor, expert_idx: int = -1): device = dst_w3_w1_weight_scale.device - w1_weight_scale = load_weight_shard( - w1_weight_scale, - module.tp_size, - module.tp_rank, - TensorParallelMode.COLUMN, - device=device) if w1_weight_scale is not None else None - w3_weight_scale = load_weight_shard( - w3_weight_scale, - module.tp_size, - module.tp_rank, - TensorParallelMode.COLUMN, - device=device) if w3_weight_scale is not None else None + w1_weight_scale = load_weight_shard(w1_weight_scale, + module.tp_size, + module.tp_rank, + TensorParallelMode.COLUMN, + device=device) + w3_weight_scale = load_weight_shard(w3_weight_scale, + module.tp_size, + module.tp_rank, + TensorParallelMode.COLUMN, + device=device) if not hasattr(module, 'tmp_cutlass_w3_w1_weight_scales'): module.tmp_cutlass_w3_w1_weight_scales = {} @@ -3720,12 +3807,10 @@ def load_expert_w3_w1_weight_scale_nvfp4( expert_entry = module.tmp_cutlass_w3_w1_weight_scales.setdefault( dict_key, {}) expert_entry['dst'] = dst_w3_w1_weight_scale - if w3_weight_scale is not None: - expert_entry['w3'] = w3_weight_scale.contiguous().view( - dst_w3_w1_weight_scale.dtype) - if w1_weight_scale is not None: - expert_entry['w1'] = w1_weight_scale.contiguous().view( - dst_w3_w1_weight_scale.dtype) + expert_entry['w3'] = w3_weight_scale.contiguous().view( + dst_w3_w1_weight_scale.dtype) + expert_entry['w1'] = w1_weight_scale.contiguous().view( + dst_w3_w1_weight_scale.dtype) def load_expert_w2_weight_scale_nvfp4(self, module: torch.nn.Module, w2_weight_scale: torch.Tensor, @@ -3750,6 +3835,9 @@ def load_expert_w2_weight_scale_nvfp4(self, module: torch.nn.Module, cast_w2_weight_scale = self._maybe_padding_shape( cast_w2_weight_scale, dst_w2_weight_scale) dst_w2_weight_scale.copy_(cast_w2_weight_scale) + if not hasattr(module, '_streamed_w2_scale_covered'): + module._streamed_w2_scale_covered = set() + module._streamed_w2_scale_covered.add(dst_w2_weight_scale.data_ptr()) @staticmethod def _maybe_padding_shape(source_tensor: torch.Tensor, @@ -3781,7 +3869,111 @@ def _maybe_padding_shape(source_tensor: torch.Tensor, # mega-format CPU staging with the load balancer. ``fc1_norm_const`` is # built before the parent deletes raw input-scale staging. # ----------------------------------------------------------------- + def _streamed_coverage(self, module: torch.nn.Module) -> Dict[str, int]: + """Per-component count of routed local experts that received data. + A w3_w1 stash entry counts only when BOTH halves arrived; the + direct-copy w2 paths are tracked via row-pointer sets. Routed + slots are told apart from EPLB shared staging by storage base. + """ + + def _stash_covered(stash_name: str, param) -> int: + base = param.data.storage().data_ptr() + stash = getattr(module, stash_name, {}) + return sum(1 for (b, _idx), e in stash.items() + if b == base and 'w1' in e and 'w3' in e) + + def _rows_covered(set_name: str, param) -> int: + # Tensor (not storage) base: the recorded row addresses came + # from ``dst.data_ptr()`` of slices of this tensor. + base = param.data.data_ptr() + end = base + param.data.numel() * param.data.element_size() + ptrs = getattr(module, set_name, ()) + return sum(1 for p in ptrs if base <= p < end) + + return { + 'w3_w1_weight': + _stash_covered('tmp_cutlass_w3_w1_weights', module.w3_w1_weight), + 'w3_w1_weight_scale': + _stash_covered('tmp_cutlass_w3_w1_weight_scales', + module.w3_w1_weight_scale), + 'w2_weight': + _rows_covered('_streamed_w2_covered', module.w2_weight), + 'w2_weight_scale': + _rows_covered('_streamed_w2_scale_covered', module.w2_weight_scale), + } + + def _check_initial_aux_scale_coverage(self, + module: torch.nn.Module) -> None: + """Reject partially populated NVFP4 auxiliary-scale families.""" + n_slots = module.expert_size_per_partition + n_experts = module.num_experts + weight_scale_2 = getattr(module, 'tmp_weight_scale_2', None) or {} + raw_input_scales = getattr(module, 'tmp_raw_input_scales', None) or {} + + families = { + 'weight_scale_2': ( + sum(1 for entry in weight_scale_2.values() if entry), + sum(1 for entry in weight_scale_2.values() + if {'w1', 'w3', 'w2'} <= set(entry)), + n_slots, + ), + 'w1/w3 input_scale': ( + sum(1 for entry in raw_input_scales.values() + if 'w1' in entry or 'w3' in entry), + sum(1 for entry in raw_input_scales.values() + if 'w1' in entry and 'w3' in entry), + n_experts, + ), + 'w2 input_scale': ( + sum(1 for entry in raw_input_scales.values() if 'w2' in entry), + sum(1 for entry in raw_input_scales.values() if 'w2' in entry), + n_experts, + ), + } + incomplete = { + name: (complete, required) + for name, (present, complete, required) in families.items() + if present and complete < required + } + if incomplete: + raise RuntimeError( + "MegaMoE-CuteDSL initial load delivered partial auxiliary " + "scale families: " + + ", ".join(f"{name}={complete}/{required}" + for name, (complete, required) in incomplete.items())) + + has_weight_scale_2 = families['weight_scale_2'][0] > 0 + has_input_scale = (families['w1/w3 input_scale'][0] > 0 + or families['w2 input_scale'][0] > 0) + if has_input_scale and not has_weight_scale_2: + raise RuntimeError( + "MegaMoE-CuteDSL initial load provided input_scale without " + "weight_scale_2; both are required to derive expert alphas.") + if has_weight_scale_2 and not families['w2 input_scale'][0]: + raise RuntimeError( + "MegaMoE-CuteDSL initial load provided weight_scale_2 without " + "w2 input_scale; both are required to derive fc2 alpha.") + def process_weights_after_loading(self, module: torch.nn.Module): + if module.w3_w1_weight.data.numel() == 0: + return + + n_slots = module.expert_size_per_partition + coverage = self._streamed_coverage(module) + incomplete = {k: v for k, v in coverage.items() if v < n_slots} + if incomplete: + raise RuntimeError( + "MegaMoE-CuteDSL initial load left source components " + f"partially covered ({n_slots} local experts required per " + "component): " + ", ".join(f"{k}={v}/{n_slots}" + for k, v in incomplete.items()) + + ". The uncovered rows are uninitialized; full initial " + "loading must provide w1+w3+w2 weights and block scales for " + "every local expert.") + self._check_initial_aux_scale_coverage(module) + for attr in ('_streamed_w2_covered', '_streamed_w2_scale_covered'): + if hasattr(module, attr): + delattr(module, attr) # ---- Cat raw w3+w1 weights ---- # Iterates BOTH routed (module.w3_w1_weight.data) and shared # (module.local_shared_w3_w1_tensors) entries: the loader keys @@ -3850,6 +4042,31 @@ def process_weights_after_loading(self, module: torch.nn.Module): if self.need_load_shared_weights(module): self._register_mega_shared_staging(module) + # ---- Release the now-dead source NVFP4 routed weights ---- + # run_moe reads only the mega buffers, and EPLB migration targets + # are mega-format too (see class docstring), so the raw set is + # freed unconditionally. Rebinding quant_scales inside the shrink + # matters: otherwise every layer's freed source scales (~11 GB/rank) + # stay reachable until the model-wide sweep rebinds them. + self._shrink_streamed_source_params(module) + + def _shrink_streamed_source_params(self, module: torch.nn.Module): + """Re-shrink materialized streamed sources to 0-element placeholders + and rebind quant_scales to them. Idempotent. + """ + for _name in self._STREAMED_SOURCE_PARAMS: + _p = getattr(module, _name, None) + if _p is not None and _p.data.numel() > 0: + # On repeat calls replace_parameter_and_save_metadata + # re-registers the ORIGINAL saved placeholder; the device + # of the tensor passed here is effectively ignored. + replace_parameter_and_save_metadata( + module, _name, + nn.Parameter(torch.empty(0, dtype=_p.data.dtype), + requires_grad=False), + module.rebuild_tensor_metadata) + self.setup_quant_scales(module) + @staticmethod def _build_fc1_norm_const_tensor(raw_input_scales: Dict, expert_ids: List[int], @@ -3887,6 +4104,10 @@ def _build_fc1_norm_const(self, module: torch.nn.Module) -> None: module.fc1_norm_const.data.copy_( scalar.expand(num_local_slots).contiguous()) return + if not any('w2' in e for e in raw_input_scales.values()): + # Weights-only reload: the stash exists but carries no w2 input + # scales; keep the previously built values. + return routed_norm_const = self._build_fc1_norm_const_tensor( raw_input_scales, @@ -4061,12 +4282,18 @@ def _build_mega_format_weights(self, module: torch.nn.Module): mega_fc2_weight_sf}`` via :meth:`_build_mega_format_buffers` (the transform pipeline itself). """ + + # CPU-staged reload sources: upload ONE layer transiently so the + # pack pipeline runs on GPU (see _materialize_source_params). + def _on_cuda(t: torch.Tensor) -> torch.Tensor: + return t if t.is_cuda else t.cuda() + mega_fc1, mega_fc1_sf, mega_fc2, mega_fc2_sf = ( self._build_mega_format_buffers( - raw_w3_w1=module.w3_w1_weight.data, - raw_w3_w1_sf=module.w3_w1_weight_scale.data, - raw_w2=module.w2_weight.data, - raw_w2_sf=module.w2_weight_scale.data, + raw_w3_w1=_on_cuda(module.w3_w1_weight.data), + raw_w3_w1_sf=_on_cuda(module.w3_w1_weight_scale.data), + raw_w2=_on_cuda(module.w2_weight.data), + raw_w2_sf=_on_cuda(module.w2_weight_scale.data), num_slots=module.expert_size_per_partition, intermediate=module.intermediate_size_per_partition, hidden=module.hidden_size, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 50aa317e9e43..e3176be119f4 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1461,6 +1461,18 @@ def trtllm_gen_fmha_jit_warmup(): resource_manager=resource_manager) torch.cuda.synchronize() + @staticmethod + def _release_megamoe_profiling_scratch(): + # MegaMoE tuning resources are shared across layers, so only the engine + # can release them after its full autotune warmup and before graph + # capture. Later eviction could invalidate a captured workspace pointer. + from ..custom_ops import cute_dsl_megamoe_custom_op as _megamoe_op + release_megamoe_scratch = getattr(_megamoe_op, + "release_megamoe_profiling_scratch", + None) + if release_megamoe_scratch is not None: + release_megamoe_scratch() + def _run_autotuner_warmup(self, resource_manager: ResourceManager): """Runs a forward pass to populate the autotuner cache.""" if not self.llm_args.enable_autotuner: @@ -1514,6 +1526,8 @@ def _run_autotuner_warmup(self, resource_manager: ResourceManager): ) AutoTuner.get().print_profiling_cache() + self._release_megamoe_profiling_scratch() + # Clear workspace buffers allocated during the autotuner forward pass. # The autotuner runs a context-only forward with max_num_tokens, which # causes the global Buffers pool to cache large MoE/GEMM workspaces. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d06faa6e9c0d..89b40b9c8039 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1382,7 +1382,8 @@ class MoeConfig(StrictBaseModel): """Configuration for MoE.""" backend: Literal[ "AUTO", "CUTLASS", "CUTEDSL", "WIDEEP", "TRTLLM", "DEEPGEMM", - "DENSEGEMM", "VANILLA", "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM"] = Field( + "DENSEGEMM", "VANILLA", "TRITON", "MARLIN", "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL"] = Field( default='AUTO', description="MoE backend to use. " "AUTO selects default backend based on model. It currently doesn\'t always give the best choice for all scenarios. The capabilities of auto selection will be improved in future releases." diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index f5d0dd8925d2..e7c78a278c12 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -931,9 +931,10 @@ "VANILLA", "TRITON", "MARLIN", - "MEGAMOE_DEEPGEMM" + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL" ], - "annotation": "Literal['AUTO', 'CUTLASS', 'CUTEDSL', 'WIDEEP', 'TRTLLM', 'DEEPGEMM', 'DENSEGEMM', 'VANILLA', 'TRITON', 'MARLIN', 'MEGAMOE_DEEPGEMM']", + "annotation": "Literal['AUTO', 'CUTLASS', 'CUTEDSL', 'WIDEEP', 'TRTLLM', 'DEEPGEMM', 'DENSEGEMM', 'VANILLA', 'TRITON', 'MARLIN', 'MEGAMOE_DEEPGEMM', 'MEGAMOE_CUTEDSL']", "converter": "", "kind": "categorical", "path": "moe_config.backend" diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index d812dcf5d84e..8c89de01873d 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -114,6 +114,7 @@ l0_b200: - unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py - unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py # ------------- MoE: test_moe_backend (by backend) --------------- + - unittest/_torch/modules/moe/test_megamoe_streaming_load.py - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe # ------------- MoE: test_single_gpu (by backend) --------------- - unittest/_torch/modules/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and not None" diff --git a/tests/microbenchmarks/bench_moe/backend.py b/tests/microbenchmarks/bench_moe/backend.py index a4b682b5ab0b..866eabdd3774 100644 --- a/tests/microbenchmarks/bench_moe/backend.py +++ b/tests/microbenchmarks/bench_moe/backend.py @@ -40,6 +40,7 @@ class MoeBackendType(str, Enum): DEEPGEMM = "DEEPGEMM" DENSEGEMM = "DENSEGEMM" MEGAMOE_DEEPGEMM = "MEGAMOE_DEEPGEMM" + MEGAMOE_CUTEDSL = "MEGAMOE_CUTEDSL" @dataclass @@ -122,4 +123,8 @@ def get_backend_class(backend_type: MoeBackendType): from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoEDeepGemm return MegaMoEDeepGemm + if backend_type == MoeBackendType.MEGAMOE_CUTEDSL: + from tensorrt_llm._torch.modules.fused_moe.mega_moe import MegaMoECuteDsl + + return MegaMoECuteDsl raise ValueError(f"unknown MoE backend {backend_type!r}") diff --git a/tests/microbenchmarks/bench_moe/build.py b/tests/microbenchmarks/bench_moe/build.py index 0715bc783b23..36a3b5e27e27 100644 --- a/tests/microbenchmarks/bench_moe/build.py +++ b/tests/microbenchmarks/bench_moe/build.py @@ -49,6 +49,7 @@ "DeepGemmFusedMoE": "DEEPGEMM", "DenseGEMMFusedMoE": "DENSEGEMM", "MegaMoEDeepGemm": "MEGAMOE_DEEPGEMM", + "MegaMoECuteDsl": "MEGAMOE_CUTEDSL", "VanillaMoE": "VANILLA", } diff --git a/tests/microbenchmarks/bench_moe/search.py b/tests/microbenchmarks/bench_moe/search.py index f692911e7ff2..2d9098de590e 100644 --- a/tests/microbenchmarks/bench_moe/search.py +++ b/tests/microbenchmarks/bench_moe/search.py @@ -31,7 +31,7 @@ from .mapping import _PARALLEL_MODE_LAYOUTS, _resolve_mapping_layout from .specs import _ALL_BACKENDS, _FORCED_COMM_ENV_VALUES, ConfigSpec, ModelSpec, SearchSpec -_FUSED_COMM_BACKENDS = frozenset({"MEGAMOE_DEEPGEMM"}) +_FUSED_COMM_BACKENDS = frozenset({"MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"}) def _is_deepep_feasible(num_ranks: int) -> bool: @@ -170,10 +170,10 @@ def is_candidate_valid( "use TEP/DEP only with other backends" ) - # MegaMoEDeepGemm is EP-only (asserts moe_tp_size == 1 in __init__); DTP/TTP are invalid. - if config.backend.upper() == "MEGAMOE_DEEPGEMM" and moe_tp > 1: + # MegaMoE backends are EP-only; DTP/TTP are invalid. + if config.backend.upper() in _FUSED_COMM_BACKENDS and moe_tp > 1: return False, ( - f"MEGAMOE_DEEPGEMM does not support MoE-TP (moe_tp_size={moe_tp}); " + f"{config.backend.upper()} does not support MoE-TP (moe_tp_size={moe_tp}); " "use DEP/TEP modes only" ) diff --git a/tests/microbenchmarks/bench_moe/timing/autotune.py b/tests/microbenchmarks/bench_moe/timing/autotune.py index f993021be407..bdcf947f101b 100644 --- a/tests/microbenchmarks/bench_moe/timing/autotune.py +++ b/tests/microbenchmarks/bench_moe/timing/autotune.py @@ -44,6 +44,7 @@ def _run_autotune( The function always restores ``AutoTuner`` singleton state on exit so that ``--fast_autotune`` set for one case does not leak into the next. """ + tactic_autotune = bool(getattr(getattr(moe, "backend", moe), "tactic_autotune", False)) tuner = AutoTuner.get() saved_warmup = tuner.warmup saved_repeat = tuner.repeat @@ -63,3 +64,9 @@ def _run_autotune( tuner.warmup = saved_warmup tuner.repeat = saved_repeat tuner.stream_delay_micro_secs = saved_stream_delay + if tactic_autotune: + from tensorrt_llm._torch.custom_ops.cute_dsl_megamoe_custom_op import ( + release_megamoe_profiling_scratch, + ) + + release_megamoe_profiling_scratch() diff --git a/tests/microbenchmarks/bench_moe/utils.py b/tests/microbenchmarks/bench_moe/utils.py index d4874f5d91e1..40e7b0d75852 100644 --- a/tests/microbenchmarks/bench_moe/utils.py +++ b/tests/microbenchmarks/bench_moe/utils.py @@ -24,7 +24,7 @@ import torch import torch.distributed as dist -from tensorrt_llm._utils import local_mpi_rank, mpi_barrier, mpi_rank +from tensorrt_llm._utils import local_mpi_rank, mpi_barrier, mpi_comm, mpi_rank from .backend import MoeBackendType @@ -65,14 +65,23 @@ def _get_free_tcp_port() -> int: def _ensure_dist_for_megamoe(moe_backend: str, rank: int, world_size: int) -> None: """Initialize the torch.distributed NCCL ProcessGroup for MegaMoE.""" - if moe_backend.upper() != MoeBackendType.MEGAMOE_DEEPGEMM.value: + if moe_backend.upper() not in ( + MoeBackendType.MEGAMOE_DEEPGEMM.value, + MoeBackendType.MEGAMOE_CUTEDSL.value, + ): return if not torch.cuda.is_available(): raise RuntimeError("CUDA required for MegaMoE backend") if dist.is_initialized(): return - os.environ.setdefault("MASTER_ADDR", "127.0.0.1") - os.environ.setdefault("MASTER_PORT", str(_get_free_tcp_port())) + rendezvous = None + if rank == 0: + master_addr = os.environ.get("MASTER_ADDR") or "127.0.0.1" + master_port = os.environ.get("MASTER_PORT") or str(_get_free_tcp_port()) + rendezvous = (master_addr, master_port) + master_addr, master_port = mpi_comm().bcast(rendezvous, root=0) + os.environ["MASTER_ADDR"] = str(master_addr) + os.environ["MASTER_PORT"] = str(master_port) os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) os.environ["LOCAL_RANK"] = str(local_mpi_rank()) diff --git a/tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py b/tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py new file mode 100644 index 000000000000..b61dc8ce204a --- /dev/null +++ b/tests/unittest/_torch/modules/moe/test_megamoe_streaming_load.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for MegaMoE-CuteDSL NVFP4 streaming source weights.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +torch = pytest.importorskip("torch") + +from torch import nn # noqa: E402 + +if TYPE_CHECKING: + from tensorrt_llm._torch.modules.fused_moe.interface import MoEWeightLoadingMode + from tensorrt_llm._torch.modules.fused_moe.quantization import NVFP4MegaMoECuteDslMethod + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), + reason="needs 1 CUDA GPU (initial-load sources are rematerialized on cuda)", +) + +NUM_EXPERTS = 4 +HIDDEN_SIZE = 256 +INTERMEDIATE_SIZE = 64 + +_STREAMED_PARAMS = ("w3_w1_weight", "w3_w1_weight_scale", "w2_weight", "w2_weight_scale") + + +def _load_classes() -> tuple[type[MoEWeightLoadingMode], type[NVFP4MegaMoECuteDslMethod]]: + from tensorrt_llm._torch.modules.fused_moe.interface import MoEWeightLoadingMode + from tensorrt_llm._torch.modules.fused_moe.quantization import NVFP4MegaMoECuteDslMethod + + return MoEWeightLoadingMode, NVFP4MegaMoECuteDslMethod + + +class _StreamingMoEModule(nn.Module): + """Minimal single-rank module for the quant-method load path.""" + + def __init__(self, weight_loading_mode: MoEWeightLoadingMode) -> None: + super().__init__() + self.num_experts = NUM_EXPERTS + self.hidden_size = HIDDEN_SIZE + self.intermediate_size_per_partition = INTERMEDIATE_SIZE + self.expand_intermediate_size_per_partition = 2 * INTERMEDIATE_SIZE + self.expert_size_per_partition = NUM_EXPERTS + self.initial_local_expert_ids = list(range(NUM_EXPERTS)) + self.tp_size = 1 + self.tp_rank = 0 + self.ep_size = 1 + self.ep_rank = 0 + self.dtype = torch.bfloat16 + self.bias = False + self.weight_loading_mode = weight_loading_mode + # No EPLB in this test: need_load_shared_weights() must be False. + self.layer_load_balancer = None + + def _add_raw_shared_weights_for_unmap(self, weight_tensors: list[torch.Tensor]) -> None: + # Only forwards to the dynamic load balancer in production; no-op here. + del weight_tensors + + +def _w13_input_scale(expert_id: int) -> float: + return 0.5 + 0.125 * expert_id + + +def _w2_input_scale(expert_id: int) -> float: + return 0.25 + 0.0625 * expert_id + + +def _make_vanilla_weights(seed: int = 20260708) -> dict[str, torch.Tensor]: + gen = torch.Generator(device="cuda").manual_seed(seed) + + def rand_u8(*shape: int) -> torch.Tensor: + return torch.randint(0, 256, shape, dtype=torch.uint8, device="cuda", generator=gen) + + def rand_fp8(*shape: int) -> torch.Tensor: + # Any byte payload works (pure moves/reinterprets); stay clear of the + # 0x7f/0xff NaN encodings for hygiene. + raw = torch.randint(1, 120, shape, dtype=torch.uint8, device="cuda", generator=gen) + return raw.view(torch.float8_e4m3fn) + + weights = {} + for e in range(NUM_EXPERTS): + weights[f"{e}.w1.weight"] = rand_u8(INTERMEDIATE_SIZE, HIDDEN_SIZE // 2) + weights[f"{e}.w3.weight"] = rand_u8(INTERMEDIATE_SIZE, HIDDEN_SIZE // 2) + weights[f"{e}.w2.weight"] = rand_u8(HIDDEN_SIZE, INTERMEDIATE_SIZE // 2) + weights[f"{e}.w1.weight_scale"] = rand_fp8(INTERMEDIATE_SIZE, HIDDEN_SIZE // 16) + weights[f"{e}.w3.weight_scale"] = rand_fp8(INTERMEDIATE_SIZE, HIDDEN_SIZE // 16) + weights[f"{e}.w2.weight_scale"] = rand_fp8(HIDDEN_SIZE, INTERMEDIATE_SIZE // 16) + # w1/w3 input scales must match per expert (parent PWAL asserts). + weights[f"{e}.w1.input_scale"] = torch.tensor(_w13_input_scale(e), dtype=torch.float32) + weights[f"{e}.w3.input_scale"] = torch.tensor(_w13_input_scale(e), dtype=torch.float32) + weights[f"{e}.w2.input_scale"] = torch.tensor(_w2_input_scale(e), dtype=torch.float32) + # w1/w3 weight_scale_2 must match per expert (reconcile warns/maxes). + ws13 = torch.tensor(0.01 * (e + 1), dtype=torch.float32) + weights[f"{e}.w1.weight_scale_2"] = ws13 + weights[f"{e}.w3.weight_scale_2"] = ws13.clone() + weights[f"{e}.w2.weight_scale_2"] = torch.tensor(0.02 * (e + 1), dtype=torch.float32) + return weights + + +def _fresh( + seed: int = 20260708, +) -> tuple[ + NVFP4MegaMoECuteDslMethod, + _StreamingMoEModule, + dict[str, torch.Tensor], +]: + mode_cls, method_cls = _load_classes() + module = _StreamingMoEModule(mode_cls.VANILLA) + method = method_cls() + with torch.device("cuda"): + method.create_weights(module) + return method, module, _make_vanilla_weights(seed) + + +def _load( + method: NVFP4MegaMoECuteDslMethod, + module: _StreamingMoEModule, + bucket: dict[str, torch.Tensor], + allow_partial_loading: bool, +) -> None: + mode_cls, _ = _load_classes() + method.load_weights( + module, bucket, mode_cls.VANILLA, allow_partial_loading=allow_partial_loading + ) + + +def _streamed_numels(module: _StreamingMoEModule) -> dict[str, int]: + return {name: getattr(module, name).data.numel() for name in _STREAMED_PARAMS} + + +def _assert_sources_freed(module: _StreamingMoEModule, context: str = "") -> None: + numels = _streamed_numels(module) + assert all(n == 0 for n in numels.values()), ( + f"streamed source params should be 0-element placeholders {context}: {numels}" + ) + + +def _expected_mega_fc2(weights: dict[str, torch.Tensor]) -> torch.Tensor: + return torch.stack([weights[f"{e}.w2.weight"] for e in range(NUM_EXPERTS)]) + + +def _expected_mega_fc1(weights: dict[str, torch.Tensor]) -> torch.Tensor: + per_slot = [] + for e in range(NUM_EXPERTS): + gate = weights[f"{e}.w1.weight"].view(INTERMEDIATE_SIZE // 16, 16, HIDDEN_SIZE // 2) + up = weights[f"{e}.w3.weight"].view(INTERMEDIATE_SIZE // 16, 16, HIDDEN_SIZE // 2) + per_slot.append( + torch.stack([gate, up], dim=1).reshape(2 * INTERMEDIATE_SIZE, HIDDEN_SIZE // 2) + ) + return torch.stack(per_slot) + + +def _expected_fc1_norm_const() -> torch.Tensor: + return torch.tensor( + [1.0 / _w2_input_scale(e) for e in range(NUM_EXPERTS)], + dtype=torch.float32, + device="cuda", + ) + + +def test_initial_streaming_load_layer_atomic() -> None: + method, module, weights = _fresh() + + # create_weights replaces streamed sources with empty placeholders. + _assert_sources_freed(module, "right after create_weights") + + _load(method, module, weights, allow_partial_loading=False) + _assert_sources_freed(module, "after the initial eager load") + + assert torch.equal(module.mega_fc2_weight.data, _expected_mega_fc2(weights)) + assert torch.equal(module.mega_fc1_weight.data, _expected_mega_fc1(weights)) + assert torch.allclose(module.fc1_norm_const.data, _expected_fc1_norm_const()) + + +def test_partial_load_rejected_before_source_materialization() -> None: + method, module, weights = _fresh() + _assert_sources_freed(module, "before partial load") + + with pytest.raises(NotImplementedError, match="only supports full initial weight loading"): + _load(method, module, weights, allow_partial_loading=True) + + _assert_sources_freed(module, "after rejected partial load") diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 060472b5b8be..8838a9c21f55 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -476,6 +476,63 @@ class DummyModule: method.post_load_weights(DummyModule()) +def _make_megamoe_cutedsl_for_ctor_test() -> MegaMoECuteDsl: + model_config = ModelConfig( + mapping=Mapping(world_size=1, rank=0, tp_size=1, moe_tp_size=1, moe_ep_size=1), + moe_backend=MoeBackendType.MEGAMOE_CUTEDSL.value, + skip_create_weights_in_init=True, + ) + return MegaMoECuteDsl( + routing_method=RenormalizeMoeRoutingMethod(top_k=2), + num_experts=8, + hidden_size=512, + intermediate_size=512, + dtype=torch.bfloat16, + model_config=model_config, + init_load_balancer=False, + ) + + +def test_megamoe_cutedsl_tuning_mode_forces_top_maxt_bucket( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Profiling scratch is sized for the largest adaptive bucket. + monkeypatch.setenv("MEGAMOE_TACTIC_AUTOTUNE", "1") + moe = _make_megamoe_cutedsl_for_ctor_test() + buckets = moe._maxt_buckets + assert len(buckets) >= 2, f"expected a multi-bucket ladder, got {buckets}" + small_hint = buckets[0] + assert moe._select_launch_max_tokens(small_hint) == buckets[0] + monkeypatch.setattr(AutoTuner.get(), "is_tuning_mode", True) + assert moe._select_launch_max_tokens(small_hint) == buckets[-1] + + +def test_megamoe_cutedsl_tactic_autotune_defaults_off( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Standard serving must not pay for the 36-tactic sweep by default. + monkeypatch.delenv("MEGAMOE_TACTIC_AUTOTUNE", raising=False) + moe = _make_megamoe_cutedsl_for_ctor_test() + assert moe.tactic_autotune is False + + +def test_enumerate_megamoe_candidate_tactics_curated_space() -> None: + from tensorrt_llm._torch.custom_ops import cute_dsl_megamoe_custom_op as megamoe_op + + decode = megamoe_op.enumerate_megamoe_candidate_tactics(1024) + prefill = megamoe_op.enumerate_megamoe_candidate_tactics(16384) + assert len(decode) == len(prefill) == 36 + assert {t[-1] for t in decode} == {(1, 1)} + assert {t[-1] for t in prefill} == {(2, 4)} + # The deterministic fallback stays inside the curated axes. + for num_tokens in (64, 4096, 16384): + megamoe_op.validate_megamoe_tactic(megamoe_op.default_megamoe_tactic(num_tokens)) + invalid_tactic = list(megamoe_op.default_megamoe_tactic(64)) + invalid_tactic[2] = 511 + with pytest.raises(ValueError, match=r"group_hint must be an int >= 512"): + megamoe_op.validate_megamoe_tactic(tuple(invalid_tactic)) + + def run_backend_moe( backend: MoE, backend_type: MoeBackendType,