diff --git a/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh b/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh index 655852e19..d30f05069 100755 --- a/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh +++ b/examples/deepseek-v4/run_deepseek_v4_pro_muon_1gpu.sh @@ -421,7 +421,6 @@ for v in DOCKER_IMAGE NVTE_FUSED_ATTN NVTE_FUSED_ATTN_CK NVTE_FUSED_ATTN_AOTRITO PRIMUS_STACK_GROUPED_WEIGHT_TRITON PRIMUS_ROPE_TRITON \ PRIMUS_SINKHORN_TRITON PRIMUS_HC_TRITON PRIMUS_INDEXER_TRITON \ PRIMUS_INDEXER_TRITON_FULL PRIMUS_V4_ROUTER_TRITON \ - PRIMUS_TURBO_FUSE_GROUPED_WGRAD PRIMUS_TURBO_FUSE_WGRAD_DEBUG \ PRIMUS_MUON_BATCHED_NS PRIMUS_COMPRESS_ROPE_CACHE PRIMUS_COMPRESS_POOL_TRITON; do ENV_ARGS+=("--env" "$v") done diff --git a/primus/backends/megatron/core/extensions/_triton/__init__.py b/primus/backends/megatron/core/extensions/_triton/__init__.py index aa31d698f..7f8d5a0bb 100644 --- a/primus/backends/megatron/core/extensions/_triton/__init__.py +++ b/primus/backends/megatron/core/extensions/_triton/__init__.py @@ -5,10 +5,4 @@ ############################################################################### """Triton kernels for the Primus Megatron extensions package. - -Currently contains: - -* :mod:`stack_grouped_weight` — plan-6 P34's fused - ``torch.stack + transpose(1, 2) + contiguous`` for the per-expert weight - tensors of :class:`PrimusTurboGroupedMLP`. """ diff --git a/primus/backends/megatron/core/extensions/_triton/inplace_add.py b/primus/backends/megatron/core/extensions/_triton/inplace_add.py new file mode 100644 index 000000000..210095790 --- /dev/null +++ b/primus/backends/megatron/core/extensions/_triton/inplace_add.py @@ -0,0 +1,63 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +"""In-place ``dst += src`` Triton kernel for the Primus Megatron extensions. + +Used by the FP8/FP4 weight-gradient bridge in +:mod:`primus.backends.megatron.core.extensions.primus_turbo` to accumulate a +freshly produced weight gradient into the persistent ``main_grad`` buffer with +a single Triton launch (fp32 accumulate), instead of Torch's ``add_`` which +tiles a large consolidated grouped-expert ``main_grad`` of ``[E, N, K]`` into +multiple ~528M-element ``vectorized_elementwise_kernel`` launches. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _inplace_add_kernel(dst_ptr, src_ptr, n_elements, BLOCK: tl.constexpr): + """In-place ``dst += src`` over a flat buffer, accumulating in fp32. + + int64 offsets so a single launch covers tensors with > 2**31 elements + (e.g. the consolidated grouped-expert ``main_grad`` of [E, N, K]), + instead of Torch's ``add_`` which tiles into multiple ~528M-element + ``vectorized_elementwise_kernel`` launches. + """ + pid = tl.program_id(axis=0).to(tl.int64) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + d = tl.load(dst_ptr + offs, mask=mask, other=0.0).to(tl.float32) + s = tl.load(src_ptr + offs, mask=mask, other=0.0).to(tl.float32) + tl.store(dst_ptr + offs, (d + s).to(dst_ptr.dtype.element_ty), mask=mask) + + +def inplace_add_triton_(dst: torch.Tensor, src: torch.Tensor) -> torch.Tensor: + """``dst.add_(src)`` via a single Triton launch (fp32 accumulate). + + Falls back to Torch's ``add_`` when the layout is unsupported (not on CUDA + / non-contiguous / shape mismatch). The write is in-place on ``dst``'s + storage, so ``dst`` must be contiguous. + """ + if not dst.is_cuda or not dst.is_contiguous() or dst.numel() != src.numel(): + return dst.add_(src) + + dst_flat = dst.view(-1) + # reshape (not view) so a non-contiguous grad is materialized contiguously. + src_flat = src.reshape(-1) + n_elements = dst_flat.numel() + BLOCK = 8192 + grid = (triton.cdiv(n_elements, BLOCK),) + _inplace_add_kernel[grid](dst_flat, src_flat, n_elements, BLOCK=BLOCK) + return dst + + +__all__ = [ + "inplace_add_triton_", +] diff --git a/primus/backends/megatron/core/extensions/_triton/stack_grouped_weight.py b/primus/backends/megatron/core/extensions/_triton/stack_grouped_weight.py deleted file mode 100644 index 53bdea962..000000000 --- a/primus/backends/megatron/core/extensions/_triton/stack_grouped_weight.py +++ /dev/null @@ -1,413 +0,0 @@ -############################################################################### -# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. -# -# See LICENSE for license information. -############################################################################### - -"""Triton-fused ``torch.stack + transpose(1, 2) + contiguous`` for per-expert -GroupedMLP weights (plan-6 P34). - -The eager implementation in :class:`PrimusTurboGroupedMLP._stack_grouped_linear_weight`: - -.. code-block:: python - - weights = [getattr(module, f"weight{i}") for i in range(E)] - return torch.stack(weights, dim=0).transpose(1, 2).contiguous() - -does **two** full passes over the per-expert weight data: - -* ``torch.stack`` allocates ``[E, K, N]`` and issues ``E`` per-expert - ``copy_`` calls (one full pass aggregate); -* ``.contiguous()`` after ``.transpose(1, 2)`` allocates a second - ``[E, N, K]`` buffer and writes the transposed copy (a second full - pass). - -At V4-Flash EP=8 widths (``E=32``, fc1: ``K=4096, N=4096``; fc2: -``K=4096, N=2048``, bf16) the P32 final trace attributes -``hipMemcpyWithStream`` **289.6 ms / 32 calls** to this op chain -(``2 stack ops per layer × 8 layers × 2 (FWD + BWD VJP) = 32``). At ~9 ms -per call writing 512 MiB the effective bandwidth is only ~57 GB/s — far -below the MI355X HBM peak — because each call serializes E small -allocations plus a separate transpose copy. - -This module collapses the two passes into one Triton kernel that: - -1. Indexes per-expert weight tensors via an ``int64`` pointer tensor - (``weight_ptrs[e] = weights[e].data_ptr()``) — Triton's - ``tl.load(...).to(tl.pointer_type(...))`` idiom (same pattern as the - upstream grouped-GEMM tutorial). -2. Per program processes a ``[BLOCK_K, BLOCK_N]`` tile of one expert, - doing a tile-level transpose: reads ``weight[e][k, n]`` (row-major - ``[K, N]``) and writes ``out[e][n, k]`` (row-major ``[E, N, K]``). -3. BWD is the inverse — reads a ``[BLOCK_N, BLOCK_K]`` tile of ``dout - [e, n, k]`` and writes ``dweight[e][k, n]``. - -The two layouts form a bijection so **no atomics are needed in either -direction**. - -Nomenclature note (matches plan-6 P34 design doc): - -* ``K`` = ``N_out`` (``weight.shape[0]`` for ``nn.Linear`` — output features) -* ``N`` = ``N_in`` (``weight.shape[1]`` for ``nn.Linear`` — input features) - -The eager output is ``[E, N, K]`` (after ``.transpose(1, 2)`` on the -stacked ``[E, K, N]``); the Triton path produces the same layout. - -Gating: routed through :class:`PrimusTurboGroupedMLP._stack_grouped_linear_weight` -when ``PRIMUS_STACK_GROUPED_WEIGHT_TRITON != "0"`` (default-on). Set to -``"0"`` to fall back to the eager ``torch.stack + transpose + contiguous`` -chain (kept in tree for A/B testing and as the reference path for the -G37 unit tests). -""" - -from __future__ import annotations - -import os -from typing import List, Tuple - -import torch -import triton -import triton.language as tl - -# --------------------------------------------------------------------------- -# Triton dtype mapping -# --------------------------------------------------------------------------- - -_TORCH_TO_TL_DTYPE = { - torch.float64: tl.float64, - torch.float32: tl.float32, - torch.float16: tl.float16, - torch.bfloat16: tl.bfloat16, -} - - -def _triton_dtype(t: torch.dtype): - try: - return _TORCH_TO_TL_DTYPE[t] - except KeyError as exc: - raise TypeError( - f"stack_grouped_weight: unsupported dtype {t}; " f"expected one of {list(_TORCH_TO_TL_DTYPE)}" - ) from exc - - -# --------------------------------------------------------------------------- -# Triton kernels -# --------------------------------------------------------------------------- - - -@triton.jit -def _stack_grouped_weight_fwd_kernel( - WEIGHT_PTRS, # [E] int64 — tl.load() yields each expert's data_ptr() - OUT, # [E, N, K] contiguous output, row-major (strides [N*K, K, 1]) - E, - K, - N, - BLOCK_K: tl.constexpr, - BLOCK_N: tl.constexpr, - DTYPE: tl.constexpr, -): - """Per-expert ``[K, N] -> [N, K]`` tile-transpose, fused across experts. - - Each program writes one ``[BLOCK_K, BLOCK_N]`` tile of one expert's - output. Grid: ``(E, ceil(K / BLOCK_K), ceil(N / BLOCK_N))``. - - Read : ``weight[expert][k, n]`` from per-expert pointer, stride ``[N, 1]``. - Write : ``out[expert][n, k]``, stride ``[E*N*K -> N*K, K, 1]``. - - Both load and store carry bounds masks so non-multiple-of-BLOCK shapes - are supported. - """ - pid_e = tl.program_id(0) - pid_k = tl.program_id(1) - pid_n = tl.program_id(2) - - src_ptr = tl.load(WEIGHT_PTRS + pid_e).to(tl.pointer_type(DTYPE)) - - offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) - offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - mask_k = offs_k < K - mask_n = offs_n < N - - src_offsets = offs_k[:, None] * N + offs_n[None, :] - tile = tl.load( - src_ptr + src_offsets, - mask=mask_k[:, None] & mask_n[None, :], - other=0, - ) - - dst_offsets = pid_e * (N * K) + offs_n[None, :] * K + offs_k[:, None] - tl.store( - OUT + dst_offsets, - tile, - mask=mask_k[:, None] & mask_n[None, :], - ) - - -@triton.jit -def _stack_grouped_weight_bwd_kernel( - DWEIGHT_PTRS, # [E] int64 — tl.load() yields each expert's grad data_ptr() - DOUT, # [E, N, K] grad tensor (contiguous, same layout as FWD OUT) - E, - K, - N, - BLOCK_K: tl.constexpr, - BLOCK_N: tl.constexpr, - DTYPE: tl.constexpr, -): - """Inverse of the FWD: reads ``dout[expert][n, k]`` and writes - ``dweight[expert][k, n]``. - - Grid mirrors the FWD; the kernel is a pure bijection memcpy so no - atomics needed, and the BLOCK tile is the same shape. - """ - pid_e = tl.program_id(0) - pid_k = tl.program_id(1) - pid_n = tl.program_id(2) - - dst_ptr = tl.load(DWEIGHT_PTRS + pid_e).to(tl.pointer_type(DTYPE)) - - offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) - offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - mask_k = offs_k < K - mask_n = offs_n < N - - src_offsets = pid_e * (N * K) + offs_n[None, :] * K + offs_k[:, None] - tile = tl.load( - DOUT + src_offsets, - mask=mask_k[:, None] & mask_n[None, :], - other=0, - ) - - dst_offsets = offs_k[:, None] * N + offs_n[None, :] - tl.store( - dst_ptr + dst_offsets, - tile, - mask=mask_k[:, None] & mask_n[None, :], - ) - - -# --------------------------------------------------------------------------- -# Block-size autotune -# -# The kernel is bandwidth-bound; only a small handful of block sizes are -# worth scanning. ``BLOCK_K = BLOCK_N = 64`` is the default — it fits -# comfortably in LDS at bf16 (64 * 64 * 2 = 8 KiB / program) and matches -# the typical tile size MI355X HBM controllers like for transpose copies. -# Larger blocks ((128, 64), (64, 128), (128, 128)) widen the per-program -# tile and reduce launch count when E is small; smaller (32, 32) is -# kept as a safety floor for the fast-tier shapes. -# --------------------------------------------------------------------------- - -_BLOCK_CANDIDATES: Tuple[Tuple[int, int], ...] = ( - (32, 32), - (64, 64), - (128, 64), - (64, 128), -) - - -def _pick_block(K: int, N: int) -> Tuple[int, int]: - """Pick a ``(BLOCK_K, BLOCK_N)`` tile that divides reasonably into K, N. - - Conservative heuristic — we do not run a full autotune sweep at module - import to keep cold start cheap. The selected tile is always one of - :data:`_BLOCK_CANDIDATES`, picked as the largest block that does not - leave more than half the program's elements masked off on the small - dimension (i.e. the last program's tile-fill is at least 50 %). - """ - best: Tuple[int, int] = (64, 64) - best_score = -1.0 - for bk, bn in _BLOCK_CANDIDATES: - # Fraction of useful work in the last program along each axis; - # exact-multiple tiles get 1.0; partial tiles get the fraction. - k_full = (K // bk) * bk - k_tail = K - k_full - k_fill = 1.0 if k_tail == 0 else max(k_tail / bk, 0.5) - n_full = (N // bn) * bn - n_tail = N - n_full - n_fill = 1.0 if n_tail == 0 else max(n_tail / bn, 0.5) - # Score = tile area × tile-fill product → prefer bigger tiles when - # fill is high. - score = (bk * bn) * k_fill * n_fill - if score > best_score: - best_score = score - best = (bk, bn) - return best - - -# --------------------------------------------------------------------------- -# autograd.Function entry point -# --------------------------------------------------------------------------- - - -class StackGroupedWeightFn(torch.autograd.Function): - """Fused ``torch.stack(weights).transpose(1, 2).contiguous()`` with - in-kernel ``[K, N] -> [N, K]`` transpose, fused across all experts. - - Inputs (variadic): - ``*weights``: ``E`` per-expert tensors, each ``[K, N]``, all same - ``dtype`` and ``device``, all contiguous (Megatron's parameter - allocator always returns contiguous; a defensive assertion is - kept in :meth:`forward` regardless). - - Output: - ``[E, N, K]`` contiguous tensor — bit-identical to the eager - ``torch.stack(weights, dim=0).transpose(1, 2).contiguous()`` chain - (the operation is a pure layout transform so there is no fp - rounding to worry about). - - BWD returns one ``[K, N]`` grad tensor per input weight; PyTorch's - autograd then writes each into ``weights[i].grad``. - """ - - @staticmethod - def forward(ctx, *weights: torch.Tensor) -> torch.Tensor: - if not weights: - raise ValueError("StackGroupedWeightFn requires at least one weight tensor") - - first = weights[0] - if first.ndim != 2: - raise ValueError( - f"StackGroupedWeightFn: each weight must be 2D, got " f"weight0.shape={tuple(first.shape)}" - ) - K, N = int(first.shape[0]), int(first.shape[1]) - dtype = first.dtype - device = first.device - - for i, w in enumerate(weights): - if w.shape != first.shape: - raise ValueError( - f"StackGroupedWeightFn: weight{i}.shape={tuple(w.shape)} " - f"differs from weight0.shape={tuple(first.shape)}" - ) - if w.dtype is not dtype: - raise TypeError( - f"StackGroupedWeightFn: weight{i}.dtype={w.dtype} " f"differs from weight0.dtype={dtype}" - ) - if w.device != device: - raise RuntimeError( - f"StackGroupedWeightFn: weight{i}.device={w.device} " - f"differs from weight0.device={device}" - ) - if not w.is_contiguous(): - raise ValueError( - f"StackGroupedWeightFn: weight{i} must be contiguous; " - "Megatron's parameter allocator always returns contiguous " - "tensors so a non-contiguous input here indicates a bug " - "upstream" - ) - - E = len(weights) - - weight_ptrs = torch.tensor([w.data_ptr() for w in weights], dtype=torch.int64, device=device) - - out = torch.empty(E, N, K, dtype=dtype, device=device) - - block_k, block_n = _pick_block(K, N) - grid = (E, triton.cdiv(K, block_k), triton.cdiv(N, block_n)) - _stack_grouped_weight_fwd_kernel[grid]( - weight_ptrs, - out, - E, - K, - N, - BLOCK_K=block_k, - BLOCK_N=block_n, - DTYPE=_triton_dtype(dtype), - ) - - ctx.E = E - ctx.K = K - ctx.N = N - ctx.dtype = dtype - ctx.device = device - ctx.block_k = block_k - ctx.block_n = block_n - return out - - @staticmethod - def backward(ctx, dout: torch.Tensor): # type: ignore[override] - E = ctx.E - K = ctx.K - N = ctx.N - dtype = ctx.dtype - device = ctx.device - - if not dout.is_contiguous(): - dout = dout.contiguous() - if tuple(dout.shape) != (E, N, K): - raise ValueError( - f"StackGroupedWeightFn backward: dout.shape={tuple(dout.shape)} " - f"!= expected (E={E}, N={N}, K={K})" - ) - if dout.dtype is not dtype: - dout = dout.to(dtype) - - dweights: List[torch.Tensor] = [torch.empty(K, N, dtype=dtype, device=device) for _ in range(E)] - dweight_ptrs = torch.tensor([dw.data_ptr() for dw in dweights], dtype=torch.int64, device=device) - - block_k, block_n = ctx.block_k, ctx.block_n - grid = (E, triton.cdiv(K, block_k), triton.cdiv(N, block_n)) - _stack_grouped_weight_bwd_kernel[grid]( - dweight_ptrs, - dout, - E, - K, - N, - BLOCK_K=block_k, - BLOCK_N=block_n, - DTYPE=_triton_dtype(dtype), - ) - - return tuple(dweights) - - -# --------------------------------------------------------------------------- -# Public Python entry points -# --------------------------------------------------------------------------- - - -_ENV_FLAG = "PRIMUS_STACK_GROUPED_WEIGHT_TRITON" - - -def is_triton_path_enabled() -> bool: - """Returns ``True`` when the Triton path is active. - - Default-on; set ``PRIMUS_STACK_GROUPED_WEIGHT_TRITON=0`` to fall back - to the eager ``torch.stack + transpose + contiguous`` chain. Treated - as a soft env (not a model-config flag) so an operator can A/B the - Triton path on a live training job without re-launching with a - different YAML. - """ - return os.environ.get(_ENV_FLAG, "1") != "0" - - -def eager_stack_grouped_weight(weights: List[torch.Tensor]) -> torch.Tensor: - """Reference implementation — the exact eager chain the Triton path - replaces. Kept exported for unit tests and for the env-flag-off path - in :class:`PrimusTurboGroupedMLP`. - """ - return torch.stack(weights, dim=0).transpose(1, 2).contiguous() - - -def stack_grouped_weight(weights: List[torch.Tensor]) -> torch.Tensor: - """Dispatch entry point used by :class:`PrimusTurboGroupedMLP`. - - Routes through the Triton path when :func:`is_triton_path_enabled` - returns True; otherwise calls :func:`eager_stack_grouped_weight`. - The Triton path uses an :class:`autograd.Function` so BWD scatters - the gradient back to each ``weights[i].grad`` via the inverse - transpose kernel; the eager path inherits PyTorch's default VJP - chain for ``torch.stack + transpose + contiguous``. - """ - if is_triton_path_enabled(): - return StackGroupedWeightFn.apply(*weights) - return eager_stack_grouped_weight(weights) - - -__all__ = [ - "StackGroupedWeightFn", - "eager_stack_grouped_weight", - "is_triton_path_enabled", - "stack_grouped_weight", -] diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index 536cd22e7..05456b295 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/primus/backends/megatron/core/extensions/primus_turbo.py @@ -3,10 +3,9 @@ # # See LICENSE for license information. ############################################################################### -import contextlib import gc -import os from contextlib import contextmanager, nullcontext +from functools import lru_cache from typing import Callable, Iterable, List, Optional, Tuple, Union import primus_turbo.pytorch as primus_turbo_torch @@ -25,6 +24,7 @@ condition_init_method, ) from megatron.core.model_parallel_config import ModelParallelConfig +from megatron.core.num_microbatches_calculator import get_num_microbatches from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.parallel_state import ( get_context_parallel_group, @@ -88,53 +88,33 @@ from primus.core.pipeline_parallel.handler.offload_handler import OFFLOAD_BUFFER try: - import triton - import triton.language as tl + pass _HAVE_TRITON = True except (ImportError, ModuleNotFoundError): _HAVE_TRITON = False -_dummy_wgrads = {} - - -if _HAVE_TRITON: - - @triton.jit - def _inplace_add_kernel(dst_ptr, src_ptr, n_elements, BLOCK: tl.constexpr): - """In-place ``dst += src`` over a flat buffer, accumulating in fp32. +from primus.backends.megatron.core.extensions._triton.inplace_add import ( + inplace_add_triton_, +) - int64 offsets so a single launch covers tensors with > 2**31 elements - (e.g. the consolidated grouped-expert ``main_grad`` of [E, N, K]), - instead of Torch's ``add_`` which tiles into multiple ~528M-element - ``vectorized_elementwise_kernel`` launches. - """ - pid = tl.program_id(axis=0).to(tl.int64) - offs = pid * BLOCK + tl.arange(0, BLOCK) - mask = offs < n_elements - d = tl.load(dst_ptr + offs, mask=mask, other=0.0).to(tl.float32) - s = tl.load(src_ptr + offs, mask=mask, other=0.0).to(tl.float32) - tl.store(dst_ptr + offs, (d + s).to(dst_ptr.dtype.element_ty), mask=mask) +_dummy_wgrads = {} -def _triton_inplace_add_(dst: torch.Tensor, src: torch.Tensor) -> torch.Tensor: - """``dst.add_(src)`` via a single Triton launch (fp32 accumulate). +@lru_cache(maxsize=1) +def _is_gfx1250() -> bool: + """Return True iff the local GPU is a gfx1250 device. - Falls back to Torch's ``add_`` when Triton is unavailable or the layout is - unsupported (non-contiguous / shape mismatch). The write is in-place on - ``dst``'s storage, so ``dst`` must be contiguous. + Used to route ``main_grad`` accumulation through the single-launch Triton + in-place add (:func:`inplace_add_triton_`) only on gfx1250, keeping Torch's + ``add_`` on every other architecture. """ - if not _HAVE_TRITON or not dst.is_cuda or not dst.is_contiguous() or dst.numel() != src.numel(): - return dst.add_(src) - - dst_flat = dst.view(-1) - # reshape (not view) so a non-contiguous grad is materialized contiguously. - src_flat = src.reshape(-1) - n_elements = dst_flat.numel() - BLOCK = 8192 - grid = (triton.cdiv(n_elements, BLOCK),) - _inplace_add_kernel[grid](dst_flat, src_flat, n_elements, BLOCK=BLOCK) - return dst + if not torch.cuda.is_available(): + return False + try: + return "gfx1250" in torch.cuda.get_device_properties(0).gcnArchName + except Exception: + return False def _get_dummy_wgrad(shape: list, dtype: torch.dtype, zero=False) -> torch.Tensor: @@ -159,25 +139,6 @@ def _get_dummy_wgrad(shape: list, dtype: torch.dtype, zero=False) -> torch.Tenso return _dummy_wgrads[key].detach() -class _MainGradShim: - """Per-expert handle for primus_turbo's ``fused_grouped_wgrad`` over a - *consolidated* ``[E, N, K]`` grouped-expert weight (OPT-1). - - ``PrimusTurboGroupedLinear`` keeps one consolidated weight with a single - ``main_grad`` block, but ``fused_grouped_wgrad`` / ``_expert_main_grad_view`` - expect a list of per-expert handles, each exposing a 2-D ``main_grad`` and a - ``grad_added_to_main_grad`` flag. These shims point at the contiguous 2-D - slices ``main_grad[i]`` of the consolidated block, so the grouped GEMM - backward accumulates each expert's wgrad straight into the right slice. - """ - - __slots__ = ("main_grad", "grad_added_to_main_grad") - - def __init__(self, main_grad_slice: torch.Tensor) -> None: - self.main_grad = main_grad_slice - self.grad_added_to_main_grad = False - - def _bridge_weight_grad( x: torch.Tensor, weight: torch.nn.Parameter, weight_buffer: PrimusTurboQuantizedTensorPair ): @@ -207,33 +168,11 @@ def backward(ctx, grad_x, grad_quantized_weight, grad_quantized_weight_trans): weight, "grad_added_to_main_grad" ), "weight.grad_added_to_main_grad don't have grad_added_to_main_grad attribute." - # NOTE: Set weight.grad_added_to_main_grad to True to avoid adding the - # quantized weight gradient to main_grad twice. - if grad_quantized_weight is None: - # OPT-1 fused path: the grouped GEMM backward already accumulated the - # expert wgrad straight into main_grad (under fused_grouped_wgrad) and - # returned grad_b=None, so there is nothing to add here -- just flag it. - weight.grad_added_to_main_grad = True + if _is_gfx1250(): + inplace_add_triton_(weight.main_grad, grad_quantized_weight) else: - # `is_gfx1250` only exists in newer primus_turbo builds (it gates a - # gfx1250-specific elementwise-add workaround). Older / feature-branch - # primus_turbo that predate it (e.g. the flydsl sparse-MLA attention branch) - # don't define it; treat a missing symbol as False so those builds still work - # on non-gfx1250 archs (gfx942 / gfx950) instead of raising ImportError. - try: - from primus_turbo.pytorch.core.utils import is_gfx1250 - - _use_triton_inplace_add = is_gfx1250() - except ImportError: - _use_triton_inplace_add = False - - if _use_triton_inplace_add: - # NOTE: The bandwith of torch's elementwise add kernel has issue. Use triton to temporary workaround for gfx1250. - _triton_inplace_add_(weight.main_grad, grad_quantized_weight) - else: - weight.main_grad.add_(grad_quantized_weight) - - weight.grad_added_to_main_grad = True + weight.main_grad.add_(grad_quantized_weight) + weight.grad_added_to_main_grad = True return grad_x, _get_dummy_wgrad(list(weight.shape), weight.dtype), None, None @@ -1062,62 +1001,83 @@ def forward_internal( or quant_config.block_scaling() ), "Turbo FP8 is enabled but quant config is not mxfp8, current scaling, or block scaling." - if is_first_microbatch: - ( - self.quantized_weight_buffer, - self.quantized_weight_t_buffer, - ) = _maybe_create_quantized_weight_buffers( + if get_num_microbatches() == 1: + out = primus_turbo_torch.ops.gemm_fp8( + x, weight, - float8_e4m3, - quant_config, - disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), ) + else: + if is_first_microbatch: + ( + self.quantized_weight_buffer, + self.quantized_weight_t_buffer, + ) = _maybe_create_quantized_weight_buffers( + weight, + float8_e4m3, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache + or quant_config.current_scaling(), + ) - x, quantized_weight = _bridge_weight_grad( - x, - weight, - PrimusTurboQuantizedTensorPair( - data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer - ), - ) - out = primus_turbo_torch.ops.gemm_fp8( - x, - quantized_weight, - trans_a=False, - trans_b=True, - out_dtype=None, - config=quant_config.data(), - ) + x, quantized_weight = _bridge_weight_grad( + x, + weight, + PrimusTurboQuantizedTensorPair( + data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer + ), + ) + out = primus_turbo_torch.ops.gemm_fp8( + x, + quantized_weight, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), + ) elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled(): quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4." - if is_first_microbatch: - ( - self.quantized_weight_buffer, - self.quantized_weight_t_buffer, - ) = _maybe_create_quantized_weight_buffers( + if get_num_microbatches() == 1: + out = primus_turbo_torch.ops.gemm_fp4( + x, weight, - float4_e2m1fn_x2, - quant_config, - disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), ) + else: + if is_first_microbatch: + ( + self.quantized_weight_buffer, + self.quantized_weight_t_buffer, + ) = _maybe_create_quantized_weight_buffers( + weight, + float4_e2m1fn_x2, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + ) - x, quantized_weight = _bridge_weight_grad( - x, - weight, - PrimusTurboQuantizedTensorPair( - data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer - ), - ) - out = primus_turbo_torch.ops.gemm_fp4( - x, - quantized_weight, - trans_a=False, - trans_b=True, - out_dtype=None, - config=quant_config.data(), - ) + x, quantized_weight = _bridge_weight_grad( + x, + weight, + PrimusTurboQuantizedTensorPair( + data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer + ), + ) + out = primus_turbo_torch.ops.gemm_fp4( + x, + quantized_weight, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), + ) else: out = primus_turbo_torch.ops.gemm(x, weight, trans_a=False, trans_b=True, out_dtype=None) @@ -1233,62 +1193,83 @@ def forward_internal( or quant_config.block_scaling() ), "Turbo FP8 is enabled but quant config is not mxfp8, current scaling, or block scaling." - if is_first_microbatch: - ( - self.quantized_weight_buffer, - self.quantized_weight_t_buffer, - ) = _maybe_create_quantized_weight_buffers( + if get_num_microbatches() == 1: + out = primus_turbo_torch.ops.gemm_fp8( + x, weight, - float8_e4m3, - quant_config, - disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), ) + else: + if is_first_microbatch: + ( + self.quantized_weight_buffer, + self.quantized_weight_t_buffer, + ) = _maybe_create_quantized_weight_buffers( + weight, + float8_e4m3, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache + or quant_config.current_scaling(), + ) - x, quantized_weight = _bridge_weight_grad( - x, - weight, - PrimusTurboQuantizedTensorPair( - data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer - ), - ) - out = primus_turbo_torch.ops.gemm_fp8( - x, - quantized_weight, - trans_a=False, - trans_b=True, - out_dtype=None, - config=quant_config.data(), - ) + x, quantized_weight = _bridge_weight_grad( + x, + weight, + PrimusTurboQuantizedTensorPair( + data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer + ), + ) + out = primus_turbo_torch.ops.gemm_fp8( + x, + quantized_weight, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), + ) elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled(): quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4." - if is_first_microbatch: - ( - self.quantized_weight_buffer, - self.quantized_weight_t_buffer, - ) = _maybe_create_quantized_weight_buffers( + if get_num_microbatches() == 1: + out = primus_turbo_torch.ops.gemm_fp4( + x, weight, - float4_e2m1fn_x2, - quant_config, - disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), ) + else: + if is_first_microbatch: + ( + self.quantized_weight_buffer, + self.quantized_weight_t_buffer, + ) = _maybe_create_quantized_weight_buffers( + weight, + float4_e2m1fn_x2, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + ) - x, quantized_weight = _bridge_weight_grad( - x, - weight, - PrimusTurboQuantizedTensorPair( - data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer - ), - ) - out = primus_turbo_torch.ops.gemm_fp4( - x, - quantized_weight, - trans_a=False, - trans_b=True, - out_dtype=None, - config=quant_config.data(), - ) + x, quantized_weight = _bridge_weight_grad( + x, + weight, + PrimusTurboQuantizedTensorPair( + data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer + ), + ) + out = primus_turbo_torch.ops.gemm_fp4( + x, + quantized_weight, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), + ) else: out = primus_turbo_torch.ops.gemm(x, weight, trans_a=False, trans_b=True, out_dtype=None) @@ -1397,62 +1378,83 @@ def forward_internal( or quant_config.block_scaling() ), "Turbo FP8 is enabled but quant config is not mxfp8, current scaling, or block scaling." - if is_first_microbatch: - ( - self.quantized_weight_buffer, - self.quantized_weight_t_buffer, - ) = _maybe_create_quantized_weight_buffers( + if get_num_microbatches() == 1: + out = primus_turbo_torch.ops.gemm_fp8( + x, weight, - float8_e4m3, - quant_config, - disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), ) + else: + if is_first_microbatch: + ( + self.quantized_weight_buffer, + self.quantized_weight_t_buffer, + ) = _maybe_create_quantized_weight_buffers( + weight, + float8_e4m3, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache + or quant_config.current_scaling(), + ) - x, quantized_weight = _bridge_weight_grad( - x, - weight, - PrimusTurboQuantizedTensorPair( - data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer - ), - ) - out = primus_turbo_torch.ops.gemm_fp8( - x, - quantized_weight, - trans_a=False, - trans_b=True, - out_dtype=None, - config=quant_config.data(), - ) + x, quantized_weight = _bridge_weight_grad( + x, + weight, + PrimusTurboQuantizedTensorPair( + data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer + ), + ) + out = primus_turbo_torch.ops.gemm_fp8( + x, + quantized_weight, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), + ) elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled(): quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4." - if is_first_microbatch: - ( - self.quantized_weight_buffer, - self.quantized_weight_t_buffer, - ) = _maybe_create_quantized_weight_buffers( + if get_num_microbatches() == 1: + out = primus_turbo_torch.ops.gemm_fp4( + x, weight, - float4_e2m1fn_x2, - quant_config, - disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), ) + else: + if is_first_microbatch: + ( + self.quantized_weight_buffer, + self.quantized_weight_t_buffer, + ) = _maybe_create_quantized_weight_buffers( + weight, + float4_e2m1fn_x2, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + ) - x, quantized_weight = _bridge_weight_grad( - x, - weight, - PrimusTurboQuantizedTensorPair( - data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer - ), - ) - out = primus_turbo_torch.ops.gemm_fp4( - x, - quantized_weight, - trans_a=False, - trans_b=True, - out_dtype=None, - config=quant_config.data(), - ) + x, quantized_weight = _bridge_weight_grad( + x, + weight, + PrimusTurboQuantizedTensorPair( + data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer + ), + ) + out = primus_turbo_torch.ops.gemm_fp4( + x, + quantized_weight, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), + ) else: out = primus_turbo_torch.ops.gemm(x, weight, trans_a=False, trans_b=True, out_dtype=None) @@ -1574,62 +1576,83 @@ def forward_internal(self, x, is_first_microbatch: bool = False): or quant_config.block_scaling() ), "Turbo FP8 is enabled but quant config is not mxfp8, current scaling, or block scaling." - if is_first_microbatch: - ( - self.quantized_weight_buffer, - self.quantized_weight_t_buffer, - ) = _maybe_create_quantized_weight_buffers( + if get_num_microbatches() == 1: + out = primus_turbo_torch.ops.gemm_fp8( + inp, weight, - float8_e4m3, - quant_config, - disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), ) + else: + if is_first_microbatch: + ( + self.quantized_weight_buffer, + self.quantized_weight_t_buffer, + ) = _maybe_create_quantized_weight_buffers( + weight, + float8_e4m3, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache + or quant_config.current_scaling(), + ) - inp, quantized_weight = _bridge_weight_grad( - inp, - weight, - PrimusTurboQuantizedTensorPair( - data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer - ), - ) - out = primus_turbo_torch.ops.gemm_fp8( - inp, - quantized_weight, - trans_a=False, - trans_b=True, - out_dtype=None, - config=quant_config.data(), - ) + inp, quantized_weight = _bridge_weight_grad( + inp, + weight, + PrimusTurboQuantizedTensorPair( + data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer + ), + ) + out = primus_turbo_torch.ops.gemm_fp8( + inp, + quantized_weight, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), + ) elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled(): quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4." - if is_first_microbatch: - ( - self.quantized_weight_buffer, - self.quantized_weight_t_buffer, - ) = _maybe_create_quantized_weight_buffers( + if get_num_microbatches() == 1: + out = primus_turbo_torch.ops.gemm_fp4( + inp, weight, - float4_e2m1fn_x2, - quant_config, - disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), ) + else: + if is_first_microbatch: + ( + self.quantized_weight_buffer, + self.quantized_weight_t_buffer, + ) = _maybe_create_quantized_weight_buffers( + weight, + float4_e2m1fn_x2, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + ) - inp, quantized_weight = _bridge_weight_grad( - inp, - weight, - PrimusTurboQuantizedTensorPair( - data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer - ), - ) - out = primus_turbo_torch.ops.gemm_fp4( - inp, - quantized_weight, - trans_a=False, - trans_b=True, - out_dtype=None, - config=quant_config.data(), - ) + inp, quantized_weight = _bridge_weight_grad( + inp, + weight, + PrimusTurboQuantizedTensorPair( + data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer + ), + ) + out = primus_turbo_torch.ops.gemm_fp4( + inp, + quantized_weight, + trans_a=False, + trans_b=True, + out_dtype=None, + config=quant_config.data(), + ) else: out = primus_turbo_torch.ops.gemm(inp, weight, trans_a=False, trans_b=True, out_dtype=None) @@ -1806,77 +1829,35 @@ def forward_internal( quant_config.mxfp8_scaling() or quant_config.current_scaling() or quant_config.block_scaling() ), "Turbo FP8 is enabled but quant config is not mxfp8, current scaling, or block scaling." - if is_first_microbatch: - ( - self.quantized_weight_buffer, - self.quantized_weight_t_buffer, - ) = _maybe_create_quantized_weight_buffers( + if get_num_microbatches() == 1: + out = primus_turbo_torch.ops.grouped_gemm_fp8( + x, weights, - float8_e4m3, - quant_config, - disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + m_splits, + trans_b=True, + config=quant_config.data(), ) - - x, quantized_weights = _bridge_weight_grad( - x, - weights, - PrimusTurboQuantizedTensorPair( - data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer - ), - ) - - # OPT-1 (opt-in, single-GPU): accumulate the expert wgrad straight into - # main_grad in the grouped GEMM backward (beta=1 ACCUMULATE) instead of - # GEMM-wgrad -> _WeightGradBridge.main_grad.add_(). Per-expert shims point - # at the consolidated [E,N,K] main_grad slices; the grouped backward - # accumulates into them and returns grad_b=None, then the bridge backward - # just flags grad_added. ONLY safe with no gradient all-reduce / - # reduce-scatter (TP=1 / DP=1 / EP=1), since grad_b=None skips the reduce. - # Needs a turbo wheel carrying fused_grouped_wgrad (else falls back). - _wgrad_ctx = contextlib.nullcontext() - _fwg_dbg = ( - os.environ.get("PRIMUS_TURBO_FUSE_WGRAD_DEBUG") == "1" - and getattr(type(self), "_fwg_logn", 0) < 10 - ) - _fwg_flag = os.environ.get("PRIMUS_TURBO_FUSE_GROUPED_WGRAD", "0") == "1" - _mg = getattr(weights, "main_grad", None) - if _fwg_flag and _mg is not None and _mg.dim() == 3: - try: - from primus_turbo.pytorch.ops.grouped_gemm_fp8 import ( - _expert_main_grad_view, - fused_grouped_wgrad, + else: + if is_first_microbatch: + ( + self.quantized_weight_buffer, + self.quantized_weight_t_buffer, + ) = _maybe_create_quantized_weight_buffers( + weights, + float8_e4m3, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache + or quant_config.current_scaling(), ) - _shims = [_MainGradShim(_mg[i]) for i in range(_mg.shape[0])] - if _fwg_dbg: - import sys - - _v = _expert_main_grad_view(_shims) - print( - f"[OPT-1] gate PASS shape={tuple(_mg.shape)} contig={_mg.is_contiguous()} " - f"stride={_mg.stride()} view={'OK' if _v is not None else 'REJECTED'}", - file=sys.stderr, - flush=True, - ) - type(self)._fwg_logn = getattr(type(self), "_fwg_logn", 0) + 1 - _wgrad_ctx = fused_grouped_wgrad(_shims) - except ImportError as _e: - if _fwg_dbg: - import sys - - print(f"[OPT-1] ImportError: {_e}", file=sys.stderr, flush=True) - type(self)._fwg_logn = getattr(type(self), "_fwg_logn", 0) + 1 - elif _fwg_dbg: - import sys - - print( - f"[OPT-1] gate FAIL flag={_fwg_flag} main_grad={'None' if _mg is None else f'dim={_mg.dim()}'}", - file=sys.stderr, - flush=True, + x, quantized_weights = _bridge_weight_grad( + x, + weights, + PrimusTurboQuantizedTensorPair( + data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer + ), ) - type(self)._fwg_logn = getattr(type(self), "_fwg_logn", 0) + 1 - with _wgrad_ctx: out = primus_turbo_torch.ops.grouped_gemm_fp8( x, quantized_weights, @@ -1885,7 +1866,44 @@ def forward_internal( config=quant_config.data(), ) elif PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp4_enabled(): - assert False, "FP4 is not supported in PrimusTurboGroupedLinear" + quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() + assert quant_config.mxfp4_scaling(), "Turbo FP4 is enabled but quant config is not mxfp4." + + if get_num_microbatches() == 1: + out = primus_turbo_torch.ops.grouped_gemm_fp4( + x, + weights, + m_splits, + trans_b=True, + config=quant_config.data(), + ) + else: + if is_first_microbatch: + ( + self.quantized_weight_buffer, + self.quantized_weight_t_buffer, + ) = _maybe_create_quantized_weight_buffers( + weights, + float4_e2m1fn_x2, + quant_config, + disable_parameter_transpose_cache=self.disable_parameter_transpose_cache, + ) + + x, quantized_weights = _bridge_weight_grad( + x, + weights, + PrimusTurboQuantizedTensorPair( + data=self.quantized_weight_buffer, data_t=self.quantized_weight_t_buffer + ), + ) + + out = primus_turbo_torch.ops.grouped_gemm_fp4( + x, + quantized_weights, + m_splits, + trans_b=True, + config=quant_config.data(), + ) else: out = primus_turbo_torch.ops.grouped_gemm(x, weights, m_splits, trans_b=True) diff --git a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py index 9aeaf2cd4..42cae07be 100644 --- a/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py +++ b/primus/backends/megatron/core/transformer/v4_attention_kernels/_triton_common/rope_interleaved_partial.py @@ -691,8 +691,7 @@ def apply_rope_from_positions( def is_triton_path_enabled() -> bool: """Return True iff the ``PRIMUS_ROPE_TRITON`` env knob is not ``"0"``. - Mirrors :func:`primus.backends.megatron.core.extensions._triton.stack_grouped_weight.is_triton_path_enabled` - (plan-6 P34). Default-on, A/B toggle via ``PRIMUS_ROPE_TRITON=0``. + Default-on, A/B toggle via ``PRIMUS_ROPE_TRITON=0``. """ return os.environ.get("PRIMUS_ROPE_TRITON", "1") != "0"