|
|
@@ -1,7067 +0,0 @@
|
|
|
-"""CuTeDSL W4A16 NVFP4/BF16 W4A16 MoE kernels."""
|
|
|
-
|
|
|
-from __future__ import annotations
|
|
|
-
|
|
|
-import math
|
|
|
-import os
|
|
|
-from dataclasses import dataclass, replace
|
|
|
-from typing import NamedTuple
|
|
|
-
|
|
|
-import cuda.bindings.driver as cuda
|
|
|
-import cutlass
|
|
|
-import cutlass.cute as cute
|
|
|
-import torch
|
|
|
-from cutlass.cutlass_dsl import Int32, Int64, Uint32
|
|
|
-
|
|
|
-from b12x.cute.compiler import KernelCompileSpec, compile as b12x_compile
|
|
|
-from b12x.cute.fp4 import (
|
|
|
- atomic_add_global_i32,
|
|
|
- bf16_mma_m16n8k16_f32,
|
|
|
- bf16_mma_rhs_fragments_as_mma_a_m16n8k16_f32,
|
|
|
- bfloat2_broadcast_lane,
|
|
|
- bfloat2_mul,
|
|
|
- bfloat2_to_float2_scaled,
|
|
|
- broadcast_f32_to_half2,
|
|
|
- broadcast_f32_to_bfloat2,
|
|
|
- cp_async4_shared_global,
|
|
|
- cp_async4_shared_global_pred,
|
|
|
- fabs_f32,
|
|
|
- fmax_f32,
|
|
|
- f16_mma_m16n8k16_f32,
|
|
|
- f16_mma_rhs_fragments_as_mma_a_m16n8k16_f32,
|
|
|
- half2_to_float2_scaled,
|
|
|
- get_ptr_as_int64,
|
|
|
- half2_mul,
|
|
|
- ld_global_acquire_i32,
|
|
|
- ld_global_nc_u32,
|
|
|
- ld_global_v4_f32,
|
|
|
- ld_shared_f32,
|
|
|
- ld_shared_i32_relaxed,
|
|
|
- ld_shared_u32,
|
|
|
- ld_shared_v2_u32,
|
|
|
- ld_shared_v4_f32,
|
|
|
- ld_shared_v4_u32,
|
|
|
- ldmatrix_m8n8x4_b16,
|
|
|
- ldmatrix_m8n8x2_b16,
|
|
|
- packed_dequant_e2m1x4_to_bfloat2x2,
|
|
|
- packed_dequant_e2m1x4_to_half2x2,
|
|
|
- packed_dequant_e4m3x4_to_bfloat2x2,
|
|
|
- packed_dequant_e4m3x4_to_half2x2,
|
|
|
- packed_dequant_e8m0x4_to_bfloat2x2,
|
|
|
- packed_dequant_e8m0x4_to_half2x2,
|
|
|
- packed_dequant_nf3x8_to_bfloat2x4,
|
|
|
- nf3_codebook_pools,
|
|
|
- pack_f32x2_to_bfloat2,
|
|
|
- pack_f32x2_to_f16x2,
|
|
|
- red_add_global_bf16x2,
|
|
|
- red_add_global_release_i32,
|
|
|
- red_max_global_f32_nonnegative,
|
|
|
- shared_ptr_to_u32,
|
|
|
- st_global_v4_u32,
|
|
|
- st_global_i32,
|
|
|
- st_global_v4_f32,
|
|
|
- st_shared_bf16_from_f32,
|
|
|
- st_shared_f16_from_f32,
|
|
|
- st_shared_f32,
|
|
|
- st_shared_i32,
|
|
|
- st_shared_u32,
|
|
|
- st_shared_v4_f32,
|
|
|
- st_shared_v4_u32,
|
|
|
- threadfence,
|
|
|
- warp_reduce,
|
|
|
-)
|
|
|
-from b12x.cute.utils import current_cuda_stream, make_ptr
|
|
|
-from b12x.moe.fused.w4a16.route_pack import (
|
|
|
- pack_topk_routes_by_expert as _pack_topk_routes_by_expert,
|
|
|
-)
|
|
|
-from b12x.moe.fused.w4a16.host import (
|
|
|
- _W4A16_ALLOWED_ROUTED_SIZES,
|
|
|
- max_packed_route_slots,
|
|
|
- packed_gemm_scratch_elements,
|
|
|
- plan_w4a16_buffers,
|
|
|
- select_route_block_size_m,
|
|
|
- validate_activation,
|
|
|
-)
|
|
|
-from b12x.cute.runtime_control import raise_if_kernel_resolution_frozen
|
|
|
-from b12x.moe.fused.activations import (
|
|
|
- SWIGLUOAI_UNINTERLEAVE,
|
|
|
- is_gated_moe_activation,
|
|
|
- normalize_moe_activation,
|
|
|
- normalize_swiglu_alpha_for_activation,
|
|
|
- normalize_swiglu_beta_for_activation,
|
|
|
- normalize_swiglu_limit_for_activation,
|
|
|
-)
|
|
|
-from b12x.moe.fused.micro import MoEMicroKernelBackend
|
|
|
-
|
|
|
-
|
|
|
-_ALLOWED_ROUTED_SIZES = _W4A16_ALLOWED_ROUTED_SIZES
|
|
|
-_PACK_FACTOR = 8
|
|
|
-_STAGES = 4
|
|
|
-_E8M0_LOGICAL_TAIL_SCALE_N_ALIGNMENT = 64
|
|
|
-_DEVICE_MAX_REG_BYTES = 255 * 1024
|
|
|
-_DEFAULT_MAX_SHARED_MEM = 101_376
|
|
|
-_WEIGHT_LAYOUTS = {"packed", "modelopt", "nf3_2p1"}
|
|
|
-_MODEL_OPT_W13_LAYOUTS = {"w13", "w31"}
|
|
|
-# NF3 codebook: 8 bf16 codepoints for the 3-bit ("nf3_2p1") weight layout. The
|
|
|
-# device dequant reads these as prmt pools (see nf3_codebook_pools); the host
|
|
|
-# injects the 4 pool words as compile-time constants into the GEMM.
|
|
|
-_NF3_CODEBOOK = (-1.0, -0.6047, -0.3563, -0.1275, 0.1275, 0.3563, 0.6047, 1.0)
|
|
|
-_SCALE_FORMATS = {
|
|
|
- "e4m3_k16": "e4m3_k16",
|
|
|
- "e8m0_k32": "e8m0_k32",
|
|
|
- "e4m3_k32": "e4m3_k32",
|
|
|
-}
|
|
|
-_E8M0_K32_FP16_GLOBAL_COMPENSATION = float(2.0**7)
|
|
|
-_E8M0_K32_BF16_GLOBAL_COMPENSATION = float(2.0**119)
|
|
|
-_MAX_DIRECT_TOPK_ROUTE_M = 6
|
|
|
-_W4A16_SMALL_M_DIRECT_MAX_M = 8
|
|
|
-
|
|
|
-# TC-decode: a small-M decode specialization that runs on the PACKED W4A16
|
|
|
-# object (the same weights/scales the prefill GEMM uses). It reuses the packed
|
|
|
-# tensor-core MMA inner loop but folds the top-k sum into the FC2 store
|
|
|
-# epilogue (dropping the separate top-k-sum launch). It never regresses vs the
|
|
|
-# packed GEMM within its supported M range, so it is ALWAYS used for the small-M
|
|
|
-# direct-topk decode sizes — there is no opt-in/opt-out switch.
|
|
|
-# TC-decode is available for the whole small-M direct-topk range. Its only hard
|
|
|
-# ceiling is the direct-topk route cap (above it, expert route-packing wins).
|
|
|
-# _TC_DECODE_M is retained for callers/tests that enumerate the supported sizes.
|
|
|
-_TC_DECODE_MAX_M = _W4A16_SMALL_M_DIRECT_MAX_M
|
|
|
-_TC_DECODE_M = tuple(range(1, _TC_DECODE_MAX_M + 1))
|
|
|
-
|
|
|
-
|
|
|
-def _m_specialization_key(size_m: int) -> int:
|
|
|
- """M bucket for branches that genuinely specialize on token count."""
|
|
|
- return 1 if int(size_m) == 1 else 0
|
|
|
-
|
|
|
-
|
|
|
-def _fake_m_for_specialization(size_m: int) -> int:
|
|
|
- return 1 if _m_specialization_key(size_m) == 1 else 2
|
|
|
-
|
|
|
-
|
|
|
-# The W4A16 launch model chooses blocks/SM from static resource usage
|
|
|
-# for each specialization. These register counts were measured from the local
|
|
|
-# SM121 JIT output and keep launch occupancy stable across refactors.
|
|
|
-_W4A16_REGS_SM121 = {
|
|
|
- (256, 1, 8, 8, True): 118,
|
|
|
- (256, 1, 16, 4, True): 118,
|
|
|
- (256, 1, 16, 8, True): 118,
|
|
|
- (256, 1, 32, 2, True): 118,
|
|
|
- (128, 1, 4, 8, True): 118,
|
|
|
- (128, 1, 8, 4, True): 120,
|
|
|
- (256, 1, 8, 8, False): 158,
|
|
|
- (128, 1, 4, 8, False): 154,
|
|
|
- (128, 1, 8, 4, False): 143,
|
|
|
- (256, 2, 16, 4, False): 212,
|
|
|
- (128, 2, 4, 8, False): 215,
|
|
|
- (128, 2, 8, 4, False): 214,
|
|
|
- (256, 3, 16, 4, False): 249,
|
|
|
- (128, 3, 4, 8, False): 249,
|
|
|
- (128, 3, 8, 4, False): 250,
|
|
|
- (256, 4, 16, 4, False): 255,
|
|
|
- (128, 4, 4, 8, False): 255,
|
|
|
- (128, 4, 8, 4, False): 255,
|
|
|
-}
|
|
|
-_SMALL_BATCH_TILE_CONFIGS = (
|
|
|
- (128, 128, 256),
|
|
|
- (64, 128, 128),
|
|
|
- (128, 64, 128),
|
|
|
-)
|
|
|
-_LARGE_BATCH_TILE_CONFIGS = (
|
|
|
- (64, 256, 256),
|
|
|
- (64, 128, 128),
|
|
|
- (128, 64, 128),
|
|
|
-)
|
|
|
-
|
|
|
-
|
|
|
-def _covering_count(total: int, quantum: int) -> int:
|
|
|
- return (total + quantum - 1) // quantum
|
|
|
-
|
|
|
-
|
|
|
-def _e8m0_logical_tail_scale_n(size_n: int) -> int:
|
|
|
- return (
|
|
|
- (int(size_n) + _E8M0_LOGICAL_TAIL_SCALE_N_ALIGNMENT - 1)
|
|
|
- // _E8M0_LOGICAL_TAIL_SCALE_N_ALIGNMENT
|
|
|
- ) * _E8M0_LOGICAL_TAIL_SCALE_N_ALIGNMENT
|
|
|
-
|
|
|
-
|
|
|
-def _normalize_swiglu_limit(swiglu_limit: float | None) -> float | None:
|
|
|
- return normalize_swiglu_limit_for_activation("silu", swiglu_limit)
|
|
|
-
|
|
|
-
|
|
|
-def _normalize_activation_swiglu_params(
|
|
|
- activation: str,
|
|
|
- swiglu_limit: float | None,
|
|
|
- swiglu_alpha: float | None,
|
|
|
- swiglu_beta: float | None,
|
|
|
-) -> tuple[float | None, float, float]:
|
|
|
- activation = normalize_moe_activation(activation)
|
|
|
- return (
|
|
|
- normalize_swiglu_limit_for_activation(activation, swiglu_limit),
|
|
|
- normalize_swiglu_alpha_for_activation(activation, swiglu_alpha),
|
|
|
- normalize_swiglu_beta_for_activation(activation, swiglu_beta),
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-def _w4a16_num_regs(
|
|
|
- *,
|
|
|
- cta_threads: int,
|
|
|
- cta_m_blocks: int,
|
|
|
- cta_n_blocks: int,
|
|
|
- cta_k_blocks: int,
|
|
|
- uses_m_block_8: bool,
|
|
|
-) -> int:
|
|
|
- key = (
|
|
|
- int(cta_threads),
|
|
|
- int(cta_m_blocks),
|
|
|
- int(cta_n_blocks),
|
|
|
- int(cta_k_blocks),
|
|
|
- bool(uses_m_block_8),
|
|
|
- )
|
|
|
- try:
|
|
|
- return _W4A16_REGS_SM121[key]
|
|
|
- except KeyError as exc:
|
|
|
- raise ValueError(
|
|
|
- f"missing W4A16 register count for NVFP4 BF16 specialization {key}"
|
|
|
- ) from exc
|
|
|
-
|
|
|
-
|
|
|
-def _shared_memory_footprint(
|
|
|
- *,
|
|
|
- cta_m_blocks: int,
|
|
|
- tile_n: int,
|
|
|
- tile_k: int,
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
-) -> int:
|
|
|
- cta_m = int(cta_m_blocks) * 16
|
|
|
- cta_n = int(tile_n)
|
|
|
- cta_k = int(tile_k)
|
|
|
- sh_block_meta_size = cta_m * 16
|
|
|
- sh_a_size = _STAGES * (cta_m * cta_k) * 2
|
|
|
- sh_b_size = _STAGES * (cta_k * cta_n // _PACK_FACTOR) * 4
|
|
|
- sh_red_size = cta_m * (cta_n + 8) * 2
|
|
|
- sh_bias_size = cta_n * 2
|
|
|
- tmp_size = min(sh_b_size, sh_red_size) + sh_bias_size
|
|
|
- tmp_size = max(max(sh_b_size, sh_red_size), tmp_size)
|
|
|
- sh_s_size = (
|
|
|
- _covering_count(cta_k, _scale_group_size(scale_format)) * cta_n * 2 * _STAGES
|
|
|
- )
|
|
|
- return tmp_size + sh_a_size + sh_s_size + sh_block_meta_size
|
|
|
-
|
|
|
-
|
|
|
-def _determine_blocks_per_sm(
|
|
|
- *,
|
|
|
- problem_m: int,
|
|
|
- problem_n: int,
|
|
|
- top_k: int,
|
|
|
- cta_threads: int,
|
|
|
- cta_m_blocks: int,
|
|
|
- tile_n: int,
|
|
|
- tile_k: int,
|
|
|
- uses_m_block_8: bool,
|
|
|
- sms: int,
|
|
|
- max_shared_mem: int,
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
-) -> int:
|
|
|
- num_regs = _w4a16_num_regs(
|
|
|
- cta_threads=cta_threads,
|
|
|
- cta_m_blocks=cta_m_blocks,
|
|
|
- cta_n_blocks=tile_n // 16,
|
|
|
- cta_k_blocks=tile_k // 16,
|
|
|
- uses_m_block_8=uses_m_block_8,
|
|
|
- )
|
|
|
- register_bytes = max(num_regs, 1) * int(cta_threads) * 4
|
|
|
- smem_bytes = _shared_memory_footprint(
|
|
|
- cta_m_blocks=cta_m_blocks,
|
|
|
- tile_n=tile_n,
|
|
|
- tile_k=tile_k,
|
|
|
- scale_format=scale_format,
|
|
|
- )
|
|
|
- blocks_per_sm_limit = min(
|
|
|
- _DEVICE_MAX_REG_BYTES // register_bytes,
|
|
|
- int(max_shared_mem) // (smem_bytes + 1536),
|
|
|
- )
|
|
|
- if uses_m_block_8:
|
|
|
- # Small-M (moe_block_size==8) TC-decode is weight-bandwidth/overhead
|
|
|
- # bound, not parallelism bound (only m*top_k route-blocks of GEMM work).
|
|
|
- # The fused FC1->activation->FC2 path crosses several grid barriers whose
|
|
|
- # tid==0 atomic-counter increment serializes across all grid_x CTAs, so an
|
|
|
- # oversized grid pays barrier-atomic latency proportional to grid_x for no
|
|
|
- # extra GEMM throughput. Pin one persistent CTA per SM to minimize the
|
|
|
- # barrier participant count while still covering the machine for the
|
|
|
- # I_tp=1024 GEMMs. The split-K persistent loop is grid_x-agnostic, so this
|
|
|
- # is numerically identical.
|
|
|
- blocks_per_sm_limit = 1
|
|
|
- elif cta_m_blocks == 1:
|
|
|
- blocks_per_sm_limit = max(min(blocks_per_sm_limit, 4), 1)
|
|
|
- else:
|
|
|
- blocks_per_sm_limit = max(min(blocks_per_sm_limit, 2), 1)
|
|
|
-
|
|
|
- work_cta_count = (int(problem_n) // int(tile_n)) * int(problem_m) * int(top_k) * 4
|
|
|
- if work_cta_count < int(sms) * blocks_per_sm_limit:
|
|
|
- blocks_per_sm_limit = max(work_cta_count // int(sms), 1)
|
|
|
- return int(blocks_per_sm_limit)
|
|
|
-
|
|
|
-
|
|
|
-def _candidate_tile_fits(
|
|
|
- *,
|
|
|
- problem_n: int,
|
|
|
- problem_k: int,
|
|
|
- cta_m_blocks: int,
|
|
|
- tile_n: int,
|
|
|
- tile_k: int,
|
|
|
- cta_threads: int,
|
|
|
- max_shared_mem: int,
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
- allow_logical_tail: bool = False,
|
|
|
-) -> bool:
|
|
|
- if int(tile_k) == -1 or int(tile_n) == -1 or int(cta_threads) == -1:
|
|
|
- return False
|
|
|
- scale_group_size = _scale_group_size(scale_format)
|
|
|
- exact_n = int(problem_n) % int(tile_n) == 0
|
|
|
- exact_k = int(problem_k) % int(tile_k) == 0
|
|
|
- exact_scale_k = int(problem_k) % scale_group_size == 0
|
|
|
- if not allow_logical_tail:
|
|
|
- if not exact_n or not exact_k:
|
|
|
- return False
|
|
|
- if not exact_scale_k or int(tile_k) % scale_group_size != 0:
|
|
|
- return False
|
|
|
- elif (
|
|
|
- _normalize_scale_format(scale_format) != "e8m0_k32"
|
|
|
- or int(problem_n) % 16 != 0
|
|
|
- or int(problem_k) % 8 != 0
|
|
|
- or int(tile_n) % 16 != 0
|
|
|
- or int(tile_k) % scale_group_size != 0
|
|
|
- ):
|
|
|
- return False
|
|
|
- if int(tile_n) < 64 or int(tile_k) < 64 or int(cta_threads) < 128:
|
|
|
- return False
|
|
|
- smem_bytes = _shared_memory_footprint(
|
|
|
- cta_m_blocks=cta_m_blocks,
|
|
|
- tile_n=tile_n,
|
|
|
- tile_k=tile_k,
|
|
|
- scale_format=scale_format,
|
|
|
- )
|
|
|
- return smem_bytes <= int(max_shared_mem)
|
|
|
-
|
|
|
-
|
|
|
-def _select_tile_config(
|
|
|
- *,
|
|
|
- problem_m: int,
|
|
|
- problem_n: int,
|
|
|
- problem_k: int,
|
|
|
- top_k: int,
|
|
|
- moe_block_size: int,
|
|
|
- sms: int,
|
|
|
- max_shared_mem: int,
|
|
|
- required_cta_threads: int | None = None,
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
- allow_logical_tail: bool = False,
|
|
|
-) -> tuple[int, int, int, int]:
|
|
|
- cta_m_blocks = _covering_count(moe_block_size, 16)
|
|
|
- uses_m_block_8 = moe_block_size == 8
|
|
|
- configs = (
|
|
|
- _LARGE_BATCH_TILE_CONFIGS if cta_m_blocks > 1 else _SMALL_BATCH_TILE_CONFIGS
|
|
|
- )
|
|
|
- best_blocks_per_sm = 0
|
|
|
- best_tile_config: tuple[int, int, int, int] | None = None
|
|
|
- for tile_k, tile_n, cta_threads in configs:
|
|
|
- if required_cta_threads is not None and int(cta_threads) != int(
|
|
|
- required_cta_threads
|
|
|
- ):
|
|
|
- continue
|
|
|
- if not _candidate_tile_fits(
|
|
|
- problem_n=problem_n,
|
|
|
- problem_k=problem_k,
|
|
|
- cta_m_blocks=cta_m_blocks,
|
|
|
- tile_n=tile_n,
|
|
|
- tile_k=tile_k,
|
|
|
- cta_threads=cta_threads,
|
|
|
- max_shared_mem=int(max_shared_mem) - 512,
|
|
|
- scale_format=scale_format,
|
|
|
- allow_logical_tail=allow_logical_tail,
|
|
|
- ):
|
|
|
- continue
|
|
|
- occupancy_problem_n = (
|
|
|
- _covering_count(int(problem_n), int(tile_n)) * int(tile_n)
|
|
|
- if allow_logical_tail
|
|
|
- else int(problem_n)
|
|
|
- )
|
|
|
- blocks_per_sm_limit = _determine_blocks_per_sm(
|
|
|
- problem_m=problem_m,
|
|
|
- problem_n=occupancy_problem_n,
|
|
|
- top_k=top_k,
|
|
|
- cta_threads=cta_threads,
|
|
|
- cta_m_blocks=cta_m_blocks,
|
|
|
- tile_n=tile_n,
|
|
|
- tile_k=tile_k,
|
|
|
- uses_m_block_8=uses_m_block_8,
|
|
|
- sms=sms,
|
|
|
- max_shared_mem=max_shared_mem,
|
|
|
- scale_format=scale_format,
|
|
|
- )
|
|
|
- if blocks_per_sm_limit > best_blocks_per_sm:
|
|
|
- best_blocks_per_sm = blocks_per_sm_limit
|
|
|
- best_tile_config = (tile_k, tile_n, cta_threads, blocks_per_sm_limit)
|
|
|
- if best_tile_config is None:
|
|
|
- cta_thread_msg = (
|
|
|
- ""
|
|
|
- if required_cta_threads is None
|
|
|
- else f", required_cta_threads={required_cta_threads}"
|
|
|
- )
|
|
|
- raise ValueError(
|
|
|
- "no valid W4A16 tile config for "
|
|
|
- f"M/N/K={problem_m}/{problem_n}/{problem_k}, moe_block_size={moe_block_size}"
|
|
|
- f"{cta_thread_msg}"
|
|
|
- )
|
|
|
- return best_tile_config
|
|
|
-
|
|
|
-
|
|
|
-@dataclass(frozen=True)
|
|
|
-class W4A16GemmCompileResult:
|
|
|
- compiled: object
|
|
|
- tile_n: int
|
|
|
- tile_k: int
|
|
|
- moe_block_size: int
|
|
|
- max_m_blocks: int
|
|
|
- blocks_per_sm: int
|
|
|
- weight_layout: str = "packed"
|
|
|
- scale_format: str = "e4m3_k16"
|
|
|
- w13_layout: str = "w13"
|
|
|
-
|
|
|
-
|
|
|
-@dataclass(frozen=True)
|
|
|
-class W4A16ActivationCompileResult:
|
|
|
- compiled: object
|
|
|
- rows: int
|
|
|
- intermediate_size: int
|
|
|
- activation: str
|
|
|
- swiglu_limit: float | None
|
|
|
- swiglu_alpha: float
|
|
|
- swiglu_beta: float
|
|
|
-
|
|
|
-
|
|
|
-@dataclass(frozen=True)
|
|
|
-class W4A16TopKSumCompileResult:
|
|
|
- compiled: object
|
|
|
- m: int
|
|
|
- topk: int
|
|
|
- hidden_size: int
|
|
|
-
|
|
|
-
|
|
|
-@dataclass(frozen=True)
|
|
|
-class W4A16FusedMoeCompileResult:
|
|
|
- compiled: object
|
|
|
- size_m: int
|
|
|
- hidden_size: int
|
|
|
- intermediate_size: int
|
|
|
- num_experts: int
|
|
|
- top_k: int
|
|
|
- activation: str
|
|
|
- apply_router_weight_on_input: bool
|
|
|
- zero_fc2_output: bool
|
|
|
- element_dtype: str
|
|
|
- fast_math: bool
|
|
|
- swiglu_limit: float | None
|
|
|
- swiglu_alpha: float
|
|
|
- swiglu_beta: float
|
|
|
- fc1_tile_n: int
|
|
|
- fc1_tile_k: int
|
|
|
- fc2_tile_n: int
|
|
|
- fc2_tile_k: int
|
|
|
- moe_block_size: int
|
|
|
- max_m_blocks: int
|
|
|
- blocks_per_sm: int
|
|
|
- weight_layout: str = "packed"
|
|
|
- w13_layout: str = "w13"
|
|
|
- direct_topk_routes: bool = False
|
|
|
- scale_format: str = "e4m3_k16"
|
|
|
- tc_decode_fused_sum: bool = False
|
|
|
- collect_activation_amax: bool = False
|
|
|
-
|
|
|
-
|
|
|
-@dataclass(frozen=True)
|
|
|
-class _W4A16GemmLaunch:
|
|
|
- kernel: W4A16GemmCompileResult
|
|
|
- c_tmp: torch.Tensor
|
|
|
-
|
|
|
-
|
|
|
-class _W4A16SmallMDirectLaunch(NamedTuple):
|
|
|
- compiled: object
|
|
|
- grid_x: int
|
|
|
- m: int
|
|
|
- hidden_size: int
|
|
|
- intermediate_size: int
|
|
|
- num_experts: int
|
|
|
- topk: int
|
|
|
- activation: str
|
|
|
- fast_math: bool
|
|
|
- topk_ids_dtype: torch.dtype
|
|
|
-
|
|
|
-
|
|
|
-class _W4A16SmallMDirectKernel(MoEMicroKernelBackend):
|
|
|
- """Decode-sized W4A16 specialization using the native ModelOpt layout."""
|
|
|
-
|
|
|
- _SUPPORTED_M = tuple(range(1, _W4A16_SMALL_M_DIRECT_MAX_M + 1))
|
|
|
-
|
|
|
- @classmethod
|
|
|
- def is_supported(
|
|
|
- cls,
|
|
|
- *,
|
|
|
- m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- topk: int,
|
|
|
- num_experts: int,
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
- ) -> bool:
|
|
|
- # E8M0 reads the shared packed scale grid. Both block axes must be /32.
|
|
|
- if scale_format == "e8m0_k32":
|
|
|
- scale_block_ok = (
|
|
|
- int(hidden_size) % 32 == 0
|
|
|
- and int(intermediate_size) % 32 == 0
|
|
|
- )
|
|
|
- else:
|
|
|
- scale_block_ok = True
|
|
|
- return (
|
|
|
- int(m) in cls._SUPPORTED_M
|
|
|
- and int(m) <= _W4A16_SMALL_M_DIRECT_MAX_M
|
|
|
- and int(hidden_size) > 0
|
|
|
- and int(hidden_size) % 128 == 0
|
|
|
- and int(intermediate_size) > 0
|
|
|
- and int(intermediate_size) % 16 == 0
|
|
|
- and scale_block_ok
|
|
|
- and 0 < int(topk) <= 32
|
|
|
- and int(num_experts) > 0
|
|
|
- and MoEMicroKernelBackend.is_supported(
|
|
|
- int(m),
|
|
|
- int(hidden_size),
|
|
|
- int(intermediate_size),
|
|
|
- int(topk),
|
|
|
- int(num_experts),
|
|
|
- )
|
|
|
- )
|
|
|
-
|
|
|
- def __init__(
|
|
|
- self,
|
|
|
- *,
|
|
|
- activation: str,
|
|
|
- fast_math: bool,
|
|
|
- share_input_across_experts: bool,
|
|
|
- share_expert_scales: bool,
|
|
|
- single_token: bool,
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
- swiglu_limit: float | None = None,
|
|
|
- swiglu_alpha: float | None = None,
|
|
|
- swiglu_beta: float | None = None,
|
|
|
- w13_layout: str = "w13",
|
|
|
- ):
|
|
|
- super().__init__(
|
|
|
- sf_vec_size=16,
|
|
|
- mma_tiler_mn=(64, 128),
|
|
|
- output_tile_count_n=1,
|
|
|
- fast_math=fast_math,
|
|
|
- activation=activation,
|
|
|
- share_input_across_experts=share_input_across_experts,
|
|
|
- share_expert_scales=share_expert_scales,
|
|
|
- single_token=single_token,
|
|
|
- dynamic_down_scale=False,
|
|
|
- w4a16_mode=True,
|
|
|
- scale_format=scale_format,
|
|
|
- swiglu_limit=swiglu_limit,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- w13_layout=w13_layout,
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-class W4A16GemmKernel:
|
|
|
- def __init__(
|
|
|
- self,
|
|
|
- *,
|
|
|
- size_m: int,
|
|
|
- size_n: int,
|
|
|
- size_k: int,
|
|
|
- num_experts: int,
|
|
|
- top_k: int,
|
|
|
- mul_topk_weights: bool,
|
|
|
- tile_n: int,
|
|
|
- tile_k: int,
|
|
|
- moe_block_size: int,
|
|
|
- max_m_blocks: int,
|
|
|
- element_dtype: str = "bf16",
|
|
|
- epilogue_activation: str | None = None,
|
|
|
- weight_layout: str = "packed",
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
- w13_layout: str = "w13",
|
|
|
- source_n_rotation: int = 0,
|
|
|
- single_token_route_fast_path: bool = False,
|
|
|
- direct_topk_routes: bool = False,
|
|
|
- fused_topk_sum: bool = False,
|
|
|
- fused_sum_topk: int = 1,
|
|
|
- ):
|
|
|
- if element_dtype not in {"bf16", "fp16"}:
|
|
|
- raise ValueError(f"unsupported element_dtype {element_dtype!r}")
|
|
|
- if weight_layout not in _WEIGHT_LAYOUTS:
|
|
|
- raise ValueError(f"unsupported W4A16 weight_layout {weight_layout!r}")
|
|
|
- scale_format = _normalize_scale_format(scale_format)
|
|
|
- if weight_layout == "modelopt":
|
|
|
- if w13_layout not in _MODEL_OPT_W13_LAYOUTS:
|
|
|
- raise ValueError(f"unsupported W4A16 w13_layout {w13_layout!r}")
|
|
|
- else:
|
|
|
- w13_layout = "packed"
|
|
|
- source_n_rotation = 0
|
|
|
- if weight_layout == "nf3_2p1":
|
|
|
- if scale_format != "e4m3_k32":
|
|
|
- raise ValueError(
|
|
|
- "nf3_2p1 W4A16 weights require scale_format='e4m3_k32'"
|
|
|
- )
|
|
|
- if element_dtype != "bf16":
|
|
|
- raise ValueError(
|
|
|
- "nf3_2p1 W4A16 weights are bf16-activation only (v1)"
|
|
|
- )
|
|
|
- if epilogue_activation not in (None, "relu2"):
|
|
|
- raise ValueError(
|
|
|
- "W4A16 GEMM epilogue activation currently supports only relu2"
|
|
|
- )
|
|
|
- if tile_n % 16 != 0 or tile_k % 16 != 0:
|
|
|
- raise ValueError("tile_n/tile_k must be multiples of 16")
|
|
|
- scale_group_size = _scale_group_size(scale_format)
|
|
|
- has_n_tile_tail = int(size_n) % int(tile_n) != 0
|
|
|
- has_k_tile_tail = int(size_k) % int(tile_k) != 0
|
|
|
- has_scale_k_tail = int(size_k) % scale_group_size != 0
|
|
|
- has_logical_tail = (
|
|
|
- has_n_tile_tail or has_k_tile_tail or has_scale_k_tail
|
|
|
- )
|
|
|
- if has_logical_tail:
|
|
|
- if weight_layout != "modelopt" or scale_format != "e8m0_k32":
|
|
|
- raise ValueError(
|
|
|
- "W4A16 logical tail tiles are only supported for native "
|
|
|
- "E8M0 K/32 weights"
|
|
|
- )
|
|
|
- if int(size_n) % 16 != 0:
|
|
|
- raise ValueError(
|
|
|
- "native E8M0 W4A16 tail tiles require size_n % 16 == 0"
|
|
|
- )
|
|
|
- if int(size_k) % 8 != 0:
|
|
|
- raise ValueError(
|
|
|
- "native E8M0 W4A16 tail tiles require size_k % 8 == 0"
|
|
|
- )
|
|
|
- if int(tile_k) % scale_group_size != 0:
|
|
|
- raise ValueError(
|
|
|
- "native E8M0 W4A16 tail tiles require tile_k multiples of 32"
|
|
|
- )
|
|
|
- else:
|
|
|
- if size_n % tile_n != 0:
|
|
|
- raise ValueError("size_n must be divisible by tile_n")
|
|
|
- if size_k % tile_k != 0:
|
|
|
- raise ValueError("size_k must be divisible by tile_k")
|
|
|
- if scale_format == "e8m0_k32" and (
|
|
|
- size_k % 32 != 0 or tile_k % 32 != 0
|
|
|
- ):
|
|
|
- raise ValueError(
|
|
|
- "E8M0 K/32 W4A16 scales require size_k/tile_k multiples of 32"
|
|
|
- )
|
|
|
- if scale_format == "e4m3_k32" and (
|
|
|
- size_k % 32 != 0 or tile_k % 32 != 0
|
|
|
- ):
|
|
|
- raise ValueError(
|
|
|
- "E4M3 K/32 W4A16 scales require size_k/tile_k multiples of 32"
|
|
|
- )
|
|
|
- if moe_block_size not in _ALLOWED_ROUTED_SIZES:
|
|
|
- raise ValueError(f"unsupported moe_block_size {moe_block_size}")
|
|
|
- if moe_block_size != 8 and moe_block_size % 16 != 0:
|
|
|
- raise ValueError("moe_block_size must be 8 or a multiple of 16")
|
|
|
- cta_threads = tile_n * tile_k // 64
|
|
|
- if cta_threads not in (128, 256):
|
|
|
- raise ValueError("W4A16 GEMM expects 128 or 256 CTA threads")
|
|
|
- self.size_m = int(size_m)
|
|
|
- self.size_n = int(size_n)
|
|
|
- self.size_k = int(size_k)
|
|
|
- self.num_experts = int(num_experts)
|
|
|
- self.top_k = int(top_k)
|
|
|
- self.mul_topk_weights = bool(mul_topk_weights)
|
|
|
- self.tile_n = int(tile_n)
|
|
|
- self.tile_k = int(tile_k)
|
|
|
- self.cta_n_blocks = int(tile_n // 16)
|
|
|
- self.cta_k_blocks = int(tile_k // 16)
|
|
|
- self.cta_threads = int(cta_threads)
|
|
|
- self.moe_block_size = int(moe_block_size)
|
|
|
- self.element_dtype = element_dtype
|
|
|
- self.is_fp16 = element_dtype == "fp16"
|
|
|
- self.epilogue_relu2 = epilogue_activation == "relu2"
|
|
|
- self.weight_layout = weight_layout
|
|
|
- self.weight_layout_nf3 = weight_layout == "nf3_2p1"
|
|
|
- self.scale_format = scale_format
|
|
|
- self.scale_format_e8m0_k32 = scale_format == "e8m0_k32"
|
|
|
- # k32 scale cadence (two K16 rows share one scale group). Shared by the
|
|
|
- # native E8M0 K/32 path and the NF3 e4m3-style K/32 path; the scale
|
|
|
- # DECODE stays keyed on scale_format_e8m0_k32 (e4m3_k32 uses the e4m3
|
|
|
- # decode arm, its 2**116 compensation lives in the global_scale tensor).
|
|
|
- self.scale_k32 = scale_format in ("e8m0_k32", "e4m3_k32")
|
|
|
- # NF3 codebook prmt pools (host-computed compile-time constants injected
|
|
|
- # into the device dequant). Only consumed by the nf3_2p1 path; cheap to
|
|
|
- # compute unconditionally.
|
|
|
- (
|
|
|
- self._nf3_cb0,
|
|
|
- self._nf3_cb1,
|
|
|
- self._nf3_cb2,
|
|
|
- self._nf3_cb3,
|
|
|
- ) = nf3_codebook_pools(_NF3_CODEBOOK)
|
|
|
- self.scale_group_size = int(scale_group_size)
|
|
|
- self.scale_k_groups = _covering_count(self.size_k, self.scale_group_size)
|
|
|
- self.n_tiles = _covering_count(self.size_n, self.tile_n)
|
|
|
- self.k_tiles = _covering_count(self.size_k, self.tile_k)
|
|
|
- self.covered_size_n = self.n_tiles * self.tile_n
|
|
|
- self.covered_size_k = self.k_tiles * self.tile_k
|
|
|
- self.has_n_tile_tail = bool(has_n_tile_tail)
|
|
|
- self.has_k_tile_tail = bool(has_k_tile_tail)
|
|
|
- self.has_scale_k_tail = bool(has_scale_k_tail)
|
|
|
- self.has_logical_tail = bool(has_logical_tail)
|
|
|
- self.scale_size_n = (
|
|
|
- _e8m0_logical_tail_scale_n(self.size_n)
|
|
|
- if self.scale_format_e8m0_k32 and self.has_n_tile_tail
|
|
|
- else self.size_n
|
|
|
- )
|
|
|
- self.scale_n_groups = self.scale_size_n // 16
|
|
|
- self.w13_layout = w13_layout
|
|
|
- self.source_n_rotation = int(source_n_rotation)
|
|
|
- self.single_token_route_fast_path = bool(single_token_route_fast_path)
|
|
|
- self.direct_topk_routes = bool(direct_topk_routes)
|
|
|
- self.fused_topk_sum = bool(fused_topk_sum)
|
|
|
- self.fused_sum_topk = int(fused_sum_topk)
|
|
|
- if self.fused_topk_sum and not self.direct_topk_routes:
|
|
|
- raise ValueError("fused_topk_sum requires direct_topk_routes")
|
|
|
- if self.fused_topk_sum and self.fused_sum_topk < 1:
|
|
|
- raise ValueError("fused_sum_topk must be >= 1")
|
|
|
- self.cta_m_blocks = int(_covering_count(moe_block_size, 16))
|
|
|
- self.uses_m_block_8 = moe_block_size == 8
|
|
|
- self.max_m_blocks = int(max_m_blocks)
|
|
|
- if torch.cuda.is_available():
|
|
|
- props = torch.cuda.get_device_properties(torch.cuda.current_device())
|
|
|
- self.sms = int(props.multi_processor_count)
|
|
|
- max_shared_mem = int(
|
|
|
- getattr(props, "shared_memory_per_block_optin", _DEFAULT_MAX_SHARED_MEM)
|
|
|
- )
|
|
|
- else:
|
|
|
- self.sms = 120
|
|
|
- max_shared_mem = _DEFAULT_MAX_SHARED_MEM
|
|
|
- self.blocks_per_sm = _determine_blocks_per_sm(
|
|
|
- problem_m=self.size_m,
|
|
|
- problem_n=self.covered_size_n,
|
|
|
- top_k=self.top_k,
|
|
|
- cta_threads=self.cta_threads,
|
|
|
- cta_m_blocks=self.cta_m_blocks,
|
|
|
- tile_n=self.tile_n,
|
|
|
- tile_k=self.tile_k,
|
|
|
- uses_m_block_8=self.uses_m_block_8,
|
|
|
- sms=self.sms,
|
|
|
- max_shared_mem=max_shared_mem,
|
|
|
- scale_format=scale_format,
|
|
|
- )
|
|
|
-
|
|
|
- # W4A16 shared-memory geometry, in int4 units unless noted.
|
|
|
- self.a_sh_stride = 16 * self.cta_k_blocks // 8
|
|
|
- self.a_sh_stage = self.a_sh_stride * (16 * self.cta_m_blocks)
|
|
|
- self.a_gl_rd_delta_o = 16 * self.cta_k_blocks // 8
|
|
|
- self.a_sh_wr_delta = self.a_sh_stride * (
|
|
|
- self.cta_threads // self.a_gl_rd_delta_o
|
|
|
- )
|
|
|
- self.a_sh_wr_iters = _covering_count(self.a_sh_stage, self.a_sh_wr_delta)
|
|
|
- self.a_sh_rd_delta_i = self.a_sh_stride * 16
|
|
|
-
|
|
|
- self.b_sh_stride = ((self.cta_n_blocks * 16) * 16 // _PACK_FACTOR) // 4
|
|
|
- self.b_thread_vecs = 1
|
|
|
- self.b_sh_stride_threads = self.b_sh_stride
|
|
|
- self.b_sh_stage = self.b_sh_stride * self.cta_k_blocks
|
|
|
- self.b_sh_wr_iters = self.b_sh_stage // self.cta_threads
|
|
|
- # NF3 ("nf3_2p1") stages 12-byte triples per 32-code unit instead of the
|
|
|
- # 16-byte int4 unit; the per-thread unit count is unchanged, only the
|
|
|
- # bytes-per-unit and the 16-byte cp.async chunk count differ.
|
|
|
- self.b_unit_bytes = 12 if self.weight_layout_nf3 else 16
|
|
|
- self.b_sh_stage_bytes = self.b_sh_stage * self.b_unit_bytes
|
|
|
- if self.weight_layout_nf3:
|
|
|
- if self.b_sh_stage_bytes % 16 != 0:
|
|
|
- raise ValueError(
|
|
|
- "nf3_2p1 B stage bytes must be a multiple of 16; "
|
|
|
- f"got {self.b_sh_stage_bytes}"
|
|
|
- )
|
|
|
- self.b_sh_chunks = self.b_sh_stage_bytes // 16
|
|
|
- self.b_sh_wr_iters_nf3 = _covering_count(
|
|
|
- self.b_sh_chunks, self.cta_threads
|
|
|
- )
|
|
|
- else:
|
|
|
- self.b_sh_chunks = self.b_sh_stage
|
|
|
- self.b_sh_wr_iters_nf3 = self.b_sh_wr_iters
|
|
|
-
|
|
|
- self.s_sh_stride = 16 * self.cta_n_blocks // 16
|
|
|
- self.s_tb_groups = (
|
|
|
- self.cta_k_blocks // 2 if self.scale_k32 else self.cta_k_blocks
|
|
|
- )
|
|
|
- self.s_sh_stage = self.s_tb_groups * self.s_sh_stride
|
|
|
- self.tb_n_warps = self.cta_n_blocks // 4
|
|
|
-
|
|
|
- sh_block_route_indices = self.moe_block_size // 4
|
|
|
- sh_rd_block_route_indices = self.moe_block_size // 4
|
|
|
- sh_block_topk_weights = self.moe_block_size // 2
|
|
|
- self.sh_valid_count_off = (
|
|
|
- sh_block_route_indices + sh_rd_block_route_indices + sh_block_topk_weights
|
|
|
- )
|
|
|
- self.sh_route_off = 0
|
|
|
- self.sh_rd_route_off = sh_block_route_indices
|
|
|
- self.sh_topk_off = sh_block_route_indices + sh_rd_block_route_indices
|
|
|
-
|
|
|
- sh_red_size = (2 * self.cta_n_blocks + 1) * 16 * self.cta_m_blocks
|
|
|
- # B region size in int4 (16-byte) units. NF3 units are 12 bytes, so the
|
|
|
- # region is 0.75x; round the total up to a 16-byte multiple so the
|
|
|
- # following SMEM regions keep their 16-byte alignment (exact for all
|
|
|
- # supported tiles since b_sh_stage is a multiple of 4).
|
|
|
- sh_b_size = _covering_count(_STAGES * self.b_sh_stage_bytes, 16)
|
|
|
- sh_size_min = min(sh_red_size, sh_b_size)
|
|
|
- sh_size_max = max(sh_red_size, sh_b_size)
|
|
|
- sh_bias_size = self.cta_n_blocks * 16 // 8
|
|
|
- sh_b_red_bias_size = max(sh_size_max, sh_size_min + sh_bias_size)
|
|
|
- self.sh_b_off = self.sh_valid_count_off
|
|
|
- self.sh_red_off = self.sh_valid_count_off
|
|
|
- self.sh_s_off = self.sh_valid_count_off + sh_b_red_bias_size
|
|
|
- self.sh_a_off = self.sh_s_off + _STAGES * self.s_sh_stage
|
|
|
- self.shared_int4 = self.sh_a_off + _STAGES * self.a_sh_stage
|
|
|
- self.shared_words = self.shared_int4 * 4
|
|
|
-
|
|
|
- @property
|
|
|
- def __cache_key__(self) -> tuple[object, ...]:
|
|
|
- return (
|
|
|
- self.size_n,
|
|
|
- self.size_k,
|
|
|
- self.covered_size_n,
|
|
|
- self.covered_size_k,
|
|
|
- self.scale_k_groups,
|
|
|
- self.has_logical_tail,
|
|
|
- self.num_experts,
|
|
|
- self.top_k,
|
|
|
- self.mul_topk_weights,
|
|
|
- self.tile_n,
|
|
|
- self.tile_k,
|
|
|
- self.cta_threads,
|
|
|
- self.moe_block_size,
|
|
|
- self.element_dtype,
|
|
|
- self.epilogue_relu2,
|
|
|
- self.weight_layout,
|
|
|
- self.scale_format,
|
|
|
- self.w13_layout,
|
|
|
- self.source_n_rotation,
|
|
|
- self.single_token_route_fast_path,
|
|
|
- self.direct_topk_routes,
|
|
|
- self.fused_topk_sum,
|
|
|
- self.fused_sum_topk,
|
|
|
- self.cta_m_blocks,
|
|
|
- self.uses_m_block_8,
|
|
|
- self.shared_words,
|
|
|
- # Launch bounds are part of the compiled kernel. Keep binaries
|
|
|
- # planned for different residency targets out of the same cache
|
|
|
- # entry even when their arithmetic geometry otherwise matches.
|
|
|
- self.blocks_per_sm,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _activation_smem_permuted_offset(self, i: Int32) -> Int32:
|
|
|
- row = i // Int32(self.a_gl_rd_delta_o)
|
|
|
- return Int32(self.a_gl_rd_delta_o) * row + (
|
|
|
- (i - row * Int32(self.a_gl_rd_delta_o)) ^ (row & Int32(7))
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _int4_addr(self, smem_base: Int32, int4_off: Int32) -> Int32:
|
|
|
- return smem_base + int4_off * Int32(16)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _dequant_e2m1x4_to_elem2x2(self, packed: Uint32):
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return packed_dequant_e2m1x4_to_half2x2(packed)
|
|
|
- return packed_dequant_e2m1x4_to_bfloat2x2(packed)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _dequant_e4m3x4_to_elem2x2(self, packed: Uint32):
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return packed_dequant_e4m3x4_to_half2x2(packed)
|
|
|
- return packed_dequant_e4m3x4_to_bfloat2x2(packed)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _dequant_scale_x4_to_elem2x2(self, packed: Uint32):
|
|
|
- if cutlass.const_expr(self.scale_format_e8m0_k32):
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return packed_dequant_e8m0x4_to_half2x2(packed)
|
|
|
- return packed_dequant_e8m0x4_to_bfloat2x2(packed)
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return packed_dequant_e4m3x4_to_half2x2(packed)
|
|
|
- return packed_dequant_e4m3x4_to_bfloat2x2(packed)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _elem2_mul(self, a: Uint32, b: Uint32) -> Uint32:
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return half2_mul(a, b)
|
|
|
- return bfloat2_mul(a, b)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _broadcast_f32_to_elem2(self, x: cutlass.Float32) -> Uint32:
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return broadcast_f32_to_half2(x)
|
|
|
- return broadcast_f32_to_bfloat2(x)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _pack_f32x2_to_elem2(self, x0: cutlass.Float32, x1: cutlass.Float32) -> Uint32:
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return pack_f32x2_to_f16x2(x0, x1)
|
|
|
- return pack_f32x2_to_bfloat2(x0, x1)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _elem2_to_f32x2(self, packed: Uint32):
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return half2_to_float2_scaled(packed, cutlass.Float32(1.0))
|
|
|
- return bfloat2_to_float2_scaled(packed, cutlass.Float32(1.0))
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _relu2_elem2(self, packed: Uint32) -> Uint32:
|
|
|
- x0, x1 = self._elem2_to_f32x2(packed)
|
|
|
- if x0 < cutlass.Float32(0.0):
|
|
|
- x0 = cutlass.Float32(0.0)
|
|
|
- if x1 < cutlass.Float32(0.0):
|
|
|
- x1 = cutlass.Float32(0.0)
|
|
|
- return self._pack_f32x2_to_elem2(x0 * x0, x1 * x1)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _st_shared_elem_from_f32(self, addr: Int32, val: cutlass.Float32):
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- st_shared_f16_from_f32(addr, val)
|
|
|
- else:
|
|
|
- st_shared_bf16_from_f32(addr, val)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _mma_m16n8k16_f32(
|
|
|
- self,
|
|
|
- d0: cutlass.Float32,
|
|
|
- d1: cutlass.Float32,
|
|
|
- d2: cutlass.Float32,
|
|
|
- d3: cutlass.Float32,
|
|
|
- a0: Uint32,
|
|
|
- a1: Uint32,
|
|
|
- a2: Uint32,
|
|
|
- a3: Uint32,
|
|
|
- b0: Uint32,
|
|
|
- b1: Uint32,
|
|
|
- ):
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return f16_mma_m16n8k16_f32(d0, d1, d2, d3, a0, a1, a2, a3, b0, b1)
|
|
|
- return bf16_mma_m16n8k16_f32(d0, d1, d2, d3, a0, a1, a2, a3, b0, b1)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _mma_rhs_fragments_as_mma_a_m16n8k16_f32(
|
|
|
- self,
|
|
|
- d0: cutlass.Float32,
|
|
|
- d1: cutlass.Float32,
|
|
|
- d2: cutlass.Float32,
|
|
|
- d3: cutlass.Float32,
|
|
|
- b0_0: Uint32,
|
|
|
- b1_0: Uint32,
|
|
|
- b0_1: Uint32,
|
|
|
- b1_1: Uint32,
|
|
|
- a0: Uint32,
|
|
|
- a1: Uint32,
|
|
|
- ):
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return f16_mma_rhs_fragments_as_mma_a_m16n8k16_f32(
|
|
|
- d0, d1, d2, d3, b0_0, b1_0, b0_1, b1_1, a0, a1
|
|
|
- )
|
|
|
- return bf16_mma_rhs_fragments_as_mma_a_m16n8k16_f32(
|
|
|
- d0, d1, d2, d3, b0_0, b1_0, b0_1, b1_1, a0, a1
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def __call__(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- b_i32_flat: cute.Tensor,
|
|
|
- c_bf16_flat: cute.Tensor,
|
|
|
- scales_i32_flat: cute.Tensor,
|
|
|
- global_scale: cute.Tensor,
|
|
|
- packed_route_indices: cute.Tensor,
|
|
|
- block_expert_ids: cute.Tensor,
|
|
|
- packed_route_count: cute.Tensor,
|
|
|
- topk_weights_flat: cute.Tensor,
|
|
|
- c_tmp_f32_flat: cute.Tensor,
|
|
|
- locks_i32_flat: cute.Tensor,
|
|
|
- active_m: cutlass.Int32,
|
|
|
- grid_x: cutlass.Int32,
|
|
|
- stream: cuda.CUstream,
|
|
|
- ):
|
|
|
- grid = (grid_x, 1, 1)
|
|
|
- self.kernel(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- c_bf16_flat,
|
|
|
- scales_i32_flat,
|
|
|
- global_scale,
|
|
|
- packed_route_indices,
|
|
|
- block_expert_ids,
|
|
|
- packed_route_count,
|
|
|
- topk_weights_flat,
|
|
|
- c_tmp_f32_flat,
|
|
|
- locks_i32_flat,
|
|
|
- active_m,
|
|
|
- ).launch(
|
|
|
- grid=grid,
|
|
|
- block=[self.cta_threads, 1, 1],
|
|
|
- min_blocks_per_mp=self.blocks_per_sm,
|
|
|
- stream=stream,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.kernel
|
|
|
- def kernel(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- b_i32_flat: cute.Tensor,
|
|
|
- c_bf16_flat: cute.Tensor,
|
|
|
- scales_i32_flat: cute.Tensor,
|
|
|
- global_scale: cute.Tensor,
|
|
|
- packed_route_indices: cute.Tensor,
|
|
|
- block_expert_ids: cute.Tensor,
|
|
|
- packed_route_count: cute.Tensor,
|
|
|
- topk_weights_flat: cute.Tensor,
|
|
|
- c_tmp_f32_flat: cute.Tensor,
|
|
|
- locks_i32_flat: cute.Tensor,
|
|
|
- active_m: cutlass.Int32,
|
|
|
- ):
|
|
|
- tidx, _, _ = cute.arch.thread_idx()
|
|
|
- bidx, _, _ = cute.arch.block_idx()
|
|
|
- tid = Int32(tidx)
|
|
|
- cta = Int32(bidx)
|
|
|
-
|
|
|
- smem = cutlass.utils.SmemAllocator()
|
|
|
-
|
|
|
- @cute.struct
|
|
|
- class Storage:
|
|
|
- words: cute.struct.Align[
|
|
|
- cute.struct.MemRange[cutlass.Uint32, self.shared_words],
|
|
|
- 1024,
|
|
|
- ]
|
|
|
-
|
|
|
- storage = smem.allocate(Storage)
|
|
|
- smem_base = shared_ptr_to_u32(storage.words.data_ptr())
|
|
|
-
|
|
|
- grid_x, _, _ = cute.arch.grid_dim()
|
|
|
- self._run_persistent_gemm(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- c_bf16_flat,
|
|
|
- scales_i32_flat,
|
|
|
- global_scale,
|
|
|
- packed_route_indices,
|
|
|
- block_expert_ids,
|
|
|
- packed_route_count,
|
|
|
- topk_weights_flat,
|
|
|
- c_tmp_f32_flat,
|
|
|
- locks_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- cta,
|
|
|
- Int32(grid_x),
|
|
|
- Int32(active_m),
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _run_persistent_gemm(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- b_i32_flat: cute.Tensor,
|
|
|
- c_bf16_flat: cute.Tensor,
|
|
|
- scales_i32_flat: cute.Tensor,
|
|
|
- global_scale: cute.Tensor,
|
|
|
- packed_route_indices: cute.Tensor,
|
|
|
- block_expert_ids: cute.Tensor,
|
|
|
- packed_route_count: cute.Tensor,
|
|
|
- topk_weights_flat: cute.Tensor,
|
|
|
- c_tmp_f32_flat: cute.Tensor,
|
|
|
- locks_i32_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- cta: Int32,
|
|
|
- grid_x: Int32,
|
|
|
- active_size_m: Int32,
|
|
|
- ):
|
|
|
- n_tiles = Int32(self.n_tiles)
|
|
|
- route_blocks = active_size_m * Int32(self.top_k)
|
|
|
- if cutlass.const_expr(not self.direct_topk_routes):
|
|
|
- route_blocks = packed_route_count[Int32(0)].to(Int32) // Int32(
|
|
|
- self.moe_block_size
|
|
|
- )
|
|
|
- k_tiles = Int32(self.k_tiles)
|
|
|
- global_mn_tiles = route_blocks * n_tiles
|
|
|
-
|
|
|
- tail_mn_tiles = global_mn_tiles
|
|
|
- full_grid_mn_iters = Int32(0)
|
|
|
- force_one_tile_per_cta = Int32(0)
|
|
|
- if cutlass.const_expr(self.uses_m_block_8):
|
|
|
- # TC-decode small-M: when every mn-tile fits inside the launched grid
|
|
|
- # (FC1 has only route_blocks*n_tiles tiles, far fewer than grid_x),
|
|
|
- # the default tail path fans each mn-tile across multiple CTAs along
|
|
|
- # K and pays a lock-serialized cross-CTA split-K finalize plus the
|
|
|
- # reduction-turn handshake. Instead give the first global_mn_tiles
|
|
|
- # CTAs exactly one full mn-tile (all k_tiles, reduce_slice_count==1,
|
|
|
- # no finalize, no lock traffic) and idle the rest. grid_x and the
|
|
|
- # grid-barrier participant count are unchanged (so FC2 coverage is
|
|
|
- # untouched); only FC1's intra-GEMM work partition changes. Numerically
|
|
|
- # identical: a single CTA computes the whole K-reduction per tile.
|
|
|
- if global_mn_tiles <= grid_x:
|
|
|
- force_one_tile_per_cta = Int32(1)
|
|
|
- if force_one_tile_per_cta != Int32(0):
|
|
|
- tail_mn_tiles = Int32(0)
|
|
|
- full_grid_mn_iters = Int32(1)
|
|
|
- elif global_mn_tiles > grid_x:
|
|
|
- tail_mn_tiles = global_mn_tiles - (global_mn_tiles // grid_x) * grid_x
|
|
|
- if tail_mn_tiles * Int32(3) <= grid_x:
|
|
|
- tail_mn_tiles += grid_x
|
|
|
- full_grid_mn_iters = (global_mn_tiles - tail_mn_tiles) // grid_x
|
|
|
-
|
|
|
- iters = (k_tiles * tail_mn_tiles + grid_x - Int32(1)) // grid_x
|
|
|
-
|
|
|
- lock_slot = Int32(0)
|
|
|
- if tail_mn_tiles >= grid_x:
|
|
|
- lock_slot = cta
|
|
|
- else:
|
|
|
- lock_slot = (iters * cta) // k_tiles - Int32(1)
|
|
|
-
|
|
|
- in_tail_region = Int32(0)
|
|
|
- has_work = Int32(1)
|
|
|
- work_mn_tile = cta
|
|
|
- reduce_k_tile = Int32(0)
|
|
|
- route_block_idx = Int32(0)
|
|
|
- output_n_tile = Int32(0)
|
|
|
- if iters == Int32(0) and full_grid_mn_iters == Int32(0):
|
|
|
- has_work = Int32(0)
|
|
|
-
|
|
|
- while has_work != Int32(0):
|
|
|
- reduce_tile_count = Int32(0)
|
|
|
- reduce_slice_count = Int32(1)
|
|
|
- reduce_slice_idx = Int32(0)
|
|
|
-
|
|
|
- if in_tail_region == Int32(0) and full_grid_mn_iters > Int32(0):
|
|
|
- route_block_idx = work_mn_tile // n_tiles
|
|
|
- output_n_tile = work_mn_tile - route_block_idx * n_tiles
|
|
|
- reduce_k_tile = Int32(0)
|
|
|
- reduce_tile_count = k_tiles
|
|
|
- full_grid_mn_iters -= Int32(1)
|
|
|
- else:
|
|
|
- if in_tail_region == Int32(0):
|
|
|
- in_tail_region = Int32(1)
|
|
|
- tail_mn_base = global_mn_tiles - tail_mn_tiles
|
|
|
- cta_iter_start = iters * cta
|
|
|
- work_mn_tile = cta_iter_start // k_tiles
|
|
|
- reduce_k_tile = cta_iter_start - work_mn_tile * k_tiles
|
|
|
- global_mn_tile = work_mn_tile + tail_mn_base
|
|
|
- route_block_idx = global_mn_tile // n_tiles
|
|
|
- output_n_tile = global_mn_tile - route_block_idx * n_tiles
|
|
|
-
|
|
|
- if work_mn_tile < tail_mn_tiles and iters > Int32(0):
|
|
|
- reduce_tile_count = iters * (cta + Int32(1)) - (
|
|
|
- k_tiles * work_mn_tile + reduce_k_tile
|
|
|
- )
|
|
|
- if reduce_tile_count < Int32(0):
|
|
|
- reduce_tile_count = Int32(0)
|
|
|
- if reduce_k_tile + reduce_tile_count > k_tiles:
|
|
|
- reduce_tile_count = k_tiles - reduce_k_tile
|
|
|
-
|
|
|
- if reduce_tile_count > Int32(0):
|
|
|
- first_reduce_boundary = iters * (
|
|
|
- (k_tiles * work_mn_tile + iters - Int32(1)) // iters
|
|
|
- )
|
|
|
- if first_reduce_boundary <= k_tiles * (work_mn_tile + Int32(1)):
|
|
|
- reduce_boundary_offset = (
|
|
|
- first_reduce_boundary - k_tiles * work_mn_tile
|
|
|
- )
|
|
|
- reduce_slice_count = (
|
|
|
- k_tiles - reduce_boundary_offset + iters - Int32(1)
|
|
|
- ) // iters
|
|
|
- if reduce_boundary_offset > Int32(0):
|
|
|
- reduce_slice_count += Int32(1)
|
|
|
- reduce_boundary_delta = iters * cta - first_reduce_boundary
|
|
|
- if reduce_boundary_delta < Int32(0):
|
|
|
- reduce_slice_idx = reduce_slice_count - Int32(1)
|
|
|
- else:
|
|
|
- if reduce_boundary_offset == Int32(
|
|
|
- 0
|
|
|
- ) and reduce_boundary_delta == Int32(0):
|
|
|
- reduce_slice_idx = reduce_slice_count - Int32(1)
|
|
|
- else:
|
|
|
- reduce_slice_idx = (
|
|
|
- reduce_slice_count
|
|
|
- - Int32(1)
|
|
|
- - reduce_boundary_delta // iters
|
|
|
- )
|
|
|
- if reduce_boundary_offset > Int32(0):
|
|
|
- reduce_slice_idx -= Int32(1)
|
|
|
-
|
|
|
- if tail_mn_tiles >= grid_x:
|
|
|
- if reduce_slice_count > Int32(
|
|
|
- 1
|
|
|
- ) and reduce_slice_idx == reduce_slice_count - Int32(1):
|
|
|
- lock_slot += Int32(1)
|
|
|
- else:
|
|
|
- lock_slot += Int32(1)
|
|
|
- else:
|
|
|
- has_work = Int32(0)
|
|
|
- else:
|
|
|
- has_work = Int32(0)
|
|
|
-
|
|
|
- if (
|
|
|
- has_work != Int32(0)
|
|
|
- and reduce_tile_count > Int32(0)
|
|
|
- and route_block_idx < route_blocks
|
|
|
- ):
|
|
|
- if cutlass.const_expr(self.direct_topk_routes):
|
|
|
- expert_idx = packed_route_indices[route_block_idx].to(Int32)
|
|
|
- else:
|
|
|
- expert_idx = block_expert_ids[route_block_idx].to(Int32)
|
|
|
- if expert_idx >= Int32(0):
|
|
|
- self._run_tile(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- c_bf16_flat,
|
|
|
- scales_i32_flat,
|
|
|
- global_scale,
|
|
|
- packed_route_indices,
|
|
|
- topk_weights_flat,
|
|
|
- c_tmp_f32_flat,
|
|
|
- locks_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- route_block_idx,
|
|
|
- expert_idx,
|
|
|
- output_n_tile,
|
|
|
- reduce_k_tile,
|
|
|
- reduce_tile_count,
|
|
|
- reduce_slice_count,
|
|
|
- reduce_slice_idx,
|
|
|
- lock_slot,
|
|
|
- active_size_m,
|
|
|
- )
|
|
|
-
|
|
|
- if has_work != Int32(0):
|
|
|
- if in_tail_region == Int32(0):
|
|
|
- work_mn_tile += grid_x
|
|
|
- else:
|
|
|
- reduce_k_tile = Int32(0)
|
|
|
- work_mn_tile += Int32(1)
|
|
|
- output_n_tile += Int32(1)
|
|
|
- if output_n_tile == n_tiles:
|
|
|
- output_n_tile = Int32(0)
|
|
|
- route_block_idx += Int32(1)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _read_moe_block_data(
|
|
|
- self,
|
|
|
- packed_route_indices: cute.Tensor,
|
|
|
- topk_weights_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- route_block_idx: Int32,
|
|
|
- global_scale_f32: cutlass.Float32,
|
|
|
- active_size_m: Int32,
|
|
|
- ) -> Int32:
|
|
|
- if cutlass.const_expr(self.direct_topk_routes):
|
|
|
- if tid == Int32(0):
|
|
|
- idx = route_block_idx
|
|
|
- st_shared_i32(smem_base + Int32(self.sh_route_off * 16), idx)
|
|
|
- rd_row = idx // Int32(self.top_k)
|
|
|
- st_shared_i32(smem_base + Int32(self.sh_rd_route_off * 16), rd_row)
|
|
|
- if cutlass.const_expr(self.mul_topk_weights):
|
|
|
- topk = topk_weights_flat[idx].to(cutlass.Float32) * global_scale_f32
|
|
|
- st_shared_u32(
|
|
|
- smem_base + Int32(self.sh_topk_off * 16),
|
|
|
- self._broadcast_f32_to_elem2(topk),
|
|
|
- )
|
|
|
- cute.arch.sync_threads()
|
|
|
- return Int32(1)
|
|
|
-
|
|
|
- if cutlass.const_expr(self.single_token_route_fast_path):
|
|
|
- if tid == Int32(0):
|
|
|
- idx = packed_route_indices[
|
|
|
- route_block_idx * Int32(self.moe_block_size)
|
|
|
- ].to(Int32)
|
|
|
- st_shared_i32(smem_base + Int32(self.sh_route_off * 16), idx)
|
|
|
- rd_row = idx // Int32(self.top_k)
|
|
|
- st_shared_i32(smem_base + Int32(self.sh_rd_route_off * 16), rd_row)
|
|
|
- if cutlass.const_expr(self.mul_topk_weights):
|
|
|
- safe_idx = idx
|
|
|
- if idx >= active_size_m * Int32(self.top_k):
|
|
|
- safe_idx = Int32(0)
|
|
|
- topk = (
|
|
|
- topk_weights_flat[safe_idx].to(cutlass.Float32)
|
|
|
- * global_scale_f32
|
|
|
- )
|
|
|
- st_shared_u32(
|
|
|
- smem_base + Int32(self.sh_topk_off * 16),
|
|
|
- self._broadcast_f32_to_elem2(topk),
|
|
|
- )
|
|
|
- cute.arch.sync_threads()
|
|
|
- return Int32(1)
|
|
|
-
|
|
|
- route_indices_int4_addr = self._int4_addr(
|
|
|
- smem_base, Int32(self.sh_route_off) + tid
|
|
|
- )
|
|
|
- route_indices_gmem = get_ptr_as_int64(
|
|
|
- packed_route_indices,
|
|
|
- route_block_idx * Int32(self.moe_block_size) + tid * Int32(4),
|
|
|
- )
|
|
|
- cp_async4_shared_global_pred(
|
|
|
- route_indices_int4_addr,
|
|
|
- route_indices_gmem,
|
|
|
- (tid < Int32(self.moe_block_size // 4)).to(Int32),
|
|
|
- )
|
|
|
- cute.arch.cp_async_commit_group()
|
|
|
- cute.arch.cp_async_wait_group(0)
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- if tid >= Int32(self.cta_threads - 32):
|
|
|
- size_per_thread = _covering_count(self.moe_block_size, 32)
|
|
|
- lane = tid - Int32(self.cta_threads - 32)
|
|
|
- local_count = Int32(0)
|
|
|
- for i in cutlass.range_constexpr(size_per_thread):
|
|
|
- j = lane * Int32(size_per_thread) + Int32(i)
|
|
|
- if j < Int32(self.moe_block_size):
|
|
|
- idx = ld_shared_i32_relaxed(
|
|
|
- smem_base + Int32(self.sh_route_off * 16) + j * Int32(4)
|
|
|
- )
|
|
|
- if idx < active_size_m * Int32(self.top_k):
|
|
|
- local_count += Int32(1)
|
|
|
- valid = cute.arch.warp_redux_sync(local_count, "add")
|
|
|
- if lane == Int32(0):
|
|
|
- st_shared_i32(smem_base + Int32(self.sh_valid_count_off * 16), valid)
|
|
|
-
|
|
|
- if tid < Int32(self.moe_block_size):
|
|
|
- idx = ld_shared_i32_relaxed(
|
|
|
- smem_base + Int32(self.sh_route_off * 16) + tid * Int32(4)
|
|
|
- )
|
|
|
- rd_row = idx // Int32(self.top_k)
|
|
|
- st_shared_i32(
|
|
|
- smem_base + Int32(self.sh_rd_route_off * 16) + tid * Int32(4),
|
|
|
- rd_row,
|
|
|
- )
|
|
|
- if cutlass.const_expr(self.mul_topk_weights):
|
|
|
- safe_idx = idx
|
|
|
- if idx >= active_size_m * Int32(self.top_k):
|
|
|
- safe_idx = Int32(0)
|
|
|
- topk = (
|
|
|
- topk_weights_flat[safe_idx].to(cutlass.Float32) * global_scale_f32
|
|
|
- )
|
|
|
- packed_topk = self._broadcast_f32_to_elem2(topk)
|
|
|
- # top-k weights are cached as packed element pairs.
|
|
|
- topk_word_addr = (
|
|
|
- smem_base + Int32(self.sh_topk_off * 16) + tid * Int32(4)
|
|
|
- )
|
|
|
- st_shared_u32(topk_word_addr, packed_topk)
|
|
|
-
|
|
|
- cute.arch.sync_threads()
|
|
|
- valid_count = ld_shared_i32_relaxed(
|
|
|
- smem_base + Int32(self.sh_valid_count_off * 16)
|
|
|
- )
|
|
|
- cute.arch.sync_threads()
|
|
|
- return valid_count
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _run_tile(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- b_i32_flat: cute.Tensor,
|
|
|
- c_bf16_flat: cute.Tensor,
|
|
|
- scales_i32_flat: cute.Tensor,
|
|
|
- global_scale: cute.Tensor,
|
|
|
- packed_route_indices: cute.Tensor,
|
|
|
- topk_weights_flat: cute.Tensor,
|
|
|
- c_tmp_f32_flat: cute.Tensor,
|
|
|
- locks_i32_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- route_block_idx: Int32,
|
|
|
- expert_idx: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- reduce_k_tile: Int32,
|
|
|
- reduce_tile_count: Int32,
|
|
|
- reduce_slice_count: Int32,
|
|
|
- reduce_slice_idx: Int32,
|
|
|
- lock_slot: Int32,
|
|
|
- active_size_m: Int32,
|
|
|
- ):
|
|
|
- if cutlass.const_expr(self.uses_m_block_8):
|
|
|
- self._run_tile_m8(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- c_bf16_flat,
|
|
|
- scales_i32_flat,
|
|
|
- global_scale,
|
|
|
- packed_route_indices,
|
|
|
- topk_weights_flat,
|
|
|
- c_tmp_f32_flat,
|
|
|
- locks_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- route_block_idx,
|
|
|
- expert_idx,
|
|
|
- output_n_tile,
|
|
|
- reduce_k_tile,
|
|
|
- reduce_tile_count,
|
|
|
- reduce_slice_count,
|
|
|
- reduce_slice_idx,
|
|
|
- lock_slot,
|
|
|
- active_size_m,
|
|
|
- )
|
|
|
- else:
|
|
|
- self._run_tile_large_m(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- c_bf16_flat,
|
|
|
- scales_i32_flat,
|
|
|
- global_scale,
|
|
|
- packed_route_indices,
|
|
|
- topk_weights_flat,
|
|
|
- c_tmp_f32_flat,
|
|
|
- locks_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- route_block_idx,
|
|
|
- expert_idx,
|
|
|
- output_n_tile,
|
|
|
- reduce_k_tile,
|
|
|
- reduce_tile_count,
|
|
|
- reduce_slice_count,
|
|
|
- reduce_slice_idx,
|
|
|
- lock_slot,
|
|
|
- active_size_m,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _tile_common_prologue(
|
|
|
- self,
|
|
|
- global_scale: cute.Tensor,
|
|
|
- packed_route_indices: cute.Tensor,
|
|
|
- topk_weights_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- route_block_idx: Int32,
|
|
|
- expert_idx: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- active_size_m: Int32,
|
|
|
- ):
|
|
|
- global_scale_f32 = global_scale[expert_idx].to(cutlass.Float32)
|
|
|
- if cutlass.const_expr(self.scale_format_e8m0_k32):
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- global_scale_f32 *= cutlass.Float32(_E8M0_K32_FP16_GLOBAL_COMPENSATION)
|
|
|
- else:
|
|
|
- global_scale_f32 *= cutlass.Float32(_E8M0_K32_BF16_GLOBAL_COMPENSATION)
|
|
|
- block_valid_rows = self._read_moe_block_data(
|
|
|
- packed_route_indices,
|
|
|
- topk_weights_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- route_block_idx,
|
|
|
- global_scale_f32,
|
|
|
- active_size_m,
|
|
|
- )
|
|
|
- (
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- ) = self._tile_stream_offsets(tid, expert_idx, output_n_tile)
|
|
|
- return (
|
|
|
- global_scale_f32,
|
|
|
- block_valid_rows,
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _tile_stream_offsets(self, tid: Int32, expert_idx: Int32, output_n_tile: Int32):
|
|
|
- a_gl_stride = Int32(self.size_k // 8)
|
|
|
- b_gl_stride = Int32(16 * self.size_n // (_PACK_FACTOR * 4))
|
|
|
- s_gl_stride = Int32(self.scale_n_groups)
|
|
|
- if cutlass.const_expr(self.scale_k32):
|
|
|
- if cutlass.const_expr(
|
|
|
- self.scale_format_e8m0_k32 and self.has_logical_tail
|
|
|
- ):
|
|
|
- scales_expert_stride = Int32(self.scale_k_groups * self.scale_n_groups)
|
|
|
- else:
|
|
|
- scales_expert_stride = Int32((self.size_n * self.size_k) // (32 * 16))
|
|
|
- else:
|
|
|
- scales_expert_stride = Int32((self.size_n * self.size_k) // (16 * 16))
|
|
|
- b_expert_off = (
|
|
|
- Int32((self.size_n * self.size_k) // (_PACK_FACTOR * 4)) * expert_idx
|
|
|
- )
|
|
|
- scales_expert_off = scales_expert_stride * expert_idx
|
|
|
-
|
|
|
- a_gl_rd_row = tid // Int32(self.a_gl_rd_delta_o)
|
|
|
- a_gl_rd_col0 = tid - a_gl_rd_row * Int32(self.a_gl_rd_delta_o)
|
|
|
- a_sh_wr = Int32(self.a_sh_stride) * (tid // Int32(self.a_gl_rd_delta_o)) + (
|
|
|
- tid - (tid // Int32(self.a_gl_rd_delta_o)) * Int32(self.a_gl_rd_delta_o)
|
|
|
- )
|
|
|
- a_rows_per_iter = Int32(self.cta_threads // self.a_gl_rd_delta_o)
|
|
|
-
|
|
|
- if cutlass.const_expr(self.cta_threads <= self.b_sh_stride):
|
|
|
- b_gl_rd_base = tid
|
|
|
- else:
|
|
|
- b_gl_rd_base = b_gl_stride * (tid // Int32(self.b_sh_stride)) + (
|
|
|
- tid % Int32(self.b_sh_stride)
|
|
|
- )
|
|
|
- b_gl_rd_base += b_expert_off + Int32(self.b_sh_stride) * output_n_tile
|
|
|
- b_sh_rd = tid
|
|
|
- b_sh_rd += (b_sh_rd // Int32(self.b_sh_stride)) * Int32(
|
|
|
- self.b_sh_stride * (self.b_sh_wr_iters - 1)
|
|
|
- )
|
|
|
-
|
|
|
- s_sh_rd = Int32(8) * ((tid // Int32(32)) % Int32(self.tb_n_warps)) + (
|
|
|
- tid & Int32(31)
|
|
|
- ) // Int32(4)
|
|
|
- return (
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _a_shared_read_offset(self, tid: Int32, lanes_per_row: cutlass.Constexpr[int]):
|
|
|
- a_sh_rd = Int32(self.a_sh_stride) * (
|
|
|
- (tid & Int32(31)) % Int32(lanes_per_row)
|
|
|
- ) + (tid & Int32(31)) // Int32(lanes_per_row)
|
|
|
- a_sh_rd += (
|
|
|
- Int32(2)
|
|
|
- * ((tid // Int32(32)) // Int32(self.tb_n_warps))
|
|
|
- * Int32(self.b_sh_wr_iters)
|
|
|
- )
|
|
|
- return a_sh_rd
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _run_tile_m8(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- b_i32_flat: cute.Tensor,
|
|
|
- c_bf16_flat: cute.Tensor,
|
|
|
- scales_i32_flat: cute.Tensor,
|
|
|
- global_scale: cute.Tensor,
|
|
|
- packed_route_indices: cute.Tensor,
|
|
|
- topk_weights_flat: cute.Tensor,
|
|
|
- c_tmp_f32_flat: cute.Tensor,
|
|
|
- locks_i32_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- route_block_idx: Int32,
|
|
|
- expert_idx: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- reduce_k_tile: Int32,
|
|
|
- reduce_tile_count: Int32,
|
|
|
- reduce_slice_count: Int32,
|
|
|
- reduce_slice_idx: Int32,
|
|
|
- lock_slot: Int32,
|
|
|
- active_size_m: Int32,
|
|
|
- ):
|
|
|
- (
|
|
|
- global_scale_f32,
|
|
|
- block_valid_rows,
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- ) = self._tile_common_prologue(
|
|
|
- global_scale,
|
|
|
- packed_route_indices,
|
|
|
- topk_weights_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- route_block_idx,
|
|
|
- expert_idx,
|
|
|
- output_n_tile,
|
|
|
- active_size_m,
|
|
|
- )
|
|
|
- a_sh_rd = self._a_shared_read_offset(tid, 8)
|
|
|
-
|
|
|
- acc = cute.make_rmem_tensor((4, 4), cutlass.Float32)
|
|
|
- acc.fill(0.0)
|
|
|
-
|
|
|
- k_tiles = reduce_tile_count
|
|
|
- self._prefetch_initial_tiles(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- scales_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- k_tiles,
|
|
|
- reduce_k_tile,
|
|
|
- block_valid_rows,
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- output_n_tile,
|
|
|
- expert_idx,
|
|
|
- )
|
|
|
-
|
|
|
- b_scale_cur = cute.make_rmem_tensor((2, 4), Uint32)
|
|
|
- b_scale_next = cute.make_rmem_tensor((2, 4), Uint32)
|
|
|
- self._load_b_scale_register_bundle(
|
|
|
- b_scale_cur,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- Int32(0),
|
|
|
- Int32(0),
|
|
|
- )
|
|
|
- a_regs_cur = cute.make_rmem_tensor((2,), Uint32)
|
|
|
- a_regs_next = cute.make_rmem_tensor((2,), Uint32)
|
|
|
- self._load_a_register_bundle(
|
|
|
- a_regs_cur,
|
|
|
- smem_base,
|
|
|
- a_sh_rd,
|
|
|
- Int32(0),
|
|
|
- Int32(0),
|
|
|
- True,
|
|
|
- )
|
|
|
- self._run_mma_pipeline(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- scales_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- acc,
|
|
|
- b_scale_cur,
|
|
|
- b_scale_next,
|
|
|
- a_regs_cur,
|
|
|
- a_regs_next,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- a_sh_rd,
|
|
|
- k_tiles,
|
|
|
- reduce_k_tile,
|
|
|
- block_valid_rows,
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- output_n_tile,
|
|
|
- expert_idx,
|
|
|
- True,
|
|
|
- )
|
|
|
-
|
|
|
- self._finish_tile(
|
|
|
- acc,
|
|
|
- c_bf16_flat,
|
|
|
- c_tmp_f32_flat,
|
|
|
- locks_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- output_n_tile,
|
|
|
- block_valid_rows,
|
|
|
- global_scale_f32,
|
|
|
- reduce_slice_count,
|
|
|
- reduce_slice_idx,
|
|
|
- lock_slot,
|
|
|
- True,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _run_tile_large_m(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- b_i32_flat: cute.Tensor,
|
|
|
- c_bf16_flat: cute.Tensor,
|
|
|
- scales_i32_flat: cute.Tensor,
|
|
|
- global_scale: cute.Tensor,
|
|
|
- packed_route_indices: cute.Tensor,
|
|
|
- topk_weights_flat: cute.Tensor,
|
|
|
- c_tmp_f32_flat: cute.Tensor,
|
|
|
- locks_i32_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- route_block_idx: Int32,
|
|
|
- expert_idx: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- reduce_k_tile: Int32,
|
|
|
- reduce_tile_count: Int32,
|
|
|
- reduce_slice_count: Int32,
|
|
|
- reduce_slice_idx: Int32,
|
|
|
- lock_slot: Int32,
|
|
|
- active_size_m: Int32,
|
|
|
- ):
|
|
|
- (
|
|
|
- global_scale_f32,
|
|
|
- block_valid_rows,
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- ) = self._tile_common_prologue(
|
|
|
- global_scale,
|
|
|
- packed_route_indices,
|
|
|
- topk_weights_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- route_block_idx,
|
|
|
- expert_idx,
|
|
|
- output_n_tile,
|
|
|
- active_size_m,
|
|
|
- )
|
|
|
- a_sh_rd = self._a_shared_read_offset(tid, 16)
|
|
|
-
|
|
|
- acc = cute.make_rmem_tensor((self.cta_m_blocks, 4, 2, 4), cutlass.Float32)
|
|
|
- acc.fill(0.0)
|
|
|
-
|
|
|
- k_tiles = reduce_tile_count
|
|
|
- self._prefetch_initial_tiles(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- scales_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- k_tiles,
|
|
|
- reduce_k_tile,
|
|
|
- block_valid_rows,
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- output_n_tile,
|
|
|
- expert_idx,
|
|
|
- )
|
|
|
-
|
|
|
- b_scale_cur = cute.make_rmem_tensor((2, 4), Uint32)
|
|
|
- b_scale_next = cute.make_rmem_tensor((2, 4), Uint32)
|
|
|
- self._load_b_scale_register_bundle(
|
|
|
- b_scale_cur,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- Int32(0),
|
|
|
- Int32(0),
|
|
|
- )
|
|
|
- a_regs = cute.make_rmem_tensor((self.cta_m_blocks, 4), Uint32)
|
|
|
- a_regs_next = cute.make_rmem_tensor((self.cta_m_blocks, 4), Uint32)
|
|
|
- self._load_a_register_bundle(
|
|
|
- a_regs,
|
|
|
- smem_base,
|
|
|
- a_sh_rd,
|
|
|
- Int32(0),
|
|
|
- Int32(0),
|
|
|
- False,
|
|
|
- )
|
|
|
- self._run_mma_pipeline(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- scales_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- acc,
|
|
|
- b_scale_cur,
|
|
|
- b_scale_next,
|
|
|
- a_regs,
|
|
|
- a_regs_next,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- a_sh_rd,
|
|
|
- k_tiles,
|
|
|
- reduce_k_tile,
|
|
|
- block_valid_rows,
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- output_n_tile,
|
|
|
- expert_idx,
|
|
|
- False,
|
|
|
- )
|
|
|
-
|
|
|
- self._finish_tile(
|
|
|
- acc,
|
|
|
- c_bf16_flat,
|
|
|
- c_tmp_f32_flat,
|
|
|
- locks_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- output_n_tile,
|
|
|
- block_valid_rows,
|
|
|
- global_scale_f32,
|
|
|
- reduce_slice_count,
|
|
|
- reduce_slice_idx,
|
|
|
- lock_slot,
|
|
|
- False,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _run_mma_pipeline(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- b_i32_flat: cute.Tensor,
|
|
|
- scales_i32_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- acc: cute.Tensor,
|
|
|
- b_scale_cur: cute.Tensor,
|
|
|
- b_scale_next: cute.Tensor,
|
|
|
- a_regs_cur: cute.Tensor,
|
|
|
- a_regs_next: cute.Tensor,
|
|
|
- b_sh_rd: Int32,
|
|
|
- s_sh_rd: Int32,
|
|
|
- a_sh_rd: Int32,
|
|
|
- k_tiles: Int32,
|
|
|
- reduce_k_tile: Int32,
|
|
|
- block_valid_rows: Int32,
|
|
|
- a_gl_stride: Int32,
|
|
|
- b_gl_stride: Int32,
|
|
|
- s_gl_stride: Int32,
|
|
|
- scales_expert_off: Int32,
|
|
|
- b_gl_rd_base: Int32,
|
|
|
- a_gl_rd_row: Int32,
|
|
|
- a_gl_rd_col0: Int32,
|
|
|
- a_sh_wr: Int32,
|
|
|
- a_rows_per_iter: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- expert_idx: Int32,
|
|
|
- uses_m_block_8: cutlass.Constexpr[bool],
|
|
|
- ):
|
|
|
- b_frag = cute.make_rmem_tensor((2, 2), Uint32)
|
|
|
- tile_idx = Int32(0)
|
|
|
- while tile_idx < k_tiles:
|
|
|
- for pipe in cutlass.range_constexpr(_STAGES):
|
|
|
- if tile_idx < k_tiles:
|
|
|
- for kk in cutlass.range_constexpr(self.b_sh_wr_iters):
|
|
|
- self._load_next_fragment_bundle(
|
|
|
- b_scale_next,
|
|
|
- a_regs_next,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- a_sh_rd,
|
|
|
- pipe,
|
|
|
- kk,
|
|
|
- tile_idx,
|
|
|
- k_tiles,
|
|
|
- uses_m_block_8,
|
|
|
- )
|
|
|
-
|
|
|
- self._prefetch_pipeline_step(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- scales_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- pipe,
|
|
|
- kk,
|
|
|
- tile_idx,
|
|
|
- k_tiles,
|
|
|
- reduce_k_tile,
|
|
|
- block_valid_rows,
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- output_n_tile,
|
|
|
- expert_idx,
|
|
|
- )
|
|
|
-
|
|
|
- for jj in cutlass.range_constexpr(4):
|
|
|
- if cutlass.const_expr(self.weight_layout_nf3):
|
|
|
- # NF3 triple: lo16 = half (jj % 2) of word
|
|
|
- # (jj // 2); hi8 = byte jj of the hi word; scale
|
|
|
- # register is the same per-jj lane as packed.
|
|
|
- lo_w = b_scale_cur[0, jj // 2]
|
|
|
- lo16 = (lo_w >> Uint32(16 * (jj % 2))) & Uint32(0xFFFF)
|
|
|
- hi8 = (b_scale_cur[0, 2] >> Uint32(8 * jj)) & Uint32(
|
|
|
- 0xFF
|
|
|
- )
|
|
|
- s = b_scale_cur[1, jj]
|
|
|
- self._scaled_dequant_b_fragment_nf3(
|
|
|
- b_frag, lo16, hi8, s
|
|
|
- )
|
|
|
- else:
|
|
|
- q, s = self._select_b_scale_register(jj, b_scale_cur)
|
|
|
- self._scaled_dequant_b_fragment(b_frag, q, s)
|
|
|
- if cutlass.const_expr(uses_m_block_8):
|
|
|
- self._mma_accumulate_m8(
|
|
|
- acc,
|
|
|
- jj,
|
|
|
- a_regs_cur,
|
|
|
- b_frag,
|
|
|
- )
|
|
|
- else:
|
|
|
- for mb in cutlass.range_constexpr(self.cta_m_blocks):
|
|
|
- self._mma_accumulate_large_m(
|
|
|
- acc,
|
|
|
- a_regs_cur,
|
|
|
- mb,
|
|
|
- jj,
|
|
|
- b_frag,
|
|
|
- )
|
|
|
-
|
|
|
- if cutlass.const_expr(uses_m_block_8):
|
|
|
- self._copy_a_register_bundle(
|
|
|
- a_regs_cur,
|
|
|
- a_regs_next,
|
|
|
- uses_m_block_8,
|
|
|
- )
|
|
|
- self._copy_b_scale_register_bundle(
|
|
|
- b_scale_cur, b_scale_next
|
|
|
- )
|
|
|
- else:
|
|
|
- self._copy_b_scale_register_bundle(
|
|
|
- b_scale_cur, b_scale_next
|
|
|
- )
|
|
|
- self._copy_a_register_bundle(
|
|
|
- a_regs_cur,
|
|
|
- a_regs_next,
|
|
|
- uses_m_block_8,
|
|
|
- )
|
|
|
- tile_idx += Int32(1)
|
|
|
- cute.arch.sync_threads()
|
|
|
- if tile_idx < k_tiles:
|
|
|
- self._load_b_scale_register_bundle(
|
|
|
- b_scale_cur,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- Int32(0),
|
|
|
- Int32(0),
|
|
|
- )
|
|
|
- self._load_a_register_bundle(
|
|
|
- a_regs_cur,
|
|
|
- smem_base,
|
|
|
- a_sh_rd,
|
|
|
- Int32(0),
|
|
|
- Int32(0),
|
|
|
- uses_m_block_8,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _finish_tile(
|
|
|
- self,
|
|
|
- acc: cute.Tensor,
|
|
|
- c_bf16_flat: cute.Tensor,
|
|
|
- c_tmp_f32_flat: cute.Tensor,
|
|
|
- locks_i32_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- block_valid_rows: Int32,
|
|
|
- global_scale_f32: cutlass.Float32,
|
|
|
- reduce_slice_count: Int32,
|
|
|
- reduce_slice_idx: Int32,
|
|
|
- lock_slot: Int32,
|
|
|
- uses_m_block_8: cutlass.Constexpr[bool],
|
|
|
- ):
|
|
|
- if cutlass.const_expr(uses_m_block_8):
|
|
|
- self._fold_cta_partials_m8(acc, smem_base, tid)
|
|
|
- else:
|
|
|
- self._fold_cta_partials_large_m(acc, smem_base, tid)
|
|
|
-
|
|
|
- if reduce_slice_count > Int32(1):
|
|
|
- self._wait_for_reduction_turn(
|
|
|
- locks_i32_flat, lock_slot, reduce_slice_idx, tid
|
|
|
- )
|
|
|
- self._combine_splitk_accumulators(
|
|
|
- acc,
|
|
|
- c_tmp_f32_flat,
|
|
|
- block_valid_rows,
|
|
|
- lock_slot,
|
|
|
- reduce_slice_idx,
|
|
|
- reduce_slice_count,
|
|
|
- tid,
|
|
|
- uses_m_block_8,
|
|
|
- )
|
|
|
- self._publish_reduction_turn(
|
|
|
- locks_i32_flat,
|
|
|
- lock_slot,
|
|
|
- reduce_slice_idx == reduce_slice_count - Int32(1),
|
|
|
- tid,
|
|
|
- )
|
|
|
-
|
|
|
- if reduce_slice_idx == reduce_slice_count - Int32(1):
|
|
|
- if cutlass.const_expr(uses_m_block_8):
|
|
|
- self._store_tile_m8(
|
|
|
- acc,
|
|
|
- c_bf16_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- output_n_tile,
|
|
|
- block_valid_rows,
|
|
|
- global_scale_f32,
|
|
|
- )
|
|
|
- else:
|
|
|
- self._store_tile_large_m(
|
|
|
- acc,
|
|
|
- c_bf16_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- output_n_tile,
|
|
|
- block_valid_rows,
|
|
|
- global_scale_f32,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _wait_for_reduction_turn(
|
|
|
- self,
|
|
|
- locks_i32_flat: cute.Tensor,
|
|
|
- lock_slot: Int32,
|
|
|
- count: Int32,
|
|
|
- tid: Int32,
|
|
|
- ):
|
|
|
- lock_addr = get_ptr_as_int64(locks_i32_flat, lock_slot)
|
|
|
- if tid == Int32(0):
|
|
|
- state = Int32(-1)
|
|
|
- while state != count:
|
|
|
- state = ld_global_acquire_i32(lock_addr)
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _publish_reduction_turn(
|
|
|
- self,
|
|
|
- locks_i32_flat: cute.Tensor,
|
|
|
- lock_slot: Int32,
|
|
|
- reset,
|
|
|
- tid: Int32,
|
|
|
- ):
|
|
|
- lock_addr = get_ptr_as_int64(locks_i32_flat, lock_slot)
|
|
|
- cute.arch.sync_threads()
|
|
|
- if tid == Int32(0):
|
|
|
- if reset:
|
|
|
- st_global_i32(lock_addr, Int32(0))
|
|
|
- else:
|
|
|
- red_add_global_release_i32(lock_addr, Int32(1))
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _merge_splitk_vec4(
|
|
|
- self,
|
|
|
- c_tmp_f32_flat: cute.Tensor,
|
|
|
- f32_off: Int32,
|
|
|
- reduce_slice_idx: Int32,
|
|
|
- reduce_slice_count: Int32,
|
|
|
- c0: cutlass.Float32,
|
|
|
- c1: cutlass.Float32,
|
|
|
- c2: cutlass.Float32,
|
|
|
- c3: cutlass.Float32,
|
|
|
- ):
|
|
|
- if reduce_slice_idx != Int32(0):
|
|
|
- r0, r1, r2, r3 = ld_global_v4_f32(get_ptr_as_int64(c_tmp_f32_flat, f32_off))
|
|
|
- c0 = c0 + r0
|
|
|
- c1 = c1 + r1
|
|
|
- c2 = c2 + r2
|
|
|
- c3 = c3 + r3
|
|
|
- if reduce_slice_idx != reduce_slice_count - Int32(1):
|
|
|
- st_global_v4_f32(
|
|
|
- get_ptr_as_int64(c_tmp_f32_flat, f32_off),
|
|
|
- c0,
|
|
|
- c1,
|
|
|
- c2,
|
|
|
- c3,
|
|
|
- )
|
|
|
- return c0, c1, c2, c3
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _combine_splitk_accumulators(
|
|
|
- self,
|
|
|
- acc: cute.Tensor,
|
|
|
- c_tmp_f32_flat: cute.Tensor,
|
|
|
- block_valid_rows: Int32,
|
|
|
- lock_slot: Int32,
|
|
|
- reduce_slice_idx: Int32,
|
|
|
- reduce_slice_count: Int32,
|
|
|
- tid: Int32,
|
|
|
- uses_m_block_8: cutlass.Constexpr[bool],
|
|
|
- ):
|
|
|
- active_threads = Int32(32 * self.tb_n_warps)
|
|
|
- c_size_int4 = Int32((self.cta_m_blocks * 16 * self.cta_n_blocks * 16) // 4)
|
|
|
- c_cur_offset = lock_slot * c_size_int4
|
|
|
- if cutlass.const_expr(uses_m_block_8):
|
|
|
- if tid < active_threads:
|
|
|
- for jj in cutlass.range_constexpr(4):
|
|
|
- k = jj * 2
|
|
|
- acc[jj, 0], acc[jj, 1], acc[jj, 2], acc[jj, 3] = (
|
|
|
- self._merge_splitk_slot(
|
|
|
- c_tmp_f32_flat,
|
|
|
- c_cur_offset,
|
|
|
- active_threads,
|
|
|
- Int32(k),
|
|
|
- tid,
|
|
|
- reduce_slice_idx,
|
|
|
- reduce_slice_count,
|
|
|
- acc[jj, 0],
|
|
|
- acc[jj, 1],
|
|
|
- acc[jj, 2],
|
|
|
- acc[jj, 3],
|
|
|
- )
|
|
|
- )
|
|
|
- else:
|
|
|
- lane_row = (tid & Int32(31)) // Int32(4)
|
|
|
- if tid < active_threads:
|
|
|
- for k in cutlass.range_constexpr(self.cta_m_blocks * 8):
|
|
|
- mb = k // 8
|
|
|
- flat_j = k % 8
|
|
|
- jj = flat_j // 2
|
|
|
- half = flat_j % 2
|
|
|
- row_valid = Int32(mb * 16) + lane_row < block_valid_rows
|
|
|
- if row_valid:
|
|
|
- (
|
|
|
- acc[mb, jj, half, 0],
|
|
|
- acc[mb, jj, half, 1],
|
|
|
- acc[mb, jj, half, 2],
|
|
|
- acc[mb, jj, half, 3],
|
|
|
- ) = self._merge_splitk_slot(
|
|
|
- c_tmp_f32_flat,
|
|
|
- c_cur_offset,
|
|
|
- active_threads,
|
|
|
- Int32(k),
|
|
|
- tid,
|
|
|
- reduce_slice_idx,
|
|
|
- reduce_slice_count,
|
|
|
- acc[mb, jj, half, 0],
|
|
|
- acc[mb, jj, half, 1],
|
|
|
- acc[mb, jj, half, 2],
|
|
|
- acc[mb, jj, half, 3],
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _merge_splitk_slot(
|
|
|
- self,
|
|
|
- c_tmp_f32_flat: cute.Tensor,
|
|
|
- c_cur_offset: Int32,
|
|
|
- active_threads: Int32,
|
|
|
- slot: Int32,
|
|
|
- tid: Int32,
|
|
|
- reduce_slice_idx: Int32,
|
|
|
- reduce_slice_count: Int32,
|
|
|
- c0: cutlass.Float32,
|
|
|
- c1: cutlass.Float32,
|
|
|
- c2: cutlass.Float32,
|
|
|
- c3: cutlass.Float32,
|
|
|
- ):
|
|
|
- int4_off = c_cur_offset + active_threads * slot + tid
|
|
|
- return self._merge_splitk_vec4(
|
|
|
- c_tmp_f32_flat,
|
|
|
- int4_off * Int32(4),
|
|
|
- reduce_slice_idx,
|
|
|
- reduce_slice_count,
|
|
|
- c0,
|
|
|
- c1,
|
|
|
- c2,
|
|
|
- c3,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _load_a_registers_large_m(
|
|
|
- self,
|
|
|
- smem_base: Int32,
|
|
|
- a_sh_rd: Int32,
|
|
|
- pipe: Int32,
|
|
|
- kk: Int32,
|
|
|
- m_block: Int32,
|
|
|
- ):
|
|
|
- a_addr = self._int4_addr(
|
|
|
- smem_base,
|
|
|
- Int32(self.sh_a_off)
|
|
|
- + pipe * Int32(self.a_sh_stage)
|
|
|
- + self._activation_smem_permuted_offset(
|
|
|
- Int32(2) * kk + m_block * Int32(self.a_sh_rd_delta_i) + a_sh_rd
|
|
|
- ),
|
|
|
- )
|
|
|
- return ldmatrix_m8n8x4_b16(a_addr)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _load_a_registers_m8(
|
|
|
- self,
|
|
|
- smem_base: Int32,
|
|
|
- a_sh_rd: Int32,
|
|
|
- pipe: Int32,
|
|
|
- kk: Int32,
|
|
|
- ):
|
|
|
- a_addr = self._int4_addr(
|
|
|
- smem_base,
|
|
|
- Int32(self.sh_a_off)
|
|
|
- + pipe * Int32(self.a_sh_stage)
|
|
|
- + self._activation_smem_permuted_offset(Int32(2) * kk + a_sh_rd),
|
|
|
- )
|
|
|
- return ldmatrix_m8n8x2_b16(a_addr)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _load_a_registers_large_m_bundle(
|
|
|
- self,
|
|
|
- regs: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- a_sh_rd: Int32,
|
|
|
- pipe: Int32,
|
|
|
- kk: Int32,
|
|
|
- ):
|
|
|
- for mb in cutlass.range_constexpr(self.cta_m_blocks):
|
|
|
- a0, a1, a2, a3 = self._load_a_registers_large_m(
|
|
|
- smem_base,
|
|
|
- a_sh_rd,
|
|
|
- pipe,
|
|
|
- kk,
|
|
|
- Int32(mb),
|
|
|
- )
|
|
|
- regs[mb, 0] = a0
|
|
|
- regs[mb, 1] = a1
|
|
|
- regs[mb, 2] = a2
|
|
|
- regs[mb, 3] = a3
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _load_a_registers_m8_bundle(
|
|
|
- self,
|
|
|
- regs: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- a_sh_rd: Int32,
|
|
|
- pipe: Int32,
|
|
|
- kk: Int32,
|
|
|
- ):
|
|
|
- a0, a1 = self._load_a_registers_m8(smem_base, a_sh_rd, pipe, kk)
|
|
|
- regs[0] = a0
|
|
|
- regs[1] = a1
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _clear_a_register_bundle_large_m(self, regs: cute.Tensor):
|
|
|
- for mb in cutlass.range_constexpr(self.cta_m_blocks):
|
|
|
- for reg in cutlass.range_constexpr(4):
|
|
|
- regs[mb, reg] = Uint32(0)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _clear_a_register_bundle_m8(self, regs: cute.Tensor):
|
|
|
- for reg in cutlass.range_constexpr(2):
|
|
|
- regs[reg] = Uint32(0)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _copy_a_register_bundle_large_m(self, dst: cute.Tensor, src: cute.Tensor):
|
|
|
- for mb in cutlass.range_constexpr(self.cta_m_blocks):
|
|
|
- for reg in cutlass.range_constexpr(4):
|
|
|
- dst[mb, reg] = src[mb, reg]
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _copy_a_register_bundle_m8(self, dst: cute.Tensor, src: cute.Tensor):
|
|
|
- for reg in cutlass.range_constexpr(2):
|
|
|
- dst[reg] = src[reg]
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _load_a_register_bundle(
|
|
|
- self,
|
|
|
- regs: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- a_sh_rd: Int32,
|
|
|
- pipe: Int32,
|
|
|
- kk: Int32,
|
|
|
- uses_m_block_8: cutlass.Constexpr[bool],
|
|
|
- ):
|
|
|
- if cutlass.const_expr(uses_m_block_8):
|
|
|
- self._load_a_registers_m8_bundle(regs, smem_base, a_sh_rd, pipe, kk)
|
|
|
- else:
|
|
|
- self._load_a_registers_large_m_bundle(regs, smem_base, a_sh_rd, pipe, kk)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _clear_a_register_bundle(
|
|
|
- self,
|
|
|
- regs: cute.Tensor,
|
|
|
- uses_m_block_8: cutlass.Constexpr[bool],
|
|
|
- ):
|
|
|
- if cutlass.const_expr(uses_m_block_8):
|
|
|
- self._clear_a_register_bundle_m8(regs)
|
|
|
- else:
|
|
|
- self._clear_a_register_bundle_large_m(regs)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _copy_a_register_bundle(
|
|
|
- self,
|
|
|
- dst: cute.Tensor,
|
|
|
- src: cute.Tensor,
|
|
|
- uses_m_block_8: cutlass.Constexpr[bool],
|
|
|
- ):
|
|
|
- if cutlass.const_expr(uses_m_block_8):
|
|
|
- self._copy_a_register_bundle_m8(dst, src)
|
|
|
- else:
|
|
|
- self._copy_a_register_bundle_large_m(dst, src)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _load_b_scale_registers(
|
|
|
- self,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- b_sh_rd: Int32,
|
|
|
- s_sh_rd: Int32,
|
|
|
- pipe: Int32,
|
|
|
- kk: Int32,
|
|
|
- ):
|
|
|
- if cutlass.const_expr(self.weight_layout == "modelopt"):
|
|
|
- q0, q1, q2, q3 = self._load_b_registers_modelopt_shared(
|
|
|
- smem_base,
|
|
|
- b_sh_rd,
|
|
|
- pipe,
|
|
|
- kk,
|
|
|
- )
|
|
|
- elif cutlass.const_expr(self.weight_layout_nf3):
|
|
|
- # NF3-MAPPING-V1: the thread's 32-code unit index within the pipe's
|
|
|
- # B region equals the stock staged-unit index (algebraically
|
|
|
- # identical to b_sh_stride*kk + b_sh_rd): cur_group_id is the K16 row
|
|
|
- # within the tile, (tid % b_sh_stride) is the N-pair. Each unit is a
|
|
|
- # 12-byte (lo0, lo1, hi) triple; three scalar u32 loads over stride-3
|
|
|
- # words are bank-conflict free (gcd(3, 32) == 1). q3 is unused.
|
|
|
- nf3_warp_id = tid // Int32(32)
|
|
|
- nf3_warp_row = nf3_warp_id // Int32(self.tb_n_warps)
|
|
|
- nf3_group = Int32(self.b_sh_wr_iters) * nf3_warp_row + kk
|
|
|
- nf3_unit = nf3_group * Int32(self.b_sh_stride) + (
|
|
|
- tid % Int32(self.b_sh_stride)
|
|
|
- )
|
|
|
- b_addr = (
|
|
|
- smem_base
|
|
|
- + Int32(self.sh_b_off * 16)
|
|
|
- + pipe * Int32(self.b_sh_stage_bytes)
|
|
|
- + nf3_unit * Int32(12)
|
|
|
- )
|
|
|
- q0 = ld_shared_u32(b_addr)
|
|
|
- q1 = ld_shared_u32(b_addr + Int32(4))
|
|
|
- q2 = ld_shared_u32(b_addr + Int32(8))
|
|
|
- q3 = Uint32(0)
|
|
|
- else:
|
|
|
- b_addr = self._int4_addr(
|
|
|
- smem_base,
|
|
|
- Int32(self.sh_b_off)
|
|
|
- + pipe * Int32(self.b_sh_stage)
|
|
|
- + Int32(self.b_sh_stride) * kk
|
|
|
- + b_sh_rd,
|
|
|
- )
|
|
|
- q0, q1, q2, q3 = ld_shared_v4_u32(b_addr)
|
|
|
-
|
|
|
- warp_id = tid // Int32(32)
|
|
|
- warp_row = warp_id // Int32(self.tb_n_warps)
|
|
|
- cur_group_id = Int32(self.b_sh_wr_iters) * warp_row + kk
|
|
|
- if cutlass.const_expr(self.scale_k32):
|
|
|
- scale_group_id = cur_group_id // Int32(2)
|
|
|
- else:
|
|
|
- scale_group_id = cur_group_id
|
|
|
- s_addr = (
|
|
|
- smem_base
|
|
|
- + Int32(self.sh_s_off * 16)
|
|
|
- + pipe * Int32(self.s_sh_stage * 16)
|
|
|
- + (s_sh_rd + scale_group_id * Int32(2 * self.s_sh_stride)) * Int32(8)
|
|
|
- )
|
|
|
- s_pack0, s_pack1 = ld_shared_v2_u32(s_addr)
|
|
|
- s0, s1 = self._dequant_scale_x4_to_elem2x2(s_pack0)
|
|
|
- s2, s3 = self._dequant_scale_x4_to_elem2x2(s_pack1)
|
|
|
- return q0, q1, q2, q3, s0, s1, s2, s3
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _load_b_scale_register_bundle(
|
|
|
- self,
|
|
|
- regs: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- b_sh_rd: Int32,
|
|
|
- s_sh_rd: Int32,
|
|
|
- pipe: Int32,
|
|
|
- kk: Int32,
|
|
|
- ):
|
|
|
- q0, q1, q2, q3, s0, s1, s2, s3 = self._load_b_scale_registers(
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- pipe,
|
|
|
- kk,
|
|
|
- )
|
|
|
- regs[0, 0] = q0
|
|
|
- regs[0, 1] = q1
|
|
|
- regs[0, 2] = q2
|
|
|
- regs[0, 3] = q3
|
|
|
- regs[1, 0] = s0
|
|
|
- regs[1, 1] = s1
|
|
|
- regs[1, 2] = s2
|
|
|
- regs[1, 3] = s3
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _clear_b_scale_register_bundle(self, regs: cute.Tensor):
|
|
|
- for row in cutlass.range_constexpr(2):
|
|
|
- for col in cutlass.range_constexpr(4):
|
|
|
- regs[row, col] = Uint32(0)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _copy_b_scale_register_bundle(self, dst: cute.Tensor, src: cute.Tensor):
|
|
|
- for row in cutlass.range_constexpr(2):
|
|
|
- for col in cutlass.range_constexpr(4):
|
|
|
- dst[row, col] = src[row, col]
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _select_b_scale_register(self, jj: cutlass.Constexpr[int], regs: cute.Tensor):
|
|
|
- return regs[0, jj], regs[1, jj]
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _load_next_fragment_bundle(
|
|
|
- self,
|
|
|
- b_scale_next: cute.Tensor,
|
|
|
- a_regs_next: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- b_sh_rd: Int32,
|
|
|
- s_sh_rd: Int32,
|
|
|
- a_sh_rd: Int32,
|
|
|
- pipe: cutlass.Constexpr[int],
|
|
|
- kk: cutlass.Constexpr[int],
|
|
|
- tile_idx: Int32,
|
|
|
- k_tiles: Int32,
|
|
|
- uses_m_block_8: cutlass.Constexpr[bool],
|
|
|
- ):
|
|
|
- self._clear_b_scale_register_bundle(b_scale_next)
|
|
|
- self._clear_a_register_bundle(a_regs_next, uses_m_block_8)
|
|
|
-
|
|
|
- if cutlass.const_expr(kk + 1 < self.b_sh_wr_iters):
|
|
|
- if tile_idx < k_tiles:
|
|
|
- self._load_b_scale_register_bundle(
|
|
|
- b_scale_next,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- Int32(pipe),
|
|
|
- Int32(kk + 1),
|
|
|
- )
|
|
|
- self._load_a_register_bundle(
|
|
|
- a_regs_next,
|
|
|
- smem_base,
|
|
|
- a_sh_rd,
|
|
|
- Int32(pipe),
|
|
|
- Int32(kk + 1),
|
|
|
- uses_m_block_8,
|
|
|
- )
|
|
|
- else:
|
|
|
- next_tile = tile_idx + Int32(1)
|
|
|
- if next_tile < k_tiles:
|
|
|
- self._load_b_scale_register_bundle(
|
|
|
- b_scale_next,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- b_sh_rd,
|
|
|
- s_sh_rd,
|
|
|
- Int32((pipe + 1) % _STAGES),
|
|
|
- Int32(0),
|
|
|
- )
|
|
|
- self._load_a_register_bundle(
|
|
|
- a_regs_next,
|
|
|
- smem_base,
|
|
|
- a_sh_rd,
|
|
|
- Int32((pipe + 1) % _STAGES),
|
|
|
- Int32(0),
|
|
|
- uses_m_block_8,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _scaled_dequant_b_fragment(self, frag: cute.Tensor, q: Uint32, s: Uint32):
|
|
|
- bq1 = q
|
|
|
- bq0 = bq1 << Uint32(8)
|
|
|
- b0_0, b0_1 = self._dequant_e2m1x4_to_elem2x2(bq0)
|
|
|
- b1_0, b1_1 = self._dequant_e2m1x4_to_elem2x2(bq1)
|
|
|
- s_lane0 = bfloat2_broadcast_lane(s, Int32(0))
|
|
|
- s_lane1 = bfloat2_broadcast_lane(s, Int32(1))
|
|
|
- b0_0 = self._elem2_mul(b0_0, s_lane0)
|
|
|
- b0_1 = self._elem2_mul(b0_1, s_lane0)
|
|
|
- b1_0 = self._elem2_mul(b1_0, s_lane1)
|
|
|
- b1_1 = self._elem2_mul(b1_1, s_lane1)
|
|
|
- frag[0, 0] = b0_0
|
|
|
- frag[0, 1] = b0_1
|
|
|
- frag[1, 0] = b1_0
|
|
|
- frag[1, 1] = b1_1
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _scaled_dequant_b_fragment_nf3(
|
|
|
- self, frag: cute.Tensor, lo16: Uint32, hi8: Uint32, s: Uint32
|
|
|
- ):
|
|
|
- # NF3: 8 3-bit codes (2 lo-bits in lo16, 1 hi-bit in hi8) -> 4 bf16x2
|
|
|
- # codebook fragments, in the SAME element order the packed path uses:
|
|
|
- # o0=(cb[c0],cb[c1]) -> frag[0,0], o1 -> frag[0,1], o2 -> frag[1,0],
|
|
|
- # o3 -> frag[1,1]. frag[0,*] takes the N-col-0 scale lane, frag[1,*] the
|
|
|
- # N-col-1 lane (identical broadcast to _scaled_dequant_b_fragment).
|
|
|
- o0, o1, o2, o3 = packed_dequant_nf3x8_to_bfloat2x4(
|
|
|
- lo16,
|
|
|
- hi8,
|
|
|
- Uint32(self._nf3_cb0),
|
|
|
- Uint32(self._nf3_cb1),
|
|
|
- Uint32(self._nf3_cb2),
|
|
|
- Uint32(self._nf3_cb3),
|
|
|
- )
|
|
|
- s_lane0 = bfloat2_broadcast_lane(s, Int32(0))
|
|
|
- s_lane1 = bfloat2_broadcast_lane(s, Int32(1))
|
|
|
- frag[0, 0] = self._elem2_mul(o0, s_lane0)
|
|
|
- frag[0, 1] = self._elem2_mul(o1, s_lane0)
|
|
|
- frag[1, 0] = self._elem2_mul(o2, s_lane1)
|
|
|
- frag[1, 1] = self._elem2_mul(o3, s_lane1)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _mma_accumulate_m8(
|
|
|
- self,
|
|
|
- acc: cute.Tensor,
|
|
|
- jj: cutlass.Constexpr[int],
|
|
|
- a_regs: cute.Tensor,
|
|
|
- b_frag: cute.Tensor,
|
|
|
- ):
|
|
|
- d0, d1, d2, d3 = self._mma_rhs_fragments_as_mma_a_m16n8k16_f32(
|
|
|
- acc[jj, 0],
|
|
|
- acc[jj, 1],
|
|
|
- acc[jj, 2],
|
|
|
- acc[jj, 3],
|
|
|
- b_frag[0, 0],
|
|
|
- b_frag[1, 0],
|
|
|
- b_frag[0, 1],
|
|
|
- b_frag[1, 1],
|
|
|
- a_regs[0],
|
|
|
- a_regs[1],
|
|
|
- )
|
|
|
- acc[jj, 0] = d0
|
|
|
- acc[jj, 1] = d1
|
|
|
- acc[jj, 2] = d2
|
|
|
- acc[jj, 3] = d3
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _mma_accumulate_large_m(
|
|
|
- self,
|
|
|
- acc: cute.Tensor,
|
|
|
- a_regs: cute.Tensor,
|
|
|
- mb: cutlass.Constexpr[int],
|
|
|
- jj: cutlass.Constexpr[int],
|
|
|
- b_frag: cute.Tensor,
|
|
|
- ):
|
|
|
- d0, d1, d2, d3 = self._mma_m16n8k16_f32(
|
|
|
- acc[mb, jj, 0, 0],
|
|
|
- acc[mb, jj, 0, 1],
|
|
|
- acc[mb, jj, 0, 2],
|
|
|
- acc[mb, jj, 0, 3],
|
|
|
- a_regs[mb, 0],
|
|
|
- a_regs[mb, 1],
|
|
|
- a_regs[mb, 2],
|
|
|
- a_regs[mb, 3],
|
|
|
- b_frag[0, 0],
|
|
|
- b_frag[0, 1],
|
|
|
- )
|
|
|
- acc[mb, jj, 0, 0] = d0
|
|
|
- acc[mb, jj, 0, 1] = d1
|
|
|
- acc[mb, jj, 0, 2] = d2
|
|
|
- acc[mb, jj, 0, 3] = d3
|
|
|
- d0, d1, d2, d3 = self._mma_m16n8k16_f32(
|
|
|
- acc[mb, jj, 1, 0],
|
|
|
- acc[mb, jj, 1, 1],
|
|
|
- acc[mb, jj, 1, 2],
|
|
|
- acc[mb, jj, 1, 3],
|
|
|
- a_regs[mb, 0],
|
|
|
- a_regs[mb, 1],
|
|
|
- a_regs[mb, 2],
|
|
|
- a_regs[mb, 3],
|
|
|
- b_frag[1, 0],
|
|
|
- b_frag[1, 1],
|
|
|
- )
|
|
|
- acc[mb, jj, 1, 0] = d0
|
|
|
- acc[mb, jj, 1, 1] = d1
|
|
|
- acc[mb, jj, 1, 2] = d2
|
|
|
- acc[mb, jj, 1, 3] = d3
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _source_n_from_logical(self, logical_n: Int32) -> Int32:
|
|
|
- source_n = logical_n
|
|
|
- if cutlass.const_expr(self.source_n_rotation != 0):
|
|
|
- source_n += Int32(self.source_n_rotation)
|
|
|
- if source_n >= Int32(self.size_n):
|
|
|
- source_n -= Int32(self.size_n)
|
|
|
- return source_n
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _stage_b_tile_modelopt_native(
|
|
|
- self,
|
|
|
- b_u8_flat: cute.Tensor,
|
|
|
- smem_addr: Int32,
|
|
|
- expert_idx: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- tile_idx: Int32,
|
|
|
- local_int4: Int32,
|
|
|
- ):
|
|
|
- chunks_per_row = Int32(self.tile_k // 32)
|
|
|
- local_n = local_int4 // chunks_per_row
|
|
|
- local_k_vec = local_int4 - local_n * chunks_per_row
|
|
|
- logical_n = output_n_tile * Int32(self.tile_n) + local_n
|
|
|
- source_n = self._source_n_from_logical(logical_n)
|
|
|
- packed_cols = Int32(self.size_k // 2)
|
|
|
- tile_byte = tile_idx * Int32(self.tile_k // 2) + local_k_vec * Int32(16)
|
|
|
- byte_offset = (
|
|
|
- Int64(expert_idx) * Int64(self.size_n * (self.size_k // 2))
|
|
|
- + Int64(source_n) * Int64(packed_cols)
|
|
|
- + Int64(tile_byte)
|
|
|
- )
|
|
|
- if cutlass.const_expr(self.has_logical_tail):
|
|
|
- valid_n = logical_n < Int32(self.size_n)
|
|
|
- valid_bytes = packed_cols - tile_byte
|
|
|
- if valid_bytes > Int32(16):
|
|
|
- valid_bytes = Int32(16)
|
|
|
- if valid_n and valid_bytes >= Int32(16) and cutlass.const_expr(
|
|
|
- not self.has_scale_k_tail
|
|
|
- ):
|
|
|
- cp_async4_shared_global(
|
|
|
- smem_addr,
|
|
|
- get_ptr_as_int64(b_u8_flat, byte_offset),
|
|
|
- )
|
|
|
- else:
|
|
|
- v0 = Uint32(0)
|
|
|
- v1 = Uint32(0)
|
|
|
- v2 = Uint32(0)
|
|
|
- v3 = Uint32(0)
|
|
|
- if valid_n and valid_bytes > Int32(0):
|
|
|
- src = get_ptr_as_int64(b_u8_flat, byte_offset)
|
|
|
- if valid_bytes >= Int32(4):
|
|
|
- v0 = ld_global_nc_u32(src)
|
|
|
- if valid_bytes >= Int32(8):
|
|
|
- v1 = ld_global_nc_u32(src + Int64(4))
|
|
|
- if valid_bytes >= Int32(12):
|
|
|
- v2 = ld_global_nc_u32(src + Int64(8))
|
|
|
- if valid_bytes >= Int32(16):
|
|
|
- v3 = ld_global_nc_u32(src + Int64(12))
|
|
|
- st_shared_v4_u32(smem_addr, v0, v1, v2, v3)
|
|
|
- else:
|
|
|
- cp_async4_shared_global(
|
|
|
- smem_addr,
|
|
|
- get_ptr_as_int64(b_u8_flat, byte_offset),
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _load_modelopt_shared_byte(
|
|
|
- self,
|
|
|
- smem_base: Int32,
|
|
|
- pipe: Int32,
|
|
|
- n_tile: Int32,
|
|
|
- k_tile: Int32,
|
|
|
- warp_id: Int32,
|
|
|
- tc_col: Int32,
|
|
|
- tc_row: Int32,
|
|
|
- n_delta: cutlass.Constexpr[int],
|
|
|
- k_delta: cutlass.Constexpr[int],
|
|
|
- ) -> Uint32:
|
|
|
- local_n = (
|
|
|
- n_tile * Int32(64)
|
|
|
- + warp_id * Int32(16)
|
|
|
- + tc_col
|
|
|
- + Int32(n_delta)
|
|
|
- )
|
|
|
- local_k = k_tile * Int32(16) + tc_row + Int32(k_delta)
|
|
|
- byte_offset = local_n * Int32(self.tile_k // 2) + local_k // Int32(2)
|
|
|
- word_byte_offset = byte_offset - (byte_offset & Int32(3))
|
|
|
- word = ld_shared_u32(
|
|
|
- smem_base
|
|
|
- + Int32(self.sh_b_off * 16)
|
|
|
- + pipe * Int32(self.b_sh_stage * 16)
|
|
|
- + word_byte_offset
|
|
|
- )
|
|
|
- shift = Uint32((byte_offset - word_byte_offset) * Int32(8))
|
|
|
- return (word >> shift) & Uint32(0xFF)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _pack_modelopt_byte_pair(
|
|
|
- self,
|
|
|
- word: Uint32,
|
|
|
- q: Uint32,
|
|
|
- low_shift: cutlass.Constexpr[int],
|
|
|
- high_shift: cutlass.Constexpr[int],
|
|
|
- ) -> Uint32:
|
|
|
- low = q & Uint32(0xF)
|
|
|
- high = (q >> Uint32(4)) & Uint32(0xF)
|
|
|
- return word | (low << Uint32(low_shift)) | (high << Uint32(high_shift))
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _load_modelopt_shared_packed_word_for_lane(
|
|
|
- self,
|
|
|
- smem_base: Int32,
|
|
|
- pipe: Int32,
|
|
|
- n_tile: Int32,
|
|
|
- k_tile: Int32,
|
|
|
- warp_id: Int32,
|
|
|
- tc_col: Int32,
|
|
|
- tc_row: Int32,
|
|
|
- ) -> Uint32:
|
|
|
- q0 = self._load_modelopt_shared_byte(
|
|
|
- smem_base, pipe, n_tile, k_tile, warp_id, tc_col, tc_row, 0, 0
|
|
|
- )
|
|
|
- q1 = self._load_modelopt_shared_byte(
|
|
|
- smem_base, pipe, n_tile, k_tile, warp_id, tc_col, tc_row, 0, 8
|
|
|
- )
|
|
|
- q2 = self._load_modelopt_shared_byte(
|
|
|
- smem_base, pipe, n_tile, k_tile, warp_id, tc_col, tc_row, 8, 0
|
|
|
- )
|
|
|
- q3 = self._load_modelopt_shared_byte(
|
|
|
- smem_base, pipe, n_tile, k_tile, warp_id, tc_col, tc_row, 8, 8
|
|
|
- )
|
|
|
- word = Uint32(0)
|
|
|
- word = self._pack_modelopt_byte_pair(word, q0, 0, 16)
|
|
|
- word = self._pack_modelopt_byte_pair(word, q1, 4, 20)
|
|
|
- word = self._pack_modelopt_byte_pair(word, q2, 8, 24)
|
|
|
- word = self._pack_modelopt_byte_pair(word, q3, 12, 28)
|
|
|
- return word
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _load_b_registers_modelopt_shared(
|
|
|
- self,
|
|
|
- smem_base: Int32,
|
|
|
- b_sh_rd: Int32,
|
|
|
- pipe: Int32,
|
|
|
- kk: Int32,
|
|
|
- ):
|
|
|
- packed_word_index = (Int32(self.b_sh_stride) * kk + b_sh_rd) * Int32(4)
|
|
|
- words_per_k_tile = Int32((self.tile_n // 64) * 128)
|
|
|
- k_tile = packed_word_index // words_per_k_tile
|
|
|
- pos_in_k_tile = packed_word_index - k_tile * words_per_k_tile
|
|
|
- n_tile = pos_in_k_tile // Int32(128)
|
|
|
- pos = pos_in_k_tile - n_tile * Int32(128)
|
|
|
- th_id = pos // Int32(4)
|
|
|
- tc_col = th_id // Int32(4)
|
|
|
- tc_row = (th_id - tc_col * Int32(4)) * Int32(2)
|
|
|
- q0 = self._load_modelopt_shared_packed_word_for_lane(
|
|
|
- smem_base, pipe, n_tile, k_tile, Int32(0), tc_col, tc_row
|
|
|
- )
|
|
|
- q1 = self._load_modelopt_shared_packed_word_for_lane(
|
|
|
- smem_base, pipe, n_tile, k_tile, Int32(1), tc_col, tc_row
|
|
|
- )
|
|
|
- q2 = self._load_modelopt_shared_packed_word_for_lane(
|
|
|
- smem_base, pipe, n_tile, k_tile, Int32(2), tc_col, tc_row
|
|
|
- )
|
|
|
- q3 = self._load_modelopt_shared_packed_word_for_lane(
|
|
|
- smem_base, pipe, n_tile, k_tile, Int32(3), tc_col, tc_row
|
|
|
- )
|
|
|
- return q0, q1, q2, q3
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _stage_k_tile_async(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- b_i32_flat: cute.Tensor,
|
|
|
- scales_i32_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- pipe: Int32,
|
|
|
- tile_idx: Int32,
|
|
|
- block_valid_rows: Int32,
|
|
|
- a_gl_stride: Int32,
|
|
|
- b_gl_stride: Int32,
|
|
|
- s_gl_stride: Int32,
|
|
|
- scales_expert_off: Int32,
|
|
|
- b_gl_rd_base: Int32,
|
|
|
- a_gl_rd_row: Int32,
|
|
|
- a_gl_rd_col0: Int32,
|
|
|
- a_sh_wr: Int32,
|
|
|
- a_rows_per_iter: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- expert_idx: Int32,
|
|
|
- ):
|
|
|
- for i in cutlass.range_constexpr(self.a_sh_wr_iters):
|
|
|
- row = a_rows_per_iter * Int32(i) + a_gl_rd_row
|
|
|
- route_index = Int32(0)
|
|
|
- if row < Int32(self.moe_block_size):
|
|
|
- route_index = ld_shared_i32_relaxed(
|
|
|
- smem_base + Int32(self.sh_rd_route_off * 16) + row * Int32(4)
|
|
|
- )
|
|
|
- a_int4 = (
|
|
|
- route_index * a_gl_stride
|
|
|
- + tile_idx * Int32(self.a_gl_rd_delta_o)
|
|
|
- + a_gl_rd_col0
|
|
|
- )
|
|
|
- a_dst = self._int4_addr(
|
|
|
- smem_base,
|
|
|
- Int32(self.sh_a_off)
|
|
|
- + pipe * Int32(self.a_sh_stage)
|
|
|
- + self._activation_smem_permuted_offset(
|
|
|
- Int32(i * self.a_sh_wr_delta) + a_sh_wr
|
|
|
- ),
|
|
|
- )
|
|
|
- if cutlass.const_expr(self.has_k_tile_tail):
|
|
|
- a_k_int4 = tile_idx * Int32(self.a_gl_rd_delta_o) + a_gl_rd_col0
|
|
|
- if row < block_valid_rows and a_k_int4 < a_gl_stride:
|
|
|
- cp_async4_shared_global(
|
|
|
- a_dst,
|
|
|
- get_ptr_as_int64(a_bf16_flat, a_int4 * Int32(8)),
|
|
|
- )
|
|
|
- else:
|
|
|
- st_shared_v4_u32(
|
|
|
- a_dst, Uint32(0), Uint32(0), Uint32(0), Uint32(0)
|
|
|
- )
|
|
|
- else:
|
|
|
- cp_async4_shared_global_pred(
|
|
|
- a_dst,
|
|
|
- get_ptr_as_int64(a_bf16_flat, a_int4 * Int32(8)),
|
|
|
- (row < block_valid_rows).to(Int32),
|
|
|
- )
|
|
|
-
|
|
|
- if cutlass.const_expr(self.weight_layout_nf3):
|
|
|
- # NF3 flat-span staging: the packer lays the per-(expert,
|
|
|
- # output_n_tile, k-stage) B block out as ONE contiguous 12-byte-triple
|
|
|
- # span (n_tile-major, then K16-row, then N-pair), so we copy it
|
|
|
- # verbatim as 16-byte cp.async chunks. SMEM byte X of the pipe's B
|
|
|
- # region == global byte X of the span, so the register read (unit*12)
|
|
|
- # indexes it directly. b_gl_rd_base is unused on this path.
|
|
|
- nf3_units_per_expert = (self.size_n * self.size_k) // 32
|
|
|
- nf3_ntile_stride = (self.size_k // 16) * (self.tile_n // 2)
|
|
|
- nf3_span_base_unit = (
|
|
|
- Int32(nf3_units_per_expert) * expert_idx
|
|
|
- + Int32(nf3_ntile_stride) * output_n_tile
|
|
|
- + tile_idx * Int32(self.cta_k_blocks * self.b_sh_stride)
|
|
|
- )
|
|
|
- for i in cutlass.range_constexpr(self.b_sh_wr_iters_nf3):
|
|
|
- nf3_chunk = Int32(i * self.cta_threads) + tid
|
|
|
- b_dst = (
|
|
|
- smem_base
|
|
|
- + Int32(self.sh_b_off * 16)
|
|
|
- + pipe * Int32(self.b_sh_stage_bytes)
|
|
|
- + nf3_chunk * Int32(16)
|
|
|
- )
|
|
|
- # int32-element index of the chunk's first word:
|
|
|
- # (span_base_unit*12 + chunk*16) / 4 = span_base_unit*3 + chunk*4.
|
|
|
- b_src_i32 = nf3_span_base_unit * Int32(3) + nf3_chunk * Int32(4)
|
|
|
- cp_async4_shared_global_pred(
|
|
|
- b_dst,
|
|
|
- get_ptr_as_int64(b_i32_flat, b_src_i32),
|
|
|
- (nf3_chunk < Int32(self.b_sh_chunks)).to(Int32),
|
|
|
- )
|
|
|
-
|
|
|
- for i in cutlass.range_constexpr(
|
|
|
- 0 if self.weight_layout_nf3 else self.b_sh_wr_iters
|
|
|
- ):
|
|
|
- b_src_int4 = (
|
|
|
- b_gl_rd_base
|
|
|
- + tile_idx * Int32(self.cta_k_blocks) * b_gl_stride
|
|
|
- + Int32(i * (self.cta_threads // self.b_sh_stride)) * b_gl_stride
|
|
|
- )
|
|
|
- b_dst = self._int4_addr(
|
|
|
- smem_base,
|
|
|
- Int32(self.sh_b_off)
|
|
|
- + pipe * Int32(self.b_sh_stage)
|
|
|
- + Int32(i * self.cta_threads)
|
|
|
- + tid,
|
|
|
- )
|
|
|
- if cutlass.const_expr(self.weight_layout == "packed"):
|
|
|
- cp_async4_shared_global(
|
|
|
- b_dst,
|
|
|
- get_ptr_as_int64(b_i32_flat, b_src_int4 * Int32(4)),
|
|
|
- )
|
|
|
- else:
|
|
|
- self._stage_b_tile_modelopt_native(
|
|
|
- b_i32_flat,
|
|
|
- b_dst,
|
|
|
- expert_idx,
|
|
|
- output_n_tile,
|
|
|
- tile_idx,
|
|
|
- Int32(i * self.cta_threads) + tid,
|
|
|
- )
|
|
|
-
|
|
|
- if tid < Int32(self.s_sh_stage):
|
|
|
- s_k_group = tile_idx * Int32(self.s_tb_groups) + tid // Int32(
|
|
|
- self.s_sh_stride
|
|
|
- )
|
|
|
- s_n_group = Int32(self.s_sh_stride) * output_n_tile + (
|
|
|
- tid % Int32(self.s_sh_stride)
|
|
|
- )
|
|
|
- s_src_int4 = scales_expert_off + s_gl_stride * s_k_group + s_n_group
|
|
|
- s_dst = self._int4_addr(
|
|
|
- smem_base,
|
|
|
- Int32(self.sh_s_off) + pipe * Int32(self.s_sh_stage) + tid,
|
|
|
- )
|
|
|
- if cutlass.const_expr(self.has_logical_tail):
|
|
|
- if s_k_group < Int32(self.scale_k_groups) and s_n_group < Int32(
|
|
|
- self.scale_n_groups
|
|
|
- ):
|
|
|
- cp_async4_shared_global(
|
|
|
- s_dst,
|
|
|
- get_ptr_as_int64(scales_i32_flat, s_src_int4 * Int32(4)),
|
|
|
- )
|
|
|
- else:
|
|
|
- st_shared_v4_u32(
|
|
|
- s_dst, Uint32(0), Uint32(0), Uint32(0), Uint32(0)
|
|
|
- )
|
|
|
- else:
|
|
|
- cp_async4_shared_global(
|
|
|
- s_dst,
|
|
|
- get_ptr_as_int64(scales_i32_flat, s_src_int4 * Int32(4)),
|
|
|
- )
|
|
|
-
|
|
|
- cute.arch.cp_async_commit_group()
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _prefetch_pipeline_step(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- b_i32_flat: cute.Tensor,
|
|
|
- scales_i32_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- pipe: cutlass.Constexpr[int],
|
|
|
- kk: cutlass.Constexpr[int],
|
|
|
- tile_idx: Int32,
|
|
|
- k_tiles: Int32,
|
|
|
- reduce_k_tile: Int32,
|
|
|
- block_valid_rows: Int32,
|
|
|
- a_gl_stride: Int32,
|
|
|
- b_gl_stride: Int32,
|
|
|
- s_gl_stride: Int32,
|
|
|
- scales_expert_off: Int32,
|
|
|
- b_gl_rd_base: Int32,
|
|
|
- a_gl_rd_row: Int32,
|
|
|
- a_gl_rd_col0: Int32,
|
|
|
- a_sh_wr: Int32,
|
|
|
- a_rows_per_iter: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- expert_idx: Int32,
|
|
|
- ):
|
|
|
- if cutlass.const_expr(kk == self.b_sh_wr_iters - 2):
|
|
|
- self._prefetch_lookahead_tile(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- scales_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- pipe,
|
|
|
- tile_idx,
|
|
|
- k_tiles,
|
|
|
- reduce_k_tile,
|
|
|
- block_valid_rows,
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- output_n_tile,
|
|
|
- expert_idx,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _prefetch_initial_tiles(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- b_i32_flat: cute.Tensor,
|
|
|
- scales_i32_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- k_tiles: Int32,
|
|
|
- reduce_k_tile: Int32,
|
|
|
- block_valid_rows: Int32,
|
|
|
- a_gl_stride: Int32,
|
|
|
- b_gl_stride: Int32,
|
|
|
- s_gl_stride: Int32,
|
|
|
- scales_expert_off: Int32,
|
|
|
- b_gl_rd_base: Int32,
|
|
|
- a_gl_rd_row: Int32,
|
|
|
- a_gl_rd_col0: Int32,
|
|
|
- a_sh_wr: Int32,
|
|
|
- a_rows_per_iter: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- expert_idx: Int32,
|
|
|
- ):
|
|
|
- for pipe in cutlass.range_constexpr(_STAGES - 1):
|
|
|
- if Int32(pipe) < k_tiles:
|
|
|
- self._stage_k_tile_async(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- scales_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- Int32(pipe),
|
|
|
- reduce_k_tile + Int32(pipe),
|
|
|
- block_valid_rows,
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- output_n_tile,
|
|
|
- expert_idx,
|
|
|
- )
|
|
|
- else:
|
|
|
- cute.arch.cp_async_commit_group()
|
|
|
- cute.arch.cp_async_wait_group(_STAGES - 2)
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _prefetch_lookahead_tile(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- b_i32_flat: cute.Tensor,
|
|
|
- scales_i32_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- pipe: cutlass.Constexpr[int],
|
|
|
- tile_idx: Int32,
|
|
|
- k_tiles: Int32,
|
|
|
- reduce_k_tile: Int32,
|
|
|
- block_valid_rows: Int32,
|
|
|
- a_gl_stride: Int32,
|
|
|
- b_gl_stride: Int32,
|
|
|
- s_gl_stride: Int32,
|
|
|
- scales_expert_off: Int32,
|
|
|
- b_gl_rd_base: Int32,
|
|
|
- a_gl_rd_row: Int32,
|
|
|
- a_gl_rd_col0: Int32,
|
|
|
- a_sh_wr: Int32,
|
|
|
- a_rows_per_iter: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- expert_idx: Int32,
|
|
|
- ):
|
|
|
- fetch_tile = tile_idx + Int32(_STAGES - 1)
|
|
|
- if fetch_tile < k_tiles:
|
|
|
- self._stage_k_tile_async(
|
|
|
- a_bf16_flat,
|
|
|
- b_i32_flat,
|
|
|
- scales_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- Int32((pipe + _STAGES - 1) % _STAGES),
|
|
|
- reduce_k_tile + fetch_tile,
|
|
|
- block_valid_rows,
|
|
|
- a_gl_stride,
|
|
|
- b_gl_stride,
|
|
|
- s_gl_stride,
|
|
|
- scales_expert_off,
|
|
|
- b_gl_rd_base,
|
|
|
- a_gl_rd_row,
|
|
|
- a_gl_rd_col0,
|
|
|
- a_sh_wr,
|
|
|
- a_rows_per_iter,
|
|
|
- output_n_tile,
|
|
|
- expert_idx,
|
|
|
- )
|
|
|
- else:
|
|
|
- cute.arch.cp_async_commit_group()
|
|
|
- cute.arch.cp_async_wait_group(_STAGES - 2)
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _reduction_offsets(self, tid: Int32):
|
|
|
- red_idx = tid // Int32(self.b_sh_stride_threads)
|
|
|
- red_sh_stride = Int32(self.b_sh_stride_threads * 4 * 2)
|
|
|
- red_sh_delta = Int32(self.b_sh_stride_threads)
|
|
|
- red_sh_rd = red_sh_stride * (tid // Int32(self.b_sh_stride_threads)) + (
|
|
|
- tid % Int32(self.b_sh_stride_threads)
|
|
|
- )
|
|
|
- return red_idx, red_sh_stride, red_sh_delta, red_sh_rd
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _fold_cta_partials_m8(self, acc: cute.Tensor, smem_base: Int32, tid: Int32):
|
|
|
- red_off = self.cta_threads // self.b_sh_stride_threads // 2
|
|
|
- if cutlass.const_expr(red_off >= 1):
|
|
|
- red_idx, red_sh_stride, red_sh_delta, red_sh_rd = self._reduction_offsets(
|
|
|
- tid
|
|
|
- )
|
|
|
- if cutlass.const_expr(red_off == 2):
|
|
|
- if Int32(2) <= red_idx and red_idx < Int32(4):
|
|
|
- for jj in cutlass.range_constexpr(4):
|
|
|
- red_sh_wr = red_sh_delta * Int32(jj * 2) + (
|
|
|
- red_sh_rd - red_sh_stride * Int32(2)
|
|
|
- )
|
|
|
- st_shared_v4_f32(
|
|
|
- self._int4_addr(
|
|
|
- smem_base, Int32(self.sh_red_off) + red_sh_wr
|
|
|
- ),
|
|
|
- acc[jj, 0],
|
|
|
- acc[jj, 1],
|
|
|
- acc[jj, 2],
|
|
|
- acc[jj, 3],
|
|
|
- )
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- if Int32(1) <= red_idx and red_idx < Int32(2):
|
|
|
- for jj in cutlass.range_constexpr(4):
|
|
|
- red_sh_wr = red_sh_delta * Int32(jj * 2) + (
|
|
|
- red_sh_rd - red_sh_stride
|
|
|
- )
|
|
|
- if cutlass.const_expr(red_off > 1):
|
|
|
- rd_addr = self._int4_addr(
|
|
|
- smem_base,
|
|
|
- Int32(self.sh_red_off)
|
|
|
- + red_sh_delta * Int32(jj * 2)
|
|
|
- + red_sh_rd,
|
|
|
- )
|
|
|
- wr_addr = self._int4_addr(
|
|
|
- smem_base, Int32(self.sh_red_off) + red_sh_wr
|
|
|
- )
|
|
|
- r0, r1, r2, r3 = ld_shared_v4_f32(rd_addr)
|
|
|
- w0, w1, w2, w3 = ld_shared_v4_f32(wr_addr)
|
|
|
- acc[jj, 0] = acc[jj, 0] + r0 + w0
|
|
|
- acc[jj, 1] = acc[jj, 1] + r1 + w1
|
|
|
- acc[jj, 2] = acc[jj, 2] + r2 + w2
|
|
|
- acc[jj, 3] = acc[jj, 3] + r3 + w3
|
|
|
- st_shared_v4_f32(
|
|
|
- self._int4_addr(smem_base, Int32(self.sh_red_off) + red_sh_wr),
|
|
|
- acc[jj, 0],
|
|
|
- acc[jj, 1],
|
|
|
- acc[jj, 2],
|
|
|
- acc[jj, 3],
|
|
|
- )
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- if red_idx == Int32(0):
|
|
|
- for jj in cutlass.range_constexpr(4):
|
|
|
- rd_addr = self._int4_addr(
|
|
|
- smem_base,
|
|
|
- Int32(self.sh_red_off)
|
|
|
- + red_sh_delta * Int32(jj * 2)
|
|
|
- + red_sh_rd,
|
|
|
- )
|
|
|
- r0, r1, r2, r3 = ld_shared_v4_f32(rd_addr)
|
|
|
- acc[jj, 0] = acc[jj, 0] + r0
|
|
|
- acc[jj, 1] = acc[jj, 1] + r1
|
|
|
- acc[jj, 2] = acc[jj, 2] + r2
|
|
|
- acc[jj, 3] = acc[jj, 3] + r3
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _output_store_cursor(self, tid: Int32, output_n_tile: Int32):
|
|
|
- c_gl_stride = Int32(self.size_n // 8)
|
|
|
- c_sh_stride = Int32(2 * self.cta_n_blocks + 1)
|
|
|
- c_gl_wr_delta = c_gl_stride * Int32(self.cta_threads // (2 * self.cta_n_blocks))
|
|
|
- c_sh_rd_delta = c_sh_stride * Int32(
|
|
|
- self.cta_threads // (2 * self.cta_n_blocks)
|
|
|
- )
|
|
|
- c_gl_wr = (
|
|
|
- c_gl_stride * (tid // Int32(2 * self.cta_n_blocks))
|
|
|
- + (tid % Int32(2 * self.cta_n_blocks))
|
|
|
- + Int32(2 * self.cta_n_blocks) * output_n_tile
|
|
|
- )
|
|
|
- c_sh_rd = c_sh_stride * (tid // Int32(2 * self.cta_n_blocks)) + (
|
|
|
- tid % Int32(2 * self.cta_n_blocks)
|
|
|
- )
|
|
|
- return c_gl_stride, c_sh_stride, c_gl_wr_delta, c_sh_rd_delta, c_gl_wr, c_sh_rd
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _output_store_cursor_tail(self, tid: Int32, output_n_tile: Int32):
|
|
|
- c_gl_stride = Int32(self.size_n // 8)
|
|
|
- c_gl_stride_covered = Int32(self.covered_size_n // 8)
|
|
|
- c_sh_stride = Int32(2 * self.cta_n_blocks + 1)
|
|
|
- c_gl_wr_delta = c_gl_stride_covered * Int32(
|
|
|
- self.cta_threads // (2 * self.cta_n_blocks)
|
|
|
- )
|
|
|
- c_sh_rd_delta = c_sh_stride * Int32(self.cta_threads // (2 * self.cta_n_blocks))
|
|
|
- c_gl_wr = (
|
|
|
- c_gl_stride_covered * (tid // Int32(2 * self.cta_n_blocks))
|
|
|
- + (tid % Int32(2 * self.cta_n_blocks))
|
|
|
- + Int32(2 * self.cta_n_blocks) * output_n_tile
|
|
|
- )
|
|
|
- c_sh_rd = c_sh_stride * (tid // Int32(2 * self.cta_n_blocks)) + (
|
|
|
- tid % Int32(2 * self.cta_n_blocks)
|
|
|
- )
|
|
|
- return (
|
|
|
- c_gl_stride,
|
|
|
- c_gl_stride_covered,
|
|
|
- c_sh_stride,
|
|
|
- c_gl_wr_delta,
|
|
|
- c_sh_rd_delta,
|
|
|
- c_gl_wr,
|
|
|
- c_sh_rd,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _drain_output_smem(
|
|
|
- self,
|
|
|
- c_bf16_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- c_gl_stride: Int32,
|
|
|
- c_gl_wr: Int32,
|
|
|
- c_gl_wr_delta: Int32,
|
|
|
- c_sh_rd: Int32,
|
|
|
- c_sh_rd_delta: Int32,
|
|
|
- block_valid_rows: Int32,
|
|
|
- store_iters: cutlass.Constexpr[int],
|
|
|
- ):
|
|
|
- for _ in cutlass.range_constexpr(store_iters):
|
|
|
- row = c_gl_wr // c_gl_stride
|
|
|
- if row < block_valid_rows:
|
|
|
- route_index = ld_shared_i32_relaxed(
|
|
|
- smem_base + Int32(self.sh_route_off * 16) + row * Int32(4)
|
|
|
- )
|
|
|
- true_idx = route_index * c_gl_stride + (c_gl_wr % c_gl_stride)
|
|
|
- q0, q1, q2, q3 = ld_shared_v4_u32(
|
|
|
- self._int4_addr(smem_base, Int32(self.sh_red_off) + c_sh_rd)
|
|
|
- )
|
|
|
- if cutlass.const_expr(self.mul_topk_weights):
|
|
|
- scale_bf2 = ld_shared_u32(
|
|
|
- smem_base + Int32(self.sh_topk_off * 16) + row * Int32(4)
|
|
|
- )
|
|
|
- q0 = self._elem2_mul(q0, scale_bf2)
|
|
|
- q1 = self._elem2_mul(q1, scale_bf2)
|
|
|
- q2 = self._elem2_mul(q2, scale_bf2)
|
|
|
- q3 = self._elem2_mul(q3, scale_bf2)
|
|
|
- if cutlass.const_expr(self.epilogue_relu2):
|
|
|
- q0 = self._relu2_elem2(q0)
|
|
|
- q1 = self._relu2_elem2(q1)
|
|
|
- q2 = self._relu2_elem2(q2)
|
|
|
- q3 = self._relu2_elem2(q3)
|
|
|
- if cutlass.const_expr(self.fused_topk_sum):
|
|
|
- # Fold the per-route partials into the per-token output in
|
|
|
- # the epilogue (drops the separate top-k-sum launch). The
|
|
|
- # output must be zeroed before launch; each route slot maps
|
|
|
- # to token = route_index // top_k. bf16x2 add lands two
|
|
|
- # consecutive hidden lanes per word.
|
|
|
- token_idx = route_index // Int32(self.fused_sum_topk)
|
|
|
- out_idx = token_idx * c_gl_stride + (c_gl_wr % c_gl_stride)
|
|
|
- out_addr = get_ptr_as_int64(c_bf16_flat, out_idx * Int32(8))
|
|
|
- red_add_global_bf16x2(out_addr, q0)
|
|
|
- red_add_global_bf16x2(out_addr + Int64(4), q1)
|
|
|
- red_add_global_bf16x2(out_addr + Int64(8), q2)
|
|
|
- red_add_global_bf16x2(out_addr + Int64(12), q3)
|
|
|
- else:
|
|
|
- st_global_v4_u32(
|
|
|
- get_ptr_as_int64(c_bf16_flat, true_idx * Int32(8)),
|
|
|
- q0,
|
|
|
- q1,
|
|
|
- q2,
|
|
|
- q3,
|
|
|
- )
|
|
|
- c_gl_wr += c_gl_wr_delta
|
|
|
- c_sh_rd += c_sh_rd_delta
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _drain_output_smem_tail(
|
|
|
- self,
|
|
|
- c_bf16_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- c_gl_stride: Int32,
|
|
|
- c_gl_stride_covered: Int32,
|
|
|
- c_gl_wr: Int32,
|
|
|
- c_gl_wr_delta: Int32,
|
|
|
- c_sh_rd: Int32,
|
|
|
- c_sh_rd_delta: Int32,
|
|
|
- block_valid_rows: Int32,
|
|
|
- store_iters: cutlass.Constexpr[int],
|
|
|
- ):
|
|
|
- for _ in cutlass.range_constexpr(store_iters):
|
|
|
- row = c_gl_wr // c_gl_stride_covered
|
|
|
- col_word = c_gl_wr - row * c_gl_stride_covered
|
|
|
- if row < block_valid_rows and col_word < c_gl_stride:
|
|
|
- route_index = ld_shared_i32_relaxed(
|
|
|
- smem_base + Int32(self.sh_route_off * 16) + row * Int32(4)
|
|
|
- )
|
|
|
- true_idx = route_index * c_gl_stride + col_word
|
|
|
- q0, q1, q2, q3 = ld_shared_v4_u32(
|
|
|
- self._int4_addr(smem_base, Int32(self.sh_red_off) + c_sh_rd)
|
|
|
- )
|
|
|
- if cutlass.const_expr(self.mul_topk_weights):
|
|
|
- scale_bf2 = ld_shared_u32(
|
|
|
- smem_base + Int32(self.sh_topk_off * 16) + row * Int32(4)
|
|
|
- )
|
|
|
- q0 = self._elem2_mul(q0, scale_bf2)
|
|
|
- q1 = self._elem2_mul(q1, scale_bf2)
|
|
|
- q2 = self._elem2_mul(q2, scale_bf2)
|
|
|
- q3 = self._elem2_mul(q3, scale_bf2)
|
|
|
- if cutlass.const_expr(self.epilogue_relu2):
|
|
|
- q0 = self._relu2_elem2(q0)
|
|
|
- q1 = self._relu2_elem2(q1)
|
|
|
- q2 = self._relu2_elem2(q2)
|
|
|
- q3 = self._relu2_elem2(q3)
|
|
|
- if cutlass.const_expr(self.fused_topk_sum):
|
|
|
- token_idx = route_index // Int32(self.fused_sum_topk)
|
|
|
- out_idx = token_idx * c_gl_stride + col_word
|
|
|
- out_addr = get_ptr_as_int64(c_bf16_flat, out_idx * Int32(8))
|
|
|
- red_add_global_bf16x2(out_addr, q0)
|
|
|
- red_add_global_bf16x2(out_addr + Int64(4), q1)
|
|
|
- red_add_global_bf16x2(out_addr + Int64(8), q2)
|
|
|
- red_add_global_bf16x2(out_addr + Int64(12), q3)
|
|
|
- else:
|
|
|
- st_global_v4_u32(
|
|
|
- get_ptr_as_int64(c_bf16_flat, true_idx * Int32(8)),
|
|
|
- q0,
|
|
|
- q1,
|
|
|
- q2,
|
|
|
- q3,
|
|
|
- )
|
|
|
- c_gl_wr += c_gl_wr_delta
|
|
|
- c_sh_rd += c_sh_rd_delta
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _store_tile_m8(
|
|
|
- self,
|
|
|
- acc: cute.Tensor,
|
|
|
- c_bf16_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- block_valid_rows: Int32,
|
|
|
- global_scale_f32: cutlass.Float32,
|
|
|
- ):
|
|
|
- if cutlass.const_expr(self.has_n_tile_tail):
|
|
|
- (
|
|
|
- c_gl_stride,
|
|
|
- c_gl_stride_covered,
|
|
|
- c_sh_stride,
|
|
|
- c_gl_wr_delta,
|
|
|
- c_sh_rd_delta,
|
|
|
- c_gl_wr,
|
|
|
- c_sh_rd,
|
|
|
- ) = self._output_store_cursor_tail(tid, output_n_tile)
|
|
|
- else:
|
|
|
- (
|
|
|
- c_gl_stride,
|
|
|
- c_sh_stride,
|
|
|
- c_gl_wr_delta,
|
|
|
- c_sh_rd_delta,
|
|
|
- c_gl_wr,
|
|
|
- c_sh_rd,
|
|
|
- ) = self._output_store_cursor(tid, output_n_tile)
|
|
|
- c_gl_stride_covered = c_gl_stride
|
|
|
- c_sh_wr = (
|
|
|
- Int32(8) * c_sh_stride * (((tid & Int32(31)) % Int32(4)) * Int32(2))
|
|
|
- + (tid & Int32(31)) // Int32(4)
|
|
|
- + Int32(64) * (tid // Int32(32))
|
|
|
- )
|
|
|
-
|
|
|
- if tid // Int32(32) < Int32(self.tb_n_warps):
|
|
|
- write_scale = cutlass.Float32(1.0)
|
|
|
- if cutlass.const_expr(not self.mul_topk_weights):
|
|
|
- write_scale = global_scale_f32
|
|
|
- for jj in cutlass.range_constexpr(4):
|
|
|
- wr = c_sh_wr + Int32(16 * jj)
|
|
|
- self._st_shared_elem_from_f32(
|
|
|
- smem_base + Int32(self.sh_red_off * 16) + (wr * Int32(2)),
|
|
|
- acc[jj, 0] * write_scale,
|
|
|
- )
|
|
|
- self._st_shared_elem_from_f32(
|
|
|
- smem_base
|
|
|
- + Int32(self.sh_red_off * 16)
|
|
|
- + ((wr + Int32(8) * c_sh_stride) * Int32(2)),
|
|
|
- acc[jj, 1] * write_scale,
|
|
|
- )
|
|
|
- self._st_shared_elem_from_f32(
|
|
|
- smem_base
|
|
|
- + Int32(self.sh_red_off * 16)
|
|
|
- + ((wr + Int32(8)) * Int32(2)),
|
|
|
- acc[jj, 2] * write_scale,
|
|
|
- )
|
|
|
- self._st_shared_elem_from_f32(
|
|
|
- smem_base
|
|
|
- + Int32(self.sh_red_off * 16)
|
|
|
- + ((wr + Int32(8) + Int32(8) * c_sh_stride) * Int32(2)),
|
|
|
- acc[jj, 3] * write_scale,
|
|
|
- )
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- store_iters = _covering_count(
|
|
|
- 16, self.cta_threads // (2 * self.cta_n_blocks)
|
|
|
- )
|
|
|
- if cutlass.const_expr(self.has_n_tile_tail):
|
|
|
- self._drain_output_smem_tail(
|
|
|
- c_bf16_flat,
|
|
|
- smem_base,
|
|
|
- c_gl_stride,
|
|
|
- c_gl_stride_covered,
|
|
|
- c_gl_wr,
|
|
|
- c_gl_wr_delta,
|
|
|
- c_sh_rd,
|
|
|
- c_sh_rd_delta,
|
|
|
- block_valid_rows,
|
|
|
- store_iters,
|
|
|
- )
|
|
|
- else:
|
|
|
- self._drain_output_smem(
|
|
|
- c_bf16_flat,
|
|
|
- smem_base,
|
|
|
- c_gl_stride,
|
|
|
- c_gl_wr,
|
|
|
- c_gl_wr_delta,
|
|
|
- c_sh_rd,
|
|
|
- c_sh_rd_delta,
|
|
|
- block_valid_rows,
|
|
|
- store_iters,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _fold_cta_partials_large_m(
|
|
|
- self, acc: cute.Tensor, smem_base: Int32, tid: Int32
|
|
|
- ):
|
|
|
- red_off = self.cta_threads // self.b_sh_stride_threads // 2
|
|
|
- if cutlass.const_expr(red_off >= 1):
|
|
|
- red_idx, red_sh_stride, red_sh_delta, red_sh_rd = self._reduction_offsets(
|
|
|
- tid
|
|
|
- )
|
|
|
-
|
|
|
- for mb in cutlass.range_constexpr(self.cta_m_blocks):
|
|
|
- if cutlass.const_expr(red_off == 2):
|
|
|
- if Int32(2) <= red_idx and red_idx < Int32(4):
|
|
|
- for flat_j in cutlass.range_constexpr(8):
|
|
|
- jj = flat_j // 2
|
|
|
- half = flat_j % 2
|
|
|
- red_sh_wr = red_sh_delta * Int32(flat_j) + (
|
|
|
- red_sh_rd - red_sh_stride * Int32(2)
|
|
|
- )
|
|
|
- st_shared_v4_f32(
|
|
|
- self._int4_addr(
|
|
|
- smem_base, Int32(self.sh_red_off) + red_sh_wr
|
|
|
- ),
|
|
|
- acc[mb, jj, half, 0],
|
|
|
- acc[mb, jj, half, 1],
|
|
|
- acc[mb, jj, half, 2],
|
|
|
- acc[mb, jj, half, 3],
|
|
|
- )
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- if Int32(1) <= red_idx and red_idx < Int32(2):
|
|
|
- for flat_j in cutlass.range_constexpr(8):
|
|
|
- jj = flat_j // 2
|
|
|
- half = flat_j % 2
|
|
|
- red_sh_wr = red_sh_delta * Int32(flat_j) + (
|
|
|
- red_sh_rd - red_sh_stride
|
|
|
- )
|
|
|
- if cutlass.const_expr(red_off > 1):
|
|
|
- rd_addr = self._int4_addr(
|
|
|
- smem_base,
|
|
|
- Int32(self.sh_red_off)
|
|
|
- + red_sh_delta * Int32(flat_j)
|
|
|
- + red_sh_rd,
|
|
|
- )
|
|
|
- wr_addr = self._int4_addr(
|
|
|
- smem_base,
|
|
|
- Int32(self.sh_red_off) + red_sh_wr,
|
|
|
- )
|
|
|
- r0, r1, r2, r3 = ld_shared_v4_f32(rd_addr)
|
|
|
- w0, w1, w2, w3 = ld_shared_v4_f32(wr_addr)
|
|
|
- acc[mb, jj, half, 0] = acc[mb, jj, half, 0] + r0 + w0
|
|
|
- acc[mb, jj, half, 1] = acc[mb, jj, half, 1] + r1 + w1
|
|
|
- acc[mb, jj, half, 2] = acc[mb, jj, half, 2] + r2 + w2
|
|
|
- acc[mb, jj, half, 3] = acc[mb, jj, half, 3] + r3 + w3
|
|
|
- st_shared_v4_f32(
|
|
|
- self._int4_addr(
|
|
|
- smem_base, Int32(self.sh_red_off) + red_sh_wr
|
|
|
- ),
|
|
|
- acc[mb, jj, half, 0],
|
|
|
- acc[mb, jj, half, 1],
|
|
|
- acc[mb, jj, half, 2],
|
|
|
- acc[mb, jj, half, 3],
|
|
|
- )
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- if red_idx == Int32(0):
|
|
|
- for flat_j in cutlass.range_constexpr(8):
|
|
|
- jj = flat_j // 2
|
|
|
- half = flat_j % 2
|
|
|
- rd_addr = self._int4_addr(
|
|
|
- smem_base,
|
|
|
- Int32(self.sh_red_off)
|
|
|
- + red_sh_delta * Int32(flat_j)
|
|
|
- + red_sh_rd,
|
|
|
- )
|
|
|
- r0, r1, r2, r3 = ld_shared_v4_f32(rd_addr)
|
|
|
- acc[mb, jj, half, 0] = acc[mb, jj, half, 0] + r0
|
|
|
- acc[mb, jj, half, 1] = acc[mb, jj, half, 1] + r1
|
|
|
- acc[mb, jj, half, 2] = acc[mb, jj, half, 2] + r2
|
|
|
- acc[mb, jj, half, 3] = acc[mb, jj, half, 3] + r3
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _write_bf16x2_shared(
|
|
|
- self,
|
|
|
- smem_base: Int32,
|
|
|
- half2_idx: Int32,
|
|
|
- c0: cutlass.Float32,
|
|
|
- c1: cutlass.Float32,
|
|
|
- write_scale: cutlass.Float32,
|
|
|
- ):
|
|
|
- packed = self._pack_f32x2_to_elem2(c0 * write_scale, c1 * write_scale)
|
|
|
- st_shared_u32(
|
|
|
- smem_base + Int32(self.sh_red_off * 16) + half2_idx * Int32(4),
|
|
|
- packed,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _store_tile_large_m(
|
|
|
- self,
|
|
|
- acc: cute.Tensor,
|
|
|
- c_bf16_flat: cute.Tensor,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- output_n_tile: Int32,
|
|
|
- block_valid_rows: Int32,
|
|
|
- global_scale_f32: cutlass.Float32,
|
|
|
- ):
|
|
|
- if cutlass.const_expr(self.has_n_tile_tail):
|
|
|
- (
|
|
|
- c_gl_stride,
|
|
|
- c_gl_stride_covered,
|
|
|
- c_sh_stride,
|
|
|
- c_gl_wr_delta,
|
|
|
- c_sh_rd_delta,
|
|
|
- c_gl_wr,
|
|
|
- c_sh_rd,
|
|
|
- ) = self._output_store_cursor_tail(tid, output_n_tile)
|
|
|
- else:
|
|
|
- (
|
|
|
- c_gl_stride,
|
|
|
- c_sh_stride,
|
|
|
- c_gl_wr_delta,
|
|
|
- c_sh_rd_delta,
|
|
|
- c_gl_wr,
|
|
|
- c_sh_rd,
|
|
|
- ) = self._output_store_cursor(tid, output_n_tile)
|
|
|
- c_gl_stride_covered = c_gl_stride
|
|
|
- c_sh_wr = (
|
|
|
- Int32(4) * c_sh_stride * ((tid & Int32(31)) // Int32(4))
|
|
|
- + (tid & Int32(31)) % Int32(4)
|
|
|
- + Int32(32) * (tid // Int32(32))
|
|
|
- )
|
|
|
-
|
|
|
- if tid // Int32(32) < Int32(self.tb_n_warps):
|
|
|
- write_scale = cutlass.Float32(1.0)
|
|
|
- if cutlass.const_expr(not self.mul_topk_weights):
|
|
|
- write_scale = global_scale_f32
|
|
|
- for mb in cutlass.range_constexpr(self.cta_m_blocks):
|
|
|
- for jj in cutlass.range_constexpr(4):
|
|
|
- wr = c_sh_wr + Int32(8 * jj)
|
|
|
- self._write_bf16x2_shared(
|
|
|
- smem_base,
|
|
|
- wr,
|
|
|
- acc[mb, jj, 0, 0],
|
|
|
- acc[mb, jj, 0, 1],
|
|
|
- write_scale,
|
|
|
- )
|
|
|
- self._write_bf16x2_shared(
|
|
|
- smem_base,
|
|
|
- wr + (Int32(4) * c_sh_stride) * Int32(8) + Int32(0),
|
|
|
- acc[mb, jj, 0, 2],
|
|
|
- acc[mb, jj, 0, 3],
|
|
|
- write_scale,
|
|
|
- )
|
|
|
- self._write_bf16x2_shared(
|
|
|
- smem_base,
|
|
|
- wr + Int32(4),
|
|
|
- acc[mb, jj, 1, 0],
|
|
|
- acc[mb, jj, 1, 1],
|
|
|
- write_scale,
|
|
|
- )
|
|
|
- self._write_bf16x2_shared(
|
|
|
- smem_base,
|
|
|
- wr + (Int32(4) * c_sh_stride) * Int32(8) + Int32(4),
|
|
|
- acc[mb, jj, 1, 2],
|
|
|
- acc[mb, jj, 1, 3],
|
|
|
- write_scale,
|
|
|
- )
|
|
|
- c_sh_wr += Int32(16 * (4 * (2 * self.cta_n_blocks + 1)))
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- store_iters = _covering_count(
|
|
|
- 16 * self.cta_m_blocks,
|
|
|
- self.cta_threads // (2 * self.cta_n_blocks),
|
|
|
- )
|
|
|
- if cutlass.const_expr(self.has_n_tile_tail):
|
|
|
- self._drain_output_smem_tail(
|
|
|
- c_bf16_flat,
|
|
|
- smem_base,
|
|
|
- c_gl_stride,
|
|
|
- c_gl_stride_covered,
|
|
|
- c_gl_wr,
|
|
|
- c_gl_wr_delta,
|
|
|
- c_sh_rd,
|
|
|
- c_sh_rd_delta,
|
|
|
- block_valid_rows,
|
|
|
- store_iters,
|
|
|
- )
|
|
|
- else:
|
|
|
- self._drain_output_smem(
|
|
|
- c_bf16_flat,
|
|
|
- smem_base,
|
|
|
- c_gl_stride,
|
|
|
- c_gl_wr,
|
|
|
- c_gl_wr_delta,
|
|
|
- c_sh_rd,
|
|
|
- c_sh_rd_delta,
|
|
|
- block_valid_rows,
|
|
|
- store_iters,
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-class W4A16FusedMoeKernel:
|
|
|
- def __init__(
|
|
|
- self,
|
|
|
- *,
|
|
|
- size_m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- num_experts: int,
|
|
|
- top_k: int,
|
|
|
- activation: str,
|
|
|
- apply_router_weight_on_input: bool,
|
|
|
- zero_fc2_output: bool,
|
|
|
- fc1_tile_n: int,
|
|
|
- fc1_tile_k: int,
|
|
|
- fc2_tile_n: int,
|
|
|
- fc2_tile_k: int,
|
|
|
- moe_block_size: int,
|
|
|
- max_m_blocks: int,
|
|
|
- element_dtype: str = "bf16",
|
|
|
- fast_math: bool = True,
|
|
|
- swiglu_limit: float | None = None,
|
|
|
- swiglu_alpha: float | None = None,
|
|
|
- swiglu_beta: float | None = None,
|
|
|
- weight_layout: str = "packed",
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
- w13_layout: str = "w13",
|
|
|
- direct_topk_routes: bool = False,
|
|
|
- tc_decode_fused_sum: bool = False,
|
|
|
- tc_zero_output: bool = True,
|
|
|
- collect_activation_amax: bool = False,
|
|
|
- ):
|
|
|
- activation = normalize_moe_activation(activation)
|
|
|
- is_gated = validate_activation(activation)
|
|
|
- swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params(
|
|
|
- activation,
|
|
|
- swiglu_limit,
|
|
|
- swiglu_alpha,
|
|
|
- swiglu_beta,
|
|
|
- )
|
|
|
- if weight_layout not in _WEIGHT_LAYOUTS:
|
|
|
- raise ValueError(f"unsupported W4A16 weight_layout {weight_layout!r}")
|
|
|
- scale_format = _normalize_scale_format(scale_format)
|
|
|
- if weight_layout == "modelopt":
|
|
|
- if w13_layout not in _MODEL_OPT_W13_LAYOUTS:
|
|
|
- raise ValueError(f"unsupported W4A16 w13_layout {w13_layout!r}")
|
|
|
- else:
|
|
|
- w13_layout = "packed"
|
|
|
- self.tc_decode_fused_sum = bool(tc_decode_fused_sum)
|
|
|
- # When two TC-decode launches share one pre-zeroed output (the NF3 hybrid
|
|
|
- # runs an NVFP4 launch then an NF3 launch into the same tensor), only the
|
|
|
- # first must zero it. Default True preserves single-launch behavior.
|
|
|
- self.tc_zero_output = bool(tc_zero_output)
|
|
|
- self.collect_activation_amax = bool(collect_activation_amax)
|
|
|
- if self.collect_activation_amax and bool(direct_topk_routes):
|
|
|
- raise ValueError("activation amax collection requires route-packed W4A16")
|
|
|
- if self.collect_activation_amax and self.tc_decode_fused_sum:
|
|
|
- raise ValueError("activation amax collection is incompatible with TC-decode")
|
|
|
- if self.tc_decode_fused_sum and not bool(direct_topk_routes):
|
|
|
- raise ValueError("tc_decode_fused_sum requires direct_topk_routes")
|
|
|
- if self.tc_decode_fused_sum and element_dtype != "bf16":
|
|
|
- raise ValueError("tc_decode_fused_sum currently requires bf16 activations")
|
|
|
- fc1_cols = int(intermediate_size) * (2 if is_gated else 1)
|
|
|
- routed_rows = int(size_m) * int(top_k)
|
|
|
- self.size_m = int(size_m)
|
|
|
- self.hidden_size = int(hidden_size)
|
|
|
- self.intermediate_size = int(intermediate_size)
|
|
|
- self.fc1_cols = int(fc1_cols)
|
|
|
- self.num_experts = int(num_experts)
|
|
|
- self.top_k = int(top_k)
|
|
|
- self.moe_block_size = int(moe_block_size)
|
|
|
- self.activation = activation
|
|
|
- self.activation_is_gated = is_gated
|
|
|
- self.activation_is_swigluoai = activation == SWIGLUOAI_UNINTERLEAVE
|
|
|
- self.has_swiglu_limit = swiglu_limit is not None
|
|
|
- self.swiglu_limit = 0.0 if swiglu_limit is None else swiglu_limit
|
|
|
- self.swiglu_alpha = float(swiglu_alpha)
|
|
|
- self.swiglu_beta = float(swiglu_beta)
|
|
|
- self.weight_layout = weight_layout
|
|
|
- self.scale_format = scale_format
|
|
|
- self.w13_layout = w13_layout
|
|
|
- self.apply_router_weight_on_input = bool(apply_router_weight_on_input)
|
|
|
- self.zero_fc2_output = bool(zero_fc2_output)
|
|
|
- self.element_dtype = element_dtype
|
|
|
- self.is_fp16 = element_dtype == "fp16"
|
|
|
- self.fast_math = bool(fast_math)
|
|
|
- self.direct_topk_routes = bool(direct_topk_routes)
|
|
|
- fc1_source_n_rotation = (
|
|
|
- int(intermediate_size)
|
|
|
- if (
|
|
|
- weight_layout == "modelopt"
|
|
|
- and w13_layout == "w13"
|
|
|
- and is_gated
|
|
|
- )
|
|
|
- else 0
|
|
|
- )
|
|
|
- self.fc1 = W4A16GemmKernel(
|
|
|
- size_m=size_m,
|
|
|
- size_n=fc1_cols,
|
|
|
- size_k=hidden_size,
|
|
|
- num_experts=num_experts,
|
|
|
- top_k=top_k,
|
|
|
- mul_topk_weights=bool(apply_router_weight_on_input),
|
|
|
- tile_n=fc1_tile_n,
|
|
|
- tile_k=fc1_tile_k,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- max_m_blocks=max_m_blocks,
|
|
|
- element_dtype=element_dtype,
|
|
|
- epilogue_activation=None if is_gated else "relu2",
|
|
|
- weight_layout=weight_layout,
|
|
|
- scale_format=scale_format,
|
|
|
- w13_layout=w13_layout,
|
|
|
- source_n_rotation=fc1_source_n_rotation,
|
|
|
- single_token_route_fast_path=size_m == 1
|
|
|
- and not self.direct_topk_routes,
|
|
|
- direct_topk_routes=self.direct_topk_routes,
|
|
|
- )
|
|
|
- self.fc2 = W4A16GemmKernel(
|
|
|
- size_m=routed_rows,
|
|
|
- size_n=hidden_size,
|
|
|
- size_k=intermediate_size,
|
|
|
- num_experts=num_experts,
|
|
|
- top_k=1,
|
|
|
- mul_topk_weights=not bool(apply_router_weight_on_input),
|
|
|
- tile_n=fc2_tile_n,
|
|
|
- tile_k=fc2_tile_k,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- max_m_blocks=max_m_blocks,
|
|
|
- element_dtype=element_dtype,
|
|
|
- weight_layout=weight_layout,
|
|
|
- scale_format=scale_format,
|
|
|
- w13_layout=w13_layout,
|
|
|
- single_token_route_fast_path=size_m == 1
|
|
|
- and not self.direct_topk_routes,
|
|
|
- direct_topk_routes=self.direct_topk_routes,
|
|
|
- fused_topk_sum=self.tc_decode_fused_sum,
|
|
|
- fused_sum_topk=int(top_k),
|
|
|
- )
|
|
|
- self.cta_threads = max(self.fc1.cta_threads, self.fc2.cta_threads)
|
|
|
- if self.fc1.cta_threads != self.fc2.cta_threads:
|
|
|
- raise ValueError(
|
|
|
- "fused W4A16 kernel expects matching FC1/FC2 thread counts"
|
|
|
- )
|
|
|
- self.sms = self.fc1.sms
|
|
|
- self.blocks_per_sm = min(self.fc1.blocks_per_sm, self.fc2.blocks_per_sm)
|
|
|
- self.shared_words = max(self.fc1.shared_words, self.fc2.shared_words)
|
|
|
- self.barrier_count_off = self.sms * 4
|
|
|
- self.barrier_sense_off = self.sms * 4 + 1
|
|
|
-
|
|
|
- @property
|
|
|
- def __cache_key__(self) -> tuple[object, ...]:
|
|
|
- return (
|
|
|
- self.hidden_size,
|
|
|
- self.intermediate_size,
|
|
|
- self.fc1_cols,
|
|
|
- self.num_experts,
|
|
|
- self.top_k,
|
|
|
- self.activation,
|
|
|
- self.activation_is_gated,
|
|
|
- self.activation_is_swigluoai,
|
|
|
- self.has_swiglu_limit,
|
|
|
- self.swiglu_limit,
|
|
|
- self.swiglu_alpha,
|
|
|
- self.swiglu_beta,
|
|
|
- self.weight_layout,
|
|
|
- self.scale_format,
|
|
|
- self.apply_router_weight_on_input,
|
|
|
- self.zero_fc2_output,
|
|
|
- self.element_dtype,
|
|
|
- self.fast_math,
|
|
|
- self.direct_topk_routes,
|
|
|
- self.tc_zero_output,
|
|
|
- self.collect_activation_amax,
|
|
|
- self.fc1.__cache_key__,
|
|
|
- self.fc2.__cache_key__,
|
|
|
- self.cta_threads,
|
|
|
- self.sms,
|
|
|
- self.shared_words,
|
|
|
- self.blocks_per_sm,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _cast_elem(self, x: cutlass.Float32):
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return cutlass.Float16(x)
|
|
|
- return cutlass.BFloat16(x)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _clamp_swiglu_inputs(
|
|
|
- self,
|
|
|
- gate: cutlass.Float32,
|
|
|
- up: cutlass.Float32,
|
|
|
- ):
|
|
|
- if cutlass.const_expr(self.has_swiglu_limit):
|
|
|
- limit = cutlass.Float32(self.swiglu_limit)
|
|
|
- neg_limit = cutlass.Float32(-self.swiglu_limit)
|
|
|
- if gate > limit:
|
|
|
- gate = limit
|
|
|
- if up > limit:
|
|
|
- up = limit
|
|
|
- if up < neg_limit:
|
|
|
- up = neg_limit
|
|
|
- return gate, up
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def __call__(
|
|
|
- self,
|
|
|
- a_bf16_ptr: cute.Pointer,
|
|
|
- w13_i32_flat: cute.Tensor,
|
|
|
- w2_i32_flat: cute.Tensor,
|
|
|
- fc1_bf16_flat: cute.Tensor,
|
|
|
- activated_bf16_flat: cute.Tensor,
|
|
|
- fc2_bf16_flat: cute.Tensor,
|
|
|
- w13_scales_i32_flat: cute.Tensor,
|
|
|
- w2_scales_i32_flat: cute.Tensor,
|
|
|
- w13_global_scale: cute.Tensor,
|
|
|
- w2_global_scale: cute.Tensor,
|
|
|
- packed_route_indices: cute.Tensor,
|
|
|
- block_expert_ids: cute.Tensor,
|
|
|
- packed_route_count: cute.Tensor,
|
|
|
- activation_amax_flat: cute.Tensor,
|
|
|
- layer_idx: cutlass.Int32,
|
|
|
- topk_weights_ptr: cute.Pointer,
|
|
|
- fc1_c_tmp_f32_flat: cute.Tensor,
|
|
|
- fc2_c_tmp_f32_flat: cute.Tensor,
|
|
|
- locks_i32_flat: cute.Tensor,
|
|
|
- active_m: cutlass.Int32,
|
|
|
- grid_x: cutlass.Int32,
|
|
|
- stream: cuda.CUstream,
|
|
|
- ):
|
|
|
- a_bf16_flat = cute.make_tensor(
|
|
|
- a_bf16_ptr,
|
|
|
- layout=cute.make_layout(
|
|
|
- (active_m * Int32(self.hidden_size),), stride=(1,)
|
|
|
- ),
|
|
|
- )
|
|
|
- topk_weights_flat = cute.make_tensor(
|
|
|
- topk_weights_ptr,
|
|
|
- layout=cute.make_layout((active_m * Int32(self.top_k),), stride=(1,)),
|
|
|
- )
|
|
|
- grid = (grid_x, 1, 1)
|
|
|
- self.kernel(
|
|
|
- a_bf16_flat,
|
|
|
- w13_i32_flat,
|
|
|
- w2_i32_flat,
|
|
|
- fc1_bf16_flat,
|
|
|
- activated_bf16_flat,
|
|
|
- fc2_bf16_flat,
|
|
|
- w13_scales_i32_flat,
|
|
|
- w2_scales_i32_flat,
|
|
|
- w13_global_scale,
|
|
|
- w2_global_scale,
|
|
|
- packed_route_indices,
|
|
|
- block_expert_ids,
|
|
|
- packed_route_count,
|
|
|
- activation_amax_flat,
|
|
|
- layer_idx,
|
|
|
- topk_weights_flat,
|
|
|
- fc1_c_tmp_f32_flat,
|
|
|
- fc2_c_tmp_f32_flat,
|
|
|
- locks_i32_flat,
|
|
|
- active_m,
|
|
|
- ).launch(
|
|
|
- grid=grid,
|
|
|
- block=[self.cta_threads, 1, 1],
|
|
|
- min_blocks_per_mp=self.blocks_per_sm,
|
|
|
- stream=stream,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.kernel
|
|
|
- def kernel(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- w13_i32_flat: cute.Tensor,
|
|
|
- w2_i32_flat: cute.Tensor,
|
|
|
- fc1_bf16_flat: cute.Tensor,
|
|
|
- activated_bf16_flat: cute.Tensor,
|
|
|
- fc2_bf16_flat: cute.Tensor,
|
|
|
- w13_scales_i32_flat: cute.Tensor,
|
|
|
- w2_scales_i32_flat: cute.Tensor,
|
|
|
- w13_global_scale: cute.Tensor,
|
|
|
- w2_global_scale: cute.Tensor,
|
|
|
- packed_route_indices: cute.Tensor,
|
|
|
- block_expert_ids: cute.Tensor,
|
|
|
- packed_route_count: cute.Tensor,
|
|
|
- activation_amax_flat: cute.Tensor,
|
|
|
- layer_idx: cutlass.Int32,
|
|
|
- topk_weights_flat: cute.Tensor,
|
|
|
- fc1_c_tmp_f32_flat: cute.Tensor,
|
|
|
- fc2_c_tmp_f32_flat: cute.Tensor,
|
|
|
- locks_i32_flat: cute.Tensor,
|
|
|
- active_m: cutlass.Int32,
|
|
|
- ):
|
|
|
- tidx, _, _ = cute.arch.thread_idx()
|
|
|
- bidx, _, _ = cute.arch.block_idx()
|
|
|
- grid_x_raw, _, _ = cute.arch.grid_dim()
|
|
|
- tid = Int32(tidx)
|
|
|
- cta = Int32(bidx)
|
|
|
- grid_x = Int32(grid_x_raw)
|
|
|
-
|
|
|
- smem = cutlass.utils.SmemAllocator()
|
|
|
-
|
|
|
- @cute.struct
|
|
|
- class Storage:
|
|
|
- words: cute.struct.Align[
|
|
|
- cute.struct.MemRange[cutlass.Uint32, self.shared_words],
|
|
|
- 1024,
|
|
|
- ]
|
|
|
-
|
|
|
- storage = smem.allocate(Storage)
|
|
|
- smem_base = shared_ptr_to_u32(storage.words.data_ptr())
|
|
|
-
|
|
|
- if cutlass.const_expr(self.tc_decode_fused_sum):
|
|
|
- # The TC-decode FC2 epilogue atomically accumulates per-route
|
|
|
- # partials directly into the per-token output, so the output must be
|
|
|
- # pre-zeroed. Previously this was a SEPARATE host-side output.zero_()
|
|
|
- # kernel launch on the latency-bound decode critical path (an extra
|
|
|
- # launch + its grid-fill memset before the fused kernel even starts).
|
|
|
- # Fold it into the fused kernel prologue here: every CTA zeroes a
|
|
|
- # grid-strided slice of the output BEFORE FC1, and the EXISTING
|
|
|
- # post-FC1 grid barrier (already required to order FC1 writes before
|
|
|
- # the activation/FC2 read) makes all zero stores globally visible
|
|
|
- # before the first FC2 atomic -- so no extra barrier is added. The
|
|
|
- # tiny m*hidden bf16 memset (decode: <=4*4096 elems) is dwarfed by
|
|
|
- # FC1's whole-K FP4-weight stream, but we delete one whole kernel
|
|
|
- # launch from the per-decode chain. The TC-decode output is per-token
|
|
|
- # (top_k routes atomically summed into the SAME token row), so the
|
|
|
- # zero span is active_m*hidden_size -- NOT the per-route
|
|
|
- # active_m*top_k*hidden_size of _zero_fc2_output.
|
|
|
- # tc_zero_output=False skips the zero (a paired earlier launch has
|
|
|
- # already zeroed the shared output); the grid barrier below is
|
|
|
- # unconditional so ordering is preserved either way.
|
|
|
- if cutlass.const_expr(self.tc_zero_output):
|
|
|
- zidx = cta * Int32(self.cta_threads) + tid
|
|
|
- zstride = grid_x * Int32(self.cta_threads)
|
|
|
- ztotal = active_m * Int32(self.hidden_size)
|
|
|
- zzero = self._cast_elem(cutlass.Float32(0.0))
|
|
|
- while zidx < ztotal:
|
|
|
- fc2_bf16_flat[zidx] = zzero
|
|
|
- zidx += zstride
|
|
|
-
|
|
|
- if cutlass.const_expr(self.activation_is_gated):
|
|
|
- self.fc1._run_persistent_gemm(
|
|
|
- a_bf16_flat,
|
|
|
- w13_i32_flat,
|
|
|
- fc1_bf16_flat,
|
|
|
- w13_scales_i32_flat,
|
|
|
- w13_global_scale,
|
|
|
- packed_route_indices,
|
|
|
- block_expert_ids,
|
|
|
- packed_route_count,
|
|
|
- topk_weights_flat,
|
|
|
- fc1_c_tmp_f32_flat,
|
|
|
- locks_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- cta,
|
|
|
- grid_x,
|
|
|
- active_m,
|
|
|
- )
|
|
|
- self._grid_barrier(locks_i32_flat, tid, grid_x)
|
|
|
- self._run_activation(
|
|
|
- fc1_bf16_flat,
|
|
|
- activated_bf16_flat,
|
|
|
- tid,
|
|
|
- cta,
|
|
|
- grid_x,
|
|
|
- active_m,
|
|
|
- )
|
|
|
- else:
|
|
|
- self.fc1._run_persistent_gemm(
|
|
|
- a_bf16_flat,
|
|
|
- w13_i32_flat,
|
|
|
- activated_bf16_flat,
|
|
|
- w13_scales_i32_flat,
|
|
|
- w13_global_scale,
|
|
|
- packed_route_indices,
|
|
|
- block_expert_ids,
|
|
|
- packed_route_count,
|
|
|
- topk_weights_flat,
|
|
|
- fc1_c_tmp_f32_flat,
|
|
|
- locks_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- cta,
|
|
|
- grid_x,
|
|
|
- active_m,
|
|
|
- )
|
|
|
- self._grid_barrier(locks_i32_flat, tid, grid_x)
|
|
|
- if cutlass.const_expr(self.collect_activation_amax):
|
|
|
- self._collect_activation_amax_epilogue(
|
|
|
- a_bf16_flat,
|
|
|
- activated_bf16_flat,
|
|
|
- packed_route_indices,
|
|
|
- block_expert_ids,
|
|
|
- packed_route_count,
|
|
|
- activation_amax_flat,
|
|
|
- layer_idx,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- cta,
|
|
|
- grid_x,
|
|
|
- active_m,
|
|
|
- )
|
|
|
- if cutlass.const_expr(self.zero_fc2_output):
|
|
|
- self._zero_fc2_output(fc2_bf16_flat, tid, cta, grid_x, active_m)
|
|
|
- self._grid_barrier(locks_i32_flat, tid, grid_x)
|
|
|
- self.fc2._run_persistent_gemm(
|
|
|
- activated_bf16_flat,
|
|
|
- w2_i32_flat,
|
|
|
- fc2_bf16_flat,
|
|
|
- w2_scales_i32_flat,
|
|
|
- w2_global_scale,
|
|
|
- packed_route_indices,
|
|
|
- block_expert_ids,
|
|
|
- packed_route_count,
|
|
|
- topk_weights_flat,
|
|
|
- fc2_c_tmp_f32_flat,
|
|
|
- locks_i32_flat,
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- cta,
|
|
|
- grid_x,
|
|
|
- active_m * Int32(self.top_k),
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _grid_barrier(
|
|
|
- self,
|
|
|
- locks_i32_flat: cute.Tensor,
|
|
|
- tid: Int32,
|
|
|
- grid_x: Int32,
|
|
|
- ):
|
|
|
- cute.arch.sync_threads()
|
|
|
- if tid == Int32(0):
|
|
|
- count_addr = get_ptr_as_int64(locks_i32_flat, Int32(self.barrier_count_off))
|
|
|
- sense_addr = get_ptr_as_int64(locks_i32_flat, Int32(self.barrier_sense_off))
|
|
|
- old_sense = ld_global_acquire_i32(sense_addr)
|
|
|
- old_count = atomic_add_global_i32(count_addr, Int32(1))
|
|
|
- if old_count == grid_x - Int32(1):
|
|
|
- st_global_i32(count_addr, Int32(0))
|
|
|
- threadfence()
|
|
|
- red_add_global_release_i32(sense_addr, Int32(1))
|
|
|
- else:
|
|
|
- sense = old_sense
|
|
|
- while sense == old_sense:
|
|
|
- sense = ld_global_acquire_i32(sense_addr)
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _zero_fc2_output(
|
|
|
- self,
|
|
|
- fc2_bf16_flat: cute.Tensor,
|
|
|
- tid: Int32,
|
|
|
- cta: Int32,
|
|
|
- grid_x: Int32,
|
|
|
- active_m: cutlass.Int32,
|
|
|
- ):
|
|
|
- idx = cta * Int32(self.cta_threads) + tid
|
|
|
- stride = grid_x * Int32(self.cta_threads)
|
|
|
- total = active_m * Int32(self.top_k * self.hidden_size)
|
|
|
- zero = self._cast_elem(cutlass.Float32(0.0))
|
|
|
- while idx < total:
|
|
|
- fc2_bf16_flat[idx] = zero
|
|
|
- idx += stride
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _reduce_and_red_activation_amax(
|
|
|
- self,
|
|
|
- local_max: cutlass.Float32,
|
|
|
- activation_amax_flat: cute.Tensor,
|
|
|
- layer_idx: cutlass.Int32,
|
|
|
- expert_idx: Int32,
|
|
|
- slot: Int32,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- ):
|
|
|
- lane = tid & Int32(31)
|
|
|
- warp_idx = tid // Int32(32)
|
|
|
- warp_amax = warp_reduce(local_max, fmax_f32)
|
|
|
- if lane == Int32(0):
|
|
|
- st_shared_f32(smem_base + warp_idx * Int32(4), warp_amax)
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- if warp_idx == Int32(0):
|
|
|
- block_amax = cutlass.Float32(0.0)
|
|
|
- if lane < Int32(self.cta_threads // 32):
|
|
|
- block_amax = ld_shared_f32(smem_base + lane * Int32(4))
|
|
|
- block_amax = warp_reduce(block_amax, fmax_f32)
|
|
|
- if lane == Int32(0) and block_amax > cutlass.Float32(0.0):
|
|
|
- out_idx = (
|
|
|
- (layer_idx * Int32(self.num_experts) + expert_idx) * Int32(2)
|
|
|
- + slot
|
|
|
- )
|
|
|
- red_max_global_f32_nonnegative(
|
|
|
- get_ptr_as_int64(activation_amax_flat, out_idx),
|
|
|
- block_amax,
|
|
|
- )
|
|
|
- cute.arch.sync_threads()
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _collect_activation_amax_epilogue(
|
|
|
- self,
|
|
|
- a_bf16_flat: cute.Tensor,
|
|
|
- activated_bf16_flat: cute.Tensor,
|
|
|
- packed_route_indices: cute.Tensor,
|
|
|
- block_expert_ids: cute.Tensor,
|
|
|
- packed_route_count: cute.Tensor,
|
|
|
- activation_amax_flat: cute.Tensor,
|
|
|
- layer_idx: cutlass.Int32,
|
|
|
- smem_base: Int32,
|
|
|
- tid: Int32,
|
|
|
- cta: Int32,
|
|
|
- grid_x: Int32,
|
|
|
- active_m: cutlass.Int32,
|
|
|
- ):
|
|
|
- live_routes = active_m * Int32(self.top_k)
|
|
|
- route_count = packed_route_count[Int32(0)].to(Int32)
|
|
|
- route_blocks = (
|
|
|
- route_count + Int32(self.moe_block_size) - Int32(1)
|
|
|
- ) // Int32(self.moe_block_size)
|
|
|
- expert_idx = cta
|
|
|
- while expert_idx < Int32(self.num_experts):
|
|
|
- local_fc1 = cutlass.Float32(0.0)
|
|
|
- local_fc2 = cutlass.Float32(0.0)
|
|
|
- has_route_block = Int32(0)
|
|
|
- block_idx = Int32(0)
|
|
|
- while block_idx < route_blocks:
|
|
|
- block_expert = block_expert_ids[block_idx].to(Int32)
|
|
|
- if block_expert == expert_idx:
|
|
|
- has_route_block = Int32(1)
|
|
|
- route_pos = block_idx * Int32(self.moe_block_size)
|
|
|
- route_stop = route_pos + Int32(self.moe_block_size)
|
|
|
- if route_stop > route_count:
|
|
|
- route_stop = route_count
|
|
|
- while route_pos < route_stop:
|
|
|
- route_idx = packed_route_indices[route_pos].to(Int32)
|
|
|
- if route_idx >= Int32(0) and route_idx < live_routes:
|
|
|
- token_idx = route_idx // Int32(self.top_k)
|
|
|
- fc1_col = tid
|
|
|
- while fc1_col < Int32(self.hidden_size):
|
|
|
- v1 = a_bf16_flat[
|
|
|
- token_idx * Int32(self.hidden_size) + fc1_col
|
|
|
- ].to(cutlass.Float32)
|
|
|
- local_fc1 = fmax_f32(local_fc1, fabs_f32(v1))
|
|
|
- fc1_col += Int32(self.cta_threads)
|
|
|
-
|
|
|
- fc2_col = tid
|
|
|
- while fc2_col < Int32(self.intermediate_size):
|
|
|
- v2 = activated_bf16_flat[
|
|
|
- route_idx * Int32(self.intermediate_size) + fc2_col
|
|
|
- ].to(cutlass.Float32)
|
|
|
- local_fc2 = fmax_f32(local_fc2, fabs_f32(v2))
|
|
|
- fc2_col += Int32(self.cta_threads)
|
|
|
- route_pos += Int32(1)
|
|
|
- block_idx += Int32(1)
|
|
|
-
|
|
|
- if has_route_block != Int32(0):
|
|
|
- self._reduce_and_red_activation_amax(
|
|
|
- local_fc1,
|
|
|
- activation_amax_flat,
|
|
|
- layer_idx,
|
|
|
- expert_idx,
|
|
|
- Int32(0),
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- )
|
|
|
- self._reduce_and_red_activation_amax(
|
|
|
- local_fc2,
|
|
|
- activation_amax_flat,
|
|
|
- layer_idx,
|
|
|
- expert_idx,
|
|
|
- Int32(1),
|
|
|
- smem_base,
|
|
|
- tid,
|
|
|
- )
|
|
|
- expert_idx += grid_x
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _run_activation(
|
|
|
- self,
|
|
|
- fc1_bf16_flat: cute.Tensor,
|
|
|
- activated_bf16_flat: cute.Tensor,
|
|
|
- tid: Int32,
|
|
|
- cta: Int32,
|
|
|
- grid_x: Int32,
|
|
|
- active_m: cutlass.Int32,
|
|
|
- ):
|
|
|
- idx = cta * Int32(self.cta_threads) + tid
|
|
|
- stride = grid_x * Int32(self.cta_threads)
|
|
|
- total = active_m * Int32(self.top_k * self.intermediate_size)
|
|
|
- while idx < total:
|
|
|
- if cutlass.const_expr(self.activation_is_gated):
|
|
|
- row = idx // Int32(self.intermediate_size)
|
|
|
- col = idx - row * Int32(self.intermediate_size)
|
|
|
- base = row * Int32(self.fc1_cols)
|
|
|
- gate = fc1_bf16_flat[base + col].to(cutlass.Float32)
|
|
|
- up = fc1_bf16_flat[base + Int32(self.intermediate_size) + col].to(
|
|
|
- cutlass.Float32
|
|
|
- )
|
|
|
- gate, up = self._clamp_swiglu_inputs(gate, up)
|
|
|
- sigmoid_arg = gate
|
|
|
- up_term = up
|
|
|
- if cutlass.const_expr(self.activation_is_swigluoai):
|
|
|
- sigmoid_arg = cutlass.Float32(self.swiglu_alpha) * gate
|
|
|
- up_term = up + cutlass.Float32(self.swiglu_beta)
|
|
|
- if cutlass.const_expr(self.fast_math):
|
|
|
- exp_neg_gate = cute.math.exp(-sigmoid_arg, fastmath=True)
|
|
|
- else:
|
|
|
- exp_neg_gate = cute.math.exp(-sigmoid_arg, fastmath=False)
|
|
|
- silu = gate / (cutlass.Float32(1.0) + exp_neg_gate)
|
|
|
- if cutlass.const_expr(self.activation_is_swigluoai):
|
|
|
- activated_bf16_flat[idx] = self._cast_elem(silu * up_term)
|
|
|
- else:
|
|
|
- activated_bf16_flat[idx] = self._cast_elem(
|
|
|
- self._cast_elem(silu) * self._cast_elem(up_term)
|
|
|
- )
|
|
|
- else:
|
|
|
- x = fc1_bf16_flat[idx].to(cutlass.Float32)
|
|
|
- if x < cutlass.Float32(0.0):
|
|
|
- x = cutlass.Float32(0.0)
|
|
|
- activated_bf16_flat[idx] = self._cast_elem(x * x)
|
|
|
- idx += stride
|
|
|
-
|
|
|
-
|
|
|
-class W4A16ActivationKernel:
|
|
|
- def __init__(
|
|
|
- self,
|
|
|
- *,
|
|
|
- rows: int,
|
|
|
- intermediate_size: int,
|
|
|
- activation: str,
|
|
|
- element_dtype: str = "bf16",
|
|
|
- fast_math: bool = True,
|
|
|
- swiglu_limit: float | None = None,
|
|
|
- swiglu_alpha: float | None = None,
|
|
|
- swiglu_beta: float | None = None,
|
|
|
- ):
|
|
|
- activation = normalize_moe_activation(activation)
|
|
|
- is_gated = validate_activation(activation)
|
|
|
- swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params(
|
|
|
- activation,
|
|
|
- swiglu_limit,
|
|
|
- swiglu_alpha,
|
|
|
- swiglu_beta,
|
|
|
- )
|
|
|
- if element_dtype not in {"bf16", "fp16"}:
|
|
|
- raise ValueError(f"unsupported element_dtype {element_dtype!r}")
|
|
|
- if rows <= 0 or intermediate_size <= 0:
|
|
|
- raise ValueError("rows and intermediate_size must be positive")
|
|
|
- self.rows = int(rows)
|
|
|
- self.intermediate_size = int(intermediate_size)
|
|
|
- self.activation = activation
|
|
|
- self.is_gated = is_gated
|
|
|
- self.is_swigluoai = activation == SWIGLUOAI_UNINTERLEAVE
|
|
|
- self.has_swiglu_limit = swiglu_limit is not None
|
|
|
- self.swiglu_limit = 0.0 if swiglu_limit is None else swiglu_limit
|
|
|
- self.swiglu_alpha = float(swiglu_alpha)
|
|
|
- self.swiglu_beta = float(swiglu_beta)
|
|
|
- self.element_dtype = element_dtype
|
|
|
- self.is_fp16 = element_dtype == "fp16"
|
|
|
- self.fast_math = bool(fast_math)
|
|
|
- self.cta_threads = 256
|
|
|
-
|
|
|
- @property
|
|
|
- def __cache_key__(self) -> tuple[object, ...]:
|
|
|
- return (
|
|
|
- self.intermediate_size,
|
|
|
- self.activation,
|
|
|
- self.is_gated,
|
|
|
- self.is_swigluoai,
|
|
|
- self.has_swiglu_limit,
|
|
|
- self.swiglu_limit,
|
|
|
- self.swiglu_alpha,
|
|
|
- self.swiglu_beta,
|
|
|
- self.element_dtype,
|
|
|
- self.fast_math,
|
|
|
- self.cta_threads,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _cast_elem(self, x: cutlass.Float32):
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return cutlass.Float16(x)
|
|
|
- return cutlass.BFloat16(x)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _clamp_swiglu_inputs(
|
|
|
- self,
|
|
|
- gate: cutlass.Float32,
|
|
|
- up: cutlass.Float32,
|
|
|
- ):
|
|
|
- if cutlass.const_expr(self.has_swiglu_limit):
|
|
|
- limit = cutlass.Float32(self.swiglu_limit)
|
|
|
- neg_limit = cutlass.Float32(-self.swiglu_limit)
|
|
|
- if gate > limit:
|
|
|
- gate = limit
|
|
|
- if up > limit:
|
|
|
- up = limit
|
|
|
- if up < neg_limit:
|
|
|
- up = neg_limit
|
|
|
- return gate, up
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def __call__(
|
|
|
- self,
|
|
|
- fc1_flat: cute.Tensor,
|
|
|
- activated_flat: cute.Tensor,
|
|
|
- active_rows: cutlass.Int32,
|
|
|
- stream: cuda.CUstream,
|
|
|
- ):
|
|
|
- total = active_rows * Int32(self.intermediate_size)
|
|
|
- grid = (_covering_count(total, self.cta_threads), 1, 1)
|
|
|
- self.kernel(fc1_flat, activated_flat, active_rows).launch(
|
|
|
- grid=grid,
|
|
|
- block=[self.cta_threads, 1, 1],
|
|
|
- stream=stream,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.kernel
|
|
|
- def kernel(
|
|
|
- self,
|
|
|
- fc1_flat: cute.Tensor,
|
|
|
- activated_flat: cute.Tensor,
|
|
|
- active_rows: cutlass.Int32,
|
|
|
- ):
|
|
|
- tidx, _, _ = cute.arch.thread_idx()
|
|
|
- bidx, _, _ = cute.arch.block_idx()
|
|
|
- idx = Int32(bidx) * Int32(self.cta_threads) + Int32(tidx)
|
|
|
- total = Int32(active_rows) * Int32(self.intermediate_size)
|
|
|
- if idx < total:
|
|
|
- if cutlass.const_expr(self.is_gated):
|
|
|
- row = idx // Int32(self.intermediate_size)
|
|
|
- col = idx - row * Int32(self.intermediate_size)
|
|
|
- base = row * Int32(2 * self.intermediate_size)
|
|
|
- gate = fc1_flat[base + col].to(cutlass.Float32)
|
|
|
- up = fc1_flat[base + Int32(self.intermediate_size) + col].to(
|
|
|
- cutlass.Float32
|
|
|
- )
|
|
|
- gate, up = self._clamp_swiglu_inputs(gate, up)
|
|
|
- sigmoid_arg = gate
|
|
|
- up_term = up
|
|
|
- if cutlass.const_expr(self.is_swigluoai):
|
|
|
- sigmoid_arg = cutlass.Float32(self.swiglu_alpha) * gate
|
|
|
- up_term = up + cutlass.Float32(self.swiglu_beta)
|
|
|
- if cutlass.const_expr(self.fast_math):
|
|
|
- exp_neg_gate = cute.math.exp(-sigmoid_arg, fastmath=True)
|
|
|
- else:
|
|
|
- exp_neg_gate = cute.math.exp(-sigmoid_arg, fastmath=False)
|
|
|
- silu = gate / (cutlass.Float32(1.0) + exp_neg_gate)
|
|
|
- if cutlass.const_expr(self.is_swigluoai):
|
|
|
- activated_flat[idx] = self._cast_elem(silu * up_term)
|
|
|
- else:
|
|
|
- activated_flat[idx] = self._cast_elem(
|
|
|
- self._cast_elem(silu) * self._cast_elem(up_term)
|
|
|
- )
|
|
|
- else:
|
|
|
- x = fc1_flat[idx].to(cutlass.Float32)
|
|
|
- if x < cutlass.Float32(0.0):
|
|
|
- x = cutlass.Float32(0.0)
|
|
|
- activated_flat[idx] = self._cast_elem(x * x)
|
|
|
-
|
|
|
-
|
|
|
-class W4A16TopKSumKernel:
|
|
|
- def __init__(
|
|
|
- self, *, topk: int, hidden_size: int, element_dtype: str = "bf16"
|
|
|
- ):
|
|
|
- if element_dtype not in {"bf16", "fp16"}:
|
|
|
- raise ValueError(f"unsupported element_dtype {element_dtype!r}")
|
|
|
- if topk <= 0 or hidden_size <= 0:
|
|
|
- raise ValueError("topk and hidden_size must be positive")
|
|
|
- self.topk = int(topk)
|
|
|
- self.hidden_size = int(hidden_size)
|
|
|
- self.element_dtype = element_dtype
|
|
|
- self.is_fp16 = element_dtype == "fp16"
|
|
|
- self.cta_threads = 256
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def _cast_elem(self, x: cutlass.Float32):
|
|
|
- if cutlass.const_expr(self.is_fp16):
|
|
|
- return cutlass.Float16(x)
|
|
|
- return cutlass.BFloat16(x)
|
|
|
-
|
|
|
- @cute.jit
|
|
|
- def __call__(
|
|
|
- self,
|
|
|
- fc2_ptr: cute.Pointer,
|
|
|
- output_ptr: cute.Pointer,
|
|
|
- active_m: cutlass.Int32,
|
|
|
- stream: cuda.CUstream,
|
|
|
- ):
|
|
|
- fc2_flat = cute.make_tensor(
|
|
|
- fc2_ptr,
|
|
|
- layout=cute.make_layout(
|
|
|
- (active_m * Int32(self.topk * self.hidden_size),), stride=(1,)
|
|
|
- ),
|
|
|
- )
|
|
|
- output_flat = cute.make_tensor(
|
|
|
- output_ptr,
|
|
|
- layout=cute.make_layout((active_m * Int32(self.hidden_size),), stride=(1,)),
|
|
|
- )
|
|
|
- total = active_m * Int32(self.hidden_size)
|
|
|
- grid = (_covering_count(total, self.cta_threads), 1, 1)
|
|
|
- self.kernel(fc2_flat, output_flat, active_m).launch(
|
|
|
- grid=grid,
|
|
|
- block=[self.cta_threads, 1, 1],
|
|
|
- stream=stream,
|
|
|
- )
|
|
|
-
|
|
|
- @cute.kernel
|
|
|
- def kernel(
|
|
|
- self,
|
|
|
- fc2_flat: cute.Tensor,
|
|
|
- output_flat: cute.Tensor,
|
|
|
- active_m: cutlass.Int32,
|
|
|
- ):
|
|
|
- tidx, _, _ = cute.arch.thread_idx()
|
|
|
- bidx, _, _ = cute.arch.block_idx()
|
|
|
- idx = Int32(bidx) * Int32(self.cta_threads) + Int32(tidx)
|
|
|
- total = active_m * Int32(self.hidden_size)
|
|
|
- if idx < total:
|
|
|
- token = idx // Int32(self.hidden_size)
|
|
|
- col = idx - token * Int32(self.hidden_size)
|
|
|
- acc = cutlass.Float32(0.0)
|
|
|
- for route in cutlass.range_constexpr(self.topk):
|
|
|
- row = token * Int32(self.topk) + Int32(route)
|
|
|
- acc += fc2_flat[row * Int32(self.hidden_size) + col].to(cutlass.Float32)
|
|
|
- output_flat[idx] = self._cast_elem(acc)
|
|
|
-
|
|
|
-
|
|
|
-_CACHE: dict[tuple, W4A16GemmCompileResult] = {}
|
|
|
-_FUSED_CACHE: dict[tuple, W4A16FusedMoeCompileResult] = {}
|
|
|
-_ACTIVATION_CACHE: dict[tuple, W4A16ActivationCompileResult] = {}
|
|
|
-_SUM_CACHE: dict[tuple, W4A16TopKSumCompileResult] = {}
|
|
|
-_SMALL_M_DIRECT_CACHE: dict[tuple, _W4A16SmallMDirectLaunch] = {}
|
|
|
-
|
|
|
-
|
|
|
-def _normalize_element_dtype(dtype: torch.dtype) -> str:
|
|
|
- if dtype == torch.bfloat16:
|
|
|
- return "bf16"
|
|
|
- if dtype == torch.float16:
|
|
|
- return "fp16"
|
|
|
- raise TypeError(f"unsupported W4A16 activation dtype {dtype}")
|
|
|
-
|
|
|
-
|
|
|
-def _normalize_scale_format(scale_format: str) -> str:
|
|
|
- try:
|
|
|
- return _SCALE_FORMATS[scale_format.lower()]
|
|
|
- except KeyError as exc:
|
|
|
- raise ValueError(
|
|
|
- "scale_format must be one of 'e4m3_k16', 'e8m0_k32', or 'e4m3_k32', "
|
|
|
- f"got {scale_format!r}"
|
|
|
- ) from exc
|
|
|
-
|
|
|
-
|
|
|
-def _scale_group_size(scale_format: str) -> int:
|
|
|
- return (
|
|
|
- 32
|
|
|
- if _normalize_scale_format(scale_format) in ("e8m0_k32", "e4m3_k32")
|
|
|
- else 16
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-def _scale_fake_int32_elements(
|
|
|
- *,
|
|
|
- num_experts: int,
|
|
|
- size_k: int,
|
|
|
- size_n: int,
|
|
|
- scale_format: str,
|
|
|
- allow_k_tail: bool = False,
|
|
|
- allow_n_tail: bool = False,
|
|
|
-) -> int:
|
|
|
- group_size = _scale_group_size(scale_format)
|
|
|
- if int(size_n) % 16 != 0:
|
|
|
- raise ValueError(
|
|
|
- f"W4A16 {scale_format} scales require size_n divisible by 16"
|
|
|
- )
|
|
|
- if int(size_k) % group_size != 0 and not allow_k_tail:
|
|
|
- raise ValueError(
|
|
|
- f"W4A16 {scale_format} scales require size_k divisible by {group_size}, "
|
|
|
- f"got {size_k}"
|
|
|
- )
|
|
|
- groups_k = (
|
|
|
- _covering_count(int(size_k), group_size)
|
|
|
- if allow_k_tail
|
|
|
- else int(size_k) // group_size
|
|
|
- )
|
|
|
- scale_size_n = (
|
|
|
- _e8m0_logical_tail_scale_n(int(size_n)) if allow_n_tail else int(size_n)
|
|
|
- )
|
|
|
- return int(num_experts) * groups_k * (scale_size_n // 4)
|
|
|
-
|
|
|
-
|
|
|
-def _cutlass_element_dtype(element_dtype: str):
|
|
|
- if element_dtype == "bf16":
|
|
|
- return cutlass.BFloat16
|
|
|
- if element_dtype == "fp16":
|
|
|
- return cutlass.Float16
|
|
|
- raise ValueError(f"unsupported element_dtype {element_dtype!r}")
|
|
|
-
|
|
|
-
|
|
|
-def _small_m_direct_supported(
|
|
|
- *,
|
|
|
- m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- num_experts: int,
|
|
|
- topk: int,
|
|
|
- activation: str,
|
|
|
- apply_router_weight_on_input: bool,
|
|
|
- swiglu_limit: float | None,
|
|
|
- swiglu_alpha: float,
|
|
|
- swiglu_beta: float,
|
|
|
- element_dtype: str,
|
|
|
- weight_layout: str,
|
|
|
- w13_layout: str,
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
- expert_map: torch.Tensor | None = None,
|
|
|
-) -> bool:
|
|
|
- if os.environ.get("B12X_W4A16_SMALL_M_DIRECT", "1") == "0":
|
|
|
- return False
|
|
|
- activation = normalize_moe_activation(activation)
|
|
|
- swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params(
|
|
|
- activation,
|
|
|
- swiglu_limit,
|
|
|
- swiglu_alpha,
|
|
|
- swiglu_beta,
|
|
|
- )
|
|
|
- if activation == SWIGLUOAI_UNINTERLEAVE and swiglu_limit is None:
|
|
|
- return False
|
|
|
- return (
|
|
|
- element_dtype == "bf16"
|
|
|
- # Packed serving weights use the W4A16 TC-decode path for small M.
|
|
|
- # The direct micro kernel is kept limited to native ModelOpt weights.
|
|
|
- and weight_layout == "modelopt"
|
|
|
- and w13_layout in ("w13", "w31")
|
|
|
- and not bool(apply_router_weight_on_input)
|
|
|
- and (swiglu_limit is None or activation in ("silu", SWIGLUOAI_UNINTERLEAVE))
|
|
|
- and expert_map is None
|
|
|
- and _W4A16SmallMDirectKernel.is_supported(
|
|
|
- m=m,
|
|
|
- hidden_size=hidden_size,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- topk=topk,
|
|
|
- num_experts=num_experts,
|
|
|
- scale_format=scale_format,
|
|
|
- )
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-def _compile_w4a16_small_m_direct(
|
|
|
- *,
|
|
|
- m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- num_experts: int,
|
|
|
- topk: int,
|
|
|
- activation: str,
|
|
|
- fast_math: bool,
|
|
|
- topk_ids_dtype: torch.dtype,
|
|
|
- device: torch.device | None,
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
- swiglu_limit: float | None = None,
|
|
|
- swiglu_alpha: float | None = None,
|
|
|
- swiglu_beta: float | None = None,
|
|
|
- w13_layout: str = "w13",
|
|
|
-) -> _W4A16SmallMDirectLaunch:
|
|
|
- if topk_ids_dtype not in (torch.int32, torch.int64):
|
|
|
- raise TypeError("small-M W4A16 direct path requires int32/int64 topk_ids")
|
|
|
- activation = normalize_moe_activation(activation)
|
|
|
- swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params(
|
|
|
- activation,
|
|
|
- swiglu_limit,
|
|
|
- swiglu_alpha,
|
|
|
- swiglu_beta,
|
|
|
- )
|
|
|
- cache_key = (
|
|
|
- "w4a16_small_m_direct",
|
|
|
- None if device is None else int(device.index or 0),
|
|
|
- int(m),
|
|
|
- int(hidden_size),
|
|
|
- int(intermediate_size),
|
|
|
- int(num_experts),
|
|
|
- int(topk),
|
|
|
- activation,
|
|
|
- bool(fast_math),
|
|
|
- topk_ids_dtype,
|
|
|
- scale_format,
|
|
|
- swiglu_limit,
|
|
|
- swiglu_alpha,
|
|
|
- swiglu_beta,
|
|
|
- w13_layout,
|
|
|
- )
|
|
|
- cached = _SMALL_M_DIRECT_CACHE.get(cache_key)
|
|
|
- if cached is not None:
|
|
|
- return cached
|
|
|
-
|
|
|
- kernel = _W4A16SmallMDirectKernel(
|
|
|
- activation=activation,
|
|
|
- fast_math=bool(fast_math),
|
|
|
- share_input_across_experts=(int(m) == 1),
|
|
|
- share_expert_scales=True,
|
|
|
- single_token=(int(m) == 1),
|
|
|
- scale_format=scale_format,
|
|
|
- swiglu_limit=swiglu_limit,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- w13_layout=w13_layout,
|
|
|
- )
|
|
|
- kernel.configure(
|
|
|
- int(m),
|
|
|
- int(hidden_size),
|
|
|
- int(intermediate_size),
|
|
|
- int(topk),
|
|
|
- int(num_experts),
|
|
|
- device=device,
|
|
|
- )
|
|
|
-
|
|
|
- def dummy(dt):
|
|
|
- return make_ptr(dt, 16, cute.AddressSpace.gmem, assumed_align=16)
|
|
|
-
|
|
|
- ids_dtype = cutlass.Int32 if topk_ids_dtype == torch.int32 else cutlass.Int64
|
|
|
- barrier_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (1,),
|
|
|
- assumed_align=4,
|
|
|
- )
|
|
|
- raise_if_kernel_resolution_frozen(
|
|
|
- "cute.compile", target=kernel, cache_key=cache_key
|
|
|
- )
|
|
|
- compiled = b12x_compile(
|
|
|
- kernel,
|
|
|
- dummy(cutlass.BFloat16),
|
|
|
- dummy(cutlass.Uint8),
|
|
|
- dummy(cutlass.Uint8),
|
|
|
- dummy(cutlass.Float32),
|
|
|
- dummy(cutlass.Float32),
|
|
|
- dummy(cutlass.Float32),
|
|
|
- dummy(cutlass.Uint32),
|
|
|
- dummy(cutlass.Uint8),
|
|
|
- dummy(cutlass.Uint8),
|
|
|
- dummy(cutlass.Float32),
|
|
|
- dummy(ids_dtype),
|
|
|
- dummy(cutlass.Float32),
|
|
|
- dummy(cutlass.BFloat16),
|
|
|
- barrier_fake,
|
|
|
- barrier_fake,
|
|
|
- Int32(m),
|
|
|
- Int32(kernel.grid_x),
|
|
|
- current_cuda_stream(),
|
|
|
- compile_spec=KernelCompileSpec.from_facts(
|
|
|
- "moe.w4a16.small_m_direct",
|
|
|
- 1,
|
|
|
- ("device_index", None if device is None else int(device.index or 0)),
|
|
|
- ("m", int(m)),
|
|
|
- ("hidden_size", int(hidden_size)),
|
|
|
- ("intermediate_size", int(intermediate_size)),
|
|
|
- ("num_experts", int(num_experts)),
|
|
|
- ("topk", int(topk)),
|
|
|
- ("activation", activation),
|
|
|
- ("fast_math", bool(fast_math)),
|
|
|
- ("topk_ids_dtype", str(topk_ids_dtype)),
|
|
|
- ("scale_format", scale_format),
|
|
|
- ("swiglu_limit", swiglu_limit),
|
|
|
- ("swiglu_alpha", swiglu_alpha),
|
|
|
- ("swiglu_beta", swiglu_beta),
|
|
|
- ("w13_layout", w13_layout),
|
|
|
- ("grid_x", int(kernel.grid_x)),
|
|
|
- ),
|
|
|
- )
|
|
|
- launch = _W4A16SmallMDirectLaunch(
|
|
|
- compiled=compiled,
|
|
|
- grid_x=int(kernel.grid_x),
|
|
|
- m=int(m),
|
|
|
- hidden_size=int(hidden_size),
|
|
|
- intermediate_size=int(intermediate_size),
|
|
|
- num_experts=int(num_experts),
|
|
|
- topk=int(topk),
|
|
|
- activation=activation,
|
|
|
- fast_math=bool(fast_math),
|
|
|
- topk_ids_dtype=topk_ids_dtype,
|
|
|
- )
|
|
|
- _SMALL_M_DIRECT_CACHE[cache_key] = launch
|
|
|
- return launch
|
|
|
-
|
|
|
-
|
|
|
-def compile_w4a16_gemm(
|
|
|
- *,
|
|
|
- size_m: int,
|
|
|
- size_n: int,
|
|
|
- size_k: int,
|
|
|
- num_experts: int,
|
|
|
- top_k: int,
|
|
|
- mul_topk_weights: bool,
|
|
|
- tile_n: int,
|
|
|
- tile_k: int,
|
|
|
- moe_block_size: int,
|
|
|
- max_m_blocks: int,
|
|
|
- element_dtype: str = "bf16",
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
-) -> W4A16GemmCompileResult:
|
|
|
- scale_format = _normalize_scale_format(scale_format)
|
|
|
- cutlass_dtype = _cutlass_element_dtype(element_dtype)
|
|
|
- if torch.cuda.is_available():
|
|
|
- device = int(torch.cuda.current_device())
|
|
|
- else:
|
|
|
- device = None
|
|
|
- kernel = W4A16GemmKernel(
|
|
|
- size_m=size_m,
|
|
|
- size_n=size_n,
|
|
|
- size_k=size_k,
|
|
|
- num_experts=num_experts,
|
|
|
- top_k=top_k,
|
|
|
- mul_topk_weights=mul_topk_weights,
|
|
|
- tile_n=tile_n,
|
|
|
- tile_k=tile_k,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- max_m_blocks=max_m_blocks,
|
|
|
- element_dtype=element_dtype,
|
|
|
- scale_format=scale_format,
|
|
|
- )
|
|
|
- cache_key = (
|
|
|
- "w4a16_gemm",
|
|
|
- device,
|
|
|
- kernel.__cache_key__,
|
|
|
- )
|
|
|
- cached = _CACHE.get(cache_key)
|
|
|
- if cached is not None:
|
|
|
- return replace(
|
|
|
- cached,
|
|
|
- max_m_blocks=max_m_blocks,
|
|
|
- blocks_per_sm=kernel.blocks_per_sm,
|
|
|
- )
|
|
|
-
|
|
|
- compile_size_m = _fake_m_for_specialization(size_m)
|
|
|
- compile_route_blocks = 1
|
|
|
- compile_route_slots = compile_route_blocks * int(moe_block_size)
|
|
|
- a_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass_dtype,
|
|
|
- (compile_size_m * size_k,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- b_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (num_experts * (size_k // 16) * (size_n // 16 * 32),),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- c_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass_dtype,
|
|
|
- (compile_size_m * top_k * size_n,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- scales_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (
|
|
|
- _scale_fake_int32_elements(
|
|
|
- num_experts=num_experts,
|
|
|
- size_k=size_k,
|
|
|
- size_n=size_n,
|
|
|
- scale_format=scale_format,
|
|
|
- ),
|
|
|
- ),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- global_scale_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Float32,
|
|
|
- (num_experts,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- packed_routes_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (compile_route_slots,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- block_experts_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (compile_route_blocks,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- route_count_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (1,),
|
|
|
- assumed_align=4,
|
|
|
- )
|
|
|
- topk_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Float32,
|
|
|
- (compile_size_m * top_k,),
|
|
|
- assumed_align=4,
|
|
|
- )
|
|
|
- c_tmp_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Float32,
|
|
|
- (
|
|
|
- max(
|
|
|
- size_n * compile_route_slots,
|
|
|
- 4 * 256 * moe_block_size * 256,
|
|
|
- ),
|
|
|
- ),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- locks_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (4 * 256,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
-
|
|
|
- raise_if_kernel_resolution_frozen(
|
|
|
- "cute.compile", target=kernel, cache_key=cache_key
|
|
|
- )
|
|
|
- compiled = b12x_compile(
|
|
|
- kernel,
|
|
|
- a_fake,
|
|
|
- b_fake,
|
|
|
- c_fake,
|
|
|
- scales_fake,
|
|
|
- global_scale_fake,
|
|
|
- packed_routes_fake,
|
|
|
- block_experts_fake,
|
|
|
- route_count_fake,
|
|
|
- topk_fake,
|
|
|
- c_tmp_fake,
|
|
|
- locks_fake,
|
|
|
- Int32(compile_size_m),
|
|
|
- Int32(1),
|
|
|
- current_cuda_stream(),
|
|
|
- compile_spec=KernelCompileSpec.from_key(
|
|
|
- "moe.w4a16.gemm",
|
|
|
- 1,
|
|
|
- cache_key,
|
|
|
- ),
|
|
|
- )
|
|
|
- result = W4A16GemmCompileResult(
|
|
|
- compiled=compiled,
|
|
|
- tile_n=tile_n,
|
|
|
- tile_k=tile_k,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- max_m_blocks=max_m_blocks,
|
|
|
- blocks_per_sm=kernel.blocks_per_sm,
|
|
|
- scale_format=scale_format,
|
|
|
- )
|
|
|
- _CACHE[cache_key] = result
|
|
|
- return result
|
|
|
-
|
|
|
-
|
|
|
-def compile_w4a16_fused_moe(
|
|
|
- *,
|
|
|
- size_m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- num_experts: int,
|
|
|
- top_k: int,
|
|
|
- activation: str,
|
|
|
- apply_router_weight_on_input: bool,
|
|
|
- zero_fc2_output: bool,
|
|
|
- moe_block_size: int,
|
|
|
- max_m_blocks: int,
|
|
|
- element_dtype: str = "bf16",
|
|
|
- fast_math: bool = True,
|
|
|
- sms: int,
|
|
|
- max_shared_mem: int,
|
|
|
- swiglu_limit: float | None = None,
|
|
|
- swiglu_alpha: float | None = None,
|
|
|
- swiglu_beta: float | None = None,
|
|
|
- weight_layout: str = "packed",
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
- w13_layout: str = "w13",
|
|
|
- direct_topk_routes: bool = False,
|
|
|
- tc_decode_fused_sum: bool = False,
|
|
|
- collect_activation_amax: bool = False,
|
|
|
- force_tile_config: tuple[int, int, int, int] | None = None,
|
|
|
-) -> W4A16FusedMoeCompileResult:
|
|
|
- scale_format = _normalize_scale_format(scale_format)
|
|
|
- cutlass_dtype = _cutlass_element_dtype(element_dtype)
|
|
|
- device = int(torch.cuda.current_device()) if torch.cuda.is_available() else None
|
|
|
- activation = normalize_moe_activation(activation)
|
|
|
- is_gated = validate_activation(activation)
|
|
|
- swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params(
|
|
|
- activation,
|
|
|
- swiglu_limit,
|
|
|
- swiglu_alpha,
|
|
|
- swiglu_beta,
|
|
|
- )
|
|
|
- if weight_layout not in _WEIGHT_LAYOUTS:
|
|
|
- raise ValueError(f"unsupported W4A16 weight_layout {weight_layout!r}")
|
|
|
- if weight_layout == "modelopt":
|
|
|
- if w13_layout not in _MODEL_OPT_W13_LAYOUTS:
|
|
|
- raise ValueError(f"unsupported W4A16 w13_layout {w13_layout!r}")
|
|
|
- else:
|
|
|
- w13_layout = "packed"
|
|
|
- direct_topk_routes = bool(direct_topk_routes)
|
|
|
- tc_decode_fused_sum = bool(tc_decode_fused_sum)
|
|
|
- collect_activation_amax = bool(collect_activation_amax)
|
|
|
- if collect_activation_amax and (direct_topk_routes or tc_decode_fused_sum):
|
|
|
- raise ValueError(
|
|
|
- "W4A16 activation amax collection requires the route-packed fused path"
|
|
|
- )
|
|
|
- # The TC-decode path validates M in {1,2,4,8} itself and uses direct-topk
|
|
|
- # routing for the whole {1,2,4,8} range, so it lifts the default decode cap.
|
|
|
- direct_topk_m_cap = (
|
|
|
- _W4A16_SMALL_M_DIRECT_MAX_M if tc_decode_fused_sum else _MAX_DIRECT_TOPK_ROUTE_M
|
|
|
- )
|
|
|
- if direct_topk_routes and (
|
|
|
- int(size_m) > direct_topk_m_cap
|
|
|
- or weight_layout not in ("packed", "nf3_2p1")
|
|
|
- or bool(zero_fc2_output)
|
|
|
- ):
|
|
|
- raise ValueError(
|
|
|
- "direct_topk_routes is only valid for small-M packed W4A16 without expert_map"
|
|
|
- )
|
|
|
- fc1_cols = int(intermediate_size) * (2 if is_gated else 1)
|
|
|
- routed_rows = int(size_m) * int(top_k)
|
|
|
- allow_native_logical_tail = (
|
|
|
- weight_layout == "modelopt"
|
|
|
- and scale_format == "e8m0_k32"
|
|
|
- and int(intermediate_size) % 32 != 0
|
|
|
- )
|
|
|
- fc1_tile_k, fc1_tile_n, fc1_cta_threads, _ = _select_tile_config(
|
|
|
- problem_m=size_m,
|
|
|
- problem_n=fc1_cols,
|
|
|
- problem_k=hidden_size,
|
|
|
- top_k=top_k,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- sms=sms,
|
|
|
- max_shared_mem=max_shared_mem,
|
|
|
- scale_format=scale_format,
|
|
|
- allow_logical_tail=allow_native_logical_tail,
|
|
|
- )
|
|
|
- fc2_tile_k, fc2_tile_n, fc2_cta_threads, _ = _select_tile_config(
|
|
|
- problem_m=routed_rows,
|
|
|
- problem_n=hidden_size,
|
|
|
- problem_k=intermediate_size,
|
|
|
- top_k=1,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- sms=sms,
|
|
|
- max_shared_mem=max_shared_mem,
|
|
|
- scale_format=scale_format,
|
|
|
- allow_logical_tail=allow_native_logical_tail,
|
|
|
- )
|
|
|
- if fc1_cta_threads != fc2_cta_threads:
|
|
|
- common_cta_threads = min(fc1_cta_threads, fc2_cta_threads)
|
|
|
- fc1_tile_k, fc1_tile_n, fc1_cta_threads, _ = _select_tile_config(
|
|
|
- problem_m=size_m,
|
|
|
- problem_n=fc1_cols,
|
|
|
- problem_k=hidden_size,
|
|
|
- top_k=top_k,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- sms=sms,
|
|
|
- max_shared_mem=max_shared_mem,
|
|
|
- required_cta_threads=common_cta_threads,
|
|
|
- scale_format=scale_format,
|
|
|
- allow_logical_tail=allow_native_logical_tail,
|
|
|
- )
|
|
|
- fc2_tile_k, fc2_tile_n, fc2_cta_threads, _ = _select_tile_config(
|
|
|
- problem_m=routed_rows,
|
|
|
- problem_n=hidden_size,
|
|
|
- problem_k=intermediate_size,
|
|
|
- top_k=1,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- sms=sms,
|
|
|
- max_shared_mem=max_shared_mem,
|
|
|
- required_cta_threads=common_cta_threads,
|
|
|
- scale_format=scale_format,
|
|
|
- allow_logical_tail=allow_native_logical_tail,
|
|
|
- )
|
|
|
- if fc1_cta_threads != fc2_cta_threads:
|
|
|
- raise ValueError(
|
|
|
- "fused W4A16 FC1/FC2 selected different thread counts: "
|
|
|
- f"{fc1_cta_threads} vs {fc2_cta_threads}"
|
|
|
- )
|
|
|
- # TC-decode FC1 wide-N override (single-wave-collapse sizes only): in the m8
|
|
|
- # fused path the host right-sizes grid_x to the FC1 mn-tile count and forces
|
|
|
- # one whole mn-tile per CTA over the full K. With the default fc1_tile_n=128,
|
|
|
- # FC1 (N=fc1_cols=2*intermediate_size) produces size_m*top_k*(fc1_cols/128)
|
|
|
- # mn-tiles. For TP=2 I_tp=1024 (fc1_cols=2048 => 16 n-tiles/route) bs=1 is
|
|
|
- # 96 tiles (<= 188 SMs, one wave); bs=2 is 192 -- JUST over the single-wave
|
|
|
- # SM cap -- forcing a 2-wave launch of grid_x=96 (half the machine idle per
|
|
|
- # wave, a serialized second FC1 wave of pure tail latency on the bandwidth-
|
|
|
- # bound decode). Widening FC1 to tile_n=256 (256-wide N slab per CTA over
|
|
|
- # full K; tile_k=64 keeps cta_threads=256) HALVES FC1's mn-tile count, so
|
|
|
- # bs=2 collapses to 96 tiles = exactly ONE wave, removing that whole second
|
|
|
- # FC1 wave. The narrower tile_k=64 is slower per-tile, so we apply this ONLY
|
|
|
- # where it turns a 2-wave launch into a 1-wave launch: the default 128-wide
|
|
|
- # FC1 spans 2 waves (sms < default_mn_tiles <= 2*sms) AND the wide tile fits
|
|
|
- # in one (default/2 <= sms). bs=1 (one wave already) and bs>=4 (still multi-
|
|
|
- # wave after halving) keep the faster default 128x128 tile.
|
|
|
- # Guarded by fc1_cols%256==0, smem-fit, and 256-thread geometry so the fused
|
|
|
- # FC1/FC2 single-thread-geometry contract is preserved.
|
|
|
- default_fc1_mn_tiles = (
|
|
|
- int(size_m) * int(top_k) * (int(fc1_cols) // int(fc1_tile_n))
|
|
|
- if fc1_tile_n > 0
|
|
|
- else 0
|
|
|
- )
|
|
|
- if (
|
|
|
- bool(tc_decode_fused_sum)
|
|
|
- and int(fc1_cols) % 256 == 0
|
|
|
- and fc1_tile_n == 128
|
|
|
- and (fc1_tile_n * fc1_tile_k) // 64 == 256
|
|
|
- and int(sms) < default_fc1_mn_tiles <= 2 * int(sms)
|
|
|
- and (default_fc1_mn_tiles // 2) <= int(sms)
|
|
|
- ):
|
|
|
- wide_fc1_tile_k = 64
|
|
|
- if _candidate_tile_fits(
|
|
|
- problem_n=fc1_cols,
|
|
|
- problem_k=hidden_size,
|
|
|
- cta_m_blocks=_covering_count(moe_block_size, 16),
|
|
|
- tile_n=256,
|
|
|
- tile_k=wide_fc1_tile_k,
|
|
|
- cta_threads=256,
|
|
|
- max_shared_mem=int(max_shared_mem) - 512,
|
|
|
- scale_format=scale_format,
|
|
|
- ):
|
|
|
- fc1_tile_n = 256
|
|
|
- fc1_tile_k = wide_fc1_tile_k
|
|
|
- fc1_cta_threads = 256
|
|
|
- # TC-decode FC2 wide-N override: in the m8 fused path the host right-sizes
|
|
|
- # grid_x to the FC1 mn-tile count (one persistent wave, no split-K). FC2
|
|
|
- # (N=hidden_size, K=intermediate_size) with the default tile_n=128 produces
|
|
|
- # route_blocks*(hidden_size/128) mn-tiles -- roughly double FC1's count --
|
|
|
- # so FC2 would need ~2 persistent waves while FC1 fits in 1; that second
|
|
|
- # FC2 wave is pure serialized tail latency on the bandwidth-bound decode.
|
|
|
- # Selecting tile_n=256 for FC2 (a 256-wide N slab per CTA over full K)
|
|
|
- # halves FC2's mn-tile count so it also fits one wave. Guarded by smem-fit,
|
|
|
- # hidden_size%256==0, and matching cta_threads so the fused single
|
|
|
- # thread-geometry contract is preserved.
|
|
|
- if (
|
|
|
- bool(tc_decode_fused_sum)
|
|
|
- and int(hidden_size) % 256 == 0
|
|
|
- and fc2_tile_n == 128
|
|
|
- and fc1_cta_threads == 256
|
|
|
- ):
|
|
|
- wide_fc2_tile_k = 64
|
|
|
- if _candidate_tile_fits(
|
|
|
- problem_n=hidden_size,
|
|
|
- problem_k=intermediate_size,
|
|
|
- cta_m_blocks=_covering_count(moe_block_size, 16),
|
|
|
- tile_n=256,
|
|
|
- tile_k=wide_fc2_tile_k,
|
|
|
- cta_threads=256,
|
|
|
- max_shared_mem=int(max_shared_mem) - 512,
|
|
|
- scale_format=scale_format,
|
|
|
- ):
|
|
|
- fc2_tile_n = 256
|
|
|
- fc2_tile_k = wide_fc2_tile_k
|
|
|
- fc2_cta_threads = 256
|
|
|
- # TC-decode FC2 ultra-wide override (perfect wave-balance with FC1): the
|
|
|
- # persistent grid_x is right-sized to FC1's mn-tile count. After the FC1/FC2
|
|
|
- # wide-N (tile_n=256) overrides, bs=2 still has FC1=route_blocks*8 tiles vs
|
|
|
- # FC2=route_blocks*16 tiles -- FC2 is DOUBLE FC1 and thus needs a second
|
|
|
- # persistent wave at grid_x sized to FC1. That FC2 second wave is pure
|
|
|
- # serialized tail latency on the bandwidth-bound decode. Widening FC2 to
|
|
|
- # tile_n=512 (a 512-wide N slab per CTA over full K) halves FC2's mn-tile
|
|
|
- # count again so FC2 == FC1's tile count and fits the SAME single wave. A
|
|
|
- # 512-wide N tile needs tile_k=32 to keep cta_threads=256 (512*32/64); that
|
|
|
- # is below the generic tile_k>=64 fits-floor, so we validate the footprint
|
|
|
- # directly here. tile_k=32 == the e8m0_k32 scale group, so cta_k_blocks=2
|
|
|
- # with one e8m0 scale group per k-tile -- the existing scale layout is a
|
|
|
- # clean covering. Fire ONLY when it drops FC2 from >1 wave to FC1's wave
|
|
|
- # count (bs=2). Numerically identical: only the FC2 output-tile width changes.
|
|
|
- fc1_mn_after = (
|
|
|
- int(size_m) * int(top_k) * (int(fc1_cols) // int(fc1_tile_n))
|
|
|
- if fc1_tile_n > 0
|
|
|
- else 0
|
|
|
- )
|
|
|
- fc2_mn_after = (
|
|
|
- int(size_m) * int(top_k) * (int(hidden_size) // int(fc2_tile_n))
|
|
|
- if fc2_tile_n > 0
|
|
|
- else 0
|
|
|
- )
|
|
|
- if (
|
|
|
- bool(tc_decode_fused_sum)
|
|
|
- and int(hidden_size) % 512 == 0
|
|
|
- and fc2_tile_n == 256
|
|
|
- and fc1_cta_threads == 256
|
|
|
- and fc1_mn_after > 0
|
|
|
- and fc2_mn_after > fc1_mn_after
|
|
|
- and (fc2_mn_after // 2) <= fc1_mn_after
|
|
|
- and fc1_mn_after <= int(sms)
|
|
|
- ):
|
|
|
- ultra_fc2_tile_k = 32
|
|
|
- ultra_smem = _shared_memory_footprint(
|
|
|
- cta_m_blocks=_covering_count(moe_block_size, 16),
|
|
|
- tile_n=512,
|
|
|
- tile_k=ultra_fc2_tile_k,
|
|
|
- scale_format=scale_format,
|
|
|
- )
|
|
|
- if (
|
|
|
- int(intermediate_size) % ultra_fc2_tile_k == 0
|
|
|
- and ultra_smem <= int(max_shared_mem) - 512
|
|
|
- ):
|
|
|
- fc2_tile_n = 512
|
|
|
- fc2_tile_k = ultra_fc2_tile_k
|
|
|
- fc2_cta_threads = 256
|
|
|
- if force_tile_config is not None:
|
|
|
- # Explicit (fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n) pin. The
|
|
|
- # NF3 ("nf3_2p1") flat-span weight layout is packed for a specific CTA
|
|
|
- # N-tile, so hybrid deployments pin ONE tile config across every m
|
|
|
- # regime instead of the m-dependent auto selection (and TC-decode
|
|
|
- # wide-N overrides) above. Overrides whatever was selected; the tiles
|
|
|
- # land in the GEMM cache keys, so no cache collision is possible.
|
|
|
- fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n = (
|
|
|
- int(v) for v in force_tile_config
|
|
|
- )
|
|
|
- fc1_cta_threads = (fc1_tile_n * fc1_tile_k) // 64
|
|
|
- fc2_cta_threads = (fc2_tile_n * fc2_tile_k) // 64
|
|
|
- if fc1_cta_threads != fc2_cta_threads:
|
|
|
- raise ValueError(
|
|
|
- "force_tile_config FC1/FC2 thread counts must match, got "
|
|
|
- f"{fc1_cta_threads} vs {fc2_cta_threads}"
|
|
|
- )
|
|
|
- for name, forced_pn, forced_pk, forced_tn, forced_tk in (
|
|
|
- ("fc1", fc1_cols, hidden_size, fc1_tile_n, fc1_tile_k),
|
|
|
- ("fc2", hidden_size, intermediate_size, fc2_tile_n, fc2_tile_k),
|
|
|
- ):
|
|
|
- if not _candidate_tile_fits(
|
|
|
- problem_n=forced_pn,
|
|
|
- problem_k=forced_pk,
|
|
|
- cta_m_blocks=_covering_count(moe_block_size, 16),
|
|
|
- tile_n=forced_tn,
|
|
|
- tile_k=forced_tk,
|
|
|
- cta_threads=fc1_cta_threads,
|
|
|
- max_shared_mem=int(max_shared_mem) - 512,
|
|
|
- scale_format=scale_format,
|
|
|
- allow_logical_tail=allow_native_logical_tail,
|
|
|
- ):
|
|
|
- raise ValueError(
|
|
|
- f"force_tile_config {name} tile "
|
|
|
- f"(tile_k={forced_tk}, tile_n={forced_tn}) does not fit "
|
|
|
- f"problem N/K={forced_pn}/{forced_pk} at "
|
|
|
- f"moe_block_size={moe_block_size}"
|
|
|
- )
|
|
|
- kernel = W4A16FusedMoeKernel(
|
|
|
- size_m=size_m,
|
|
|
- hidden_size=hidden_size,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- num_experts=num_experts,
|
|
|
- top_k=top_k,
|
|
|
- activation=activation,
|
|
|
- apply_router_weight_on_input=apply_router_weight_on_input,
|
|
|
- zero_fc2_output=zero_fc2_output,
|
|
|
- fc1_tile_n=fc1_tile_n,
|
|
|
- fc1_tile_k=fc1_tile_k,
|
|
|
- fc2_tile_n=fc2_tile_n,
|
|
|
- fc2_tile_k=fc2_tile_k,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- max_m_blocks=max_m_blocks,
|
|
|
- element_dtype=element_dtype,
|
|
|
- fast_math=fast_math,
|
|
|
- swiglu_limit=swiglu_limit,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- weight_layout=weight_layout,
|
|
|
- scale_format=scale_format,
|
|
|
- w13_layout=w13_layout,
|
|
|
- direct_topk_routes=direct_topk_routes,
|
|
|
- tc_decode_fused_sum=tc_decode_fused_sum,
|
|
|
- collect_activation_amax=collect_activation_amax,
|
|
|
- )
|
|
|
- cache_key = (
|
|
|
- "w4a16_fused_moe",
|
|
|
- device,
|
|
|
- kernel.__cache_key__,
|
|
|
- )
|
|
|
- cached = _FUSED_CACHE.get(cache_key)
|
|
|
- if cached is not None:
|
|
|
- return replace(
|
|
|
- cached,
|
|
|
- size_m=size_m,
|
|
|
- max_m_blocks=max_m_blocks,
|
|
|
- blocks_per_sm=kernel.blocks_per_sm,
|
|
|
- )
|
|
|
-
|
|
|
- if (not collect_activation_amax) and _small_m_direct_supported(
|
|
|
- m=size_m,
|
|
|
- hidden_size=hidden_size,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- num_experts=num_experts,
|
|
|
- topk=top_k,
|
|
|
- activation=activation,
|
|
|
- apply_router_weight_on_input=bool(apply_router_weight_on_input),
|
|
|
- swiglu_limit=swiglu_limit,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- element_dtype=element_dtype,
|
|
|
- weight_layout=weight_layout,
|
|
|
- w13_layout=w13_layout,
|
|
|
- scale_format=scale_format,
|
|
|
- ):
|
|
|
- for ids_dtype in (torch.int32, torch.int64):
|
|
|
- _compile_w4a16_small_m_direct(
|
|
|
- m=size_m,
|
|
|
- hidden_size=hidden_size,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- num_experts=num_experts,
|
|
|
- topk=top_k,
|
|
|
- activation=activation,
|
|
|
- fast_math=fast_math,
|
|
|
- topk_ids_dtype=ids_dtype,
|
|
|
- scale_format=scale_format,
|
|
|
- swiglu_limit=swiglu_limit,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- w13_layout=w13_layout,
|
|
|
- device=torch.device("cuda", device) if device is not None else None,
|
|
|
- )
|
|
|
-
|
|
|
- compile_size_m = _fake_m_for_specialization(size_m)
|
|
|
- compile_routed_rows = int(compile_size_m) * int(top_k)
|
|
|
- compile_route_blocks = compile_routed_rows if direct_topk_routes else 1
|
|
|
- compile_route_slots = compile_route_blocks * int(moe_block_size)
|
|
|
- packed_route_fake_elements = (
|
|
|
- compile_routed_rows
|
|
|
- if direct_topk_routes
|
|
|
- else compile_route_slots
|
|
|
- )
|
|
|
- a_fake = make_ptr(cutlass_dtype, 16, cute.AddressSpace.gmem, assumed_align=16)
|
|
|
- if weight_layout == "modelopt":
|
|
|
- w13_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Uint8,
|
|
|
- (num_experts * fc1_cols * (hidden_size // 2),),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- w2_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Uint8,
|
|
|
- (num_experts * hidden_size * (intermediate_size // 2),),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- elif weight_layout == "nf3_2p1":
|
|
|
- # NF3: int32, 3 words per 32-code unit; (size_n // 2) units per K16 row.
|
|
|
- w13_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (num_experts * (hidden_size // 16) * (fc1_cols // 2) * 3,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- w2_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (num_experts * (intermediate_size // 16) * (hidden_size // 2) * 3,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- else:
|
|
|
- w13_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (num_experts * (hidden_size // 16) * (fc1_cols // 16 * 32),),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- w2_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (num_experts * (intermediate_size // 16) * (hidden_size // 16 * 32),),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- fc1_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass_dtype,
|
|
|
- (compile_routed_rows * fc1_cols,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- activated_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass_dtype,
|
|
|
- (compile_routed_rows * intermediate_size,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- fc2_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass_dtype,
|
|
|
- (compile_routed_rows * hidden_size,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- w13_scales_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (
|
|
|
- _scale_fake_int32_elements(
|
|
|
- num_experts=num_experts,
|
|
|
- size_k=hidden_size,
|
|
|
- size_n=fc1_cols,
|
|
|
- scale_format=scale_format,
|
|
|
- allow_n_tail=allow_native_logical_tail,
|
|
|
- ),
|
|
|
- ),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- w2_scales_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (
|
|
|
- _scale_fake_int32_elements(
|
|
|
- num_experts=num_experts,
|
|
|
- size_k=intermediate_size,
|
|
|
- size_n=hidden_size,
|
|
|
- scale_format=scale_format,
|
|
|
- allow_k_tail=allow_native_logical_tail,
|
|
|
- ),
|
|
|
- ),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- w13_global_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Float32,
|
|
|
- (num_experts,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- w2_global_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Float32,
|
|
|
- (num_experts,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- packed_routes_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (packed_route_fake_elements,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- block_experts_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (compile_route_blocks,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- route_count_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (1,),
|
|
|
- assumed_align=4,
|
|
|
- )
|
|
|
- activation_amax_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Float32,
|
|
|
- (num_experts * 2,),
|
|
|
- assumed_align=4,
|
|
|
- )
|
|
|
- topk_fake = make_ptr(cutlass.Float32, 4, cute.AddressSpace.gmem, assumed_align=4)
|
|
|
- fc1_c_tmp_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Float32,
|
|
|
- (
|
|
|
- max(
|
|
|
- fc1_cols * compile_route_slots,
|
|
|
- 4 * 256 * moe_block_size * 256,
|
|
|
- ),
|
|
|
- ),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- fc2_c_tmp_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Float32,
|
|
|
- (
|
|
|
- max(
|
|
|
- hidden_size * compile_route_slots,
|
|
|
- 4 * 256 * moe_block_size * 256,
|
|
|
- ),
|
|
|
- ),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- locks_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass.Int32,
|
|
|
- (4 * 256 + 2,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
-
|
|
|
- raise_if_kernel_resolution_frozen(
|
|
|
- "cute.compile", target=kernel, cache_key=cache_key
|
|
|
- )
|
|
|
- compiled = b12x_compile(
|
|
|
- kernel,
|
|
|
- a_fake,
|
|
|
- w13_fake,
|
|
|
- w2_fake,
|
|
|
- fc1_fake,
|
|
|
- activated_fake,
|
|
|
- fc2_fake,
|
|
|
- w13_scales_fake,
|
|
|
- w2_scales_fake,
|
|
|
- w13_global_fake,
|
|
|
- w2_global_fake,
|
|
|
- packed_routes_fake,
|
|
|
- block_experts_fake,
|
|
|
- route_count_fake,
|
|
|
- activation_amax_fake,
|
|
|
- 0,
|
|
|
- topk_fake,
|
|
|
- fc1_c_tmp_fake,
|
|
|
- fc2_c_tmp_fake,
|
|
|
- locks_fake,
|
|
|
- 1,
|
|
|
- 1,
|
|
|
- current_cuda_stream(),
|
|
|
- compile_spec=KernelCompileSpec.from_key(
|
|
|
- "moe.w4a16.fused_moe",
|
|
|
- 1,
|
|
|
- cache_key,
|
|
|
- ),
|
|
|
- )
|
|
|
- result = W4A16FusedMoeCompileResult(
|
|
|
- compiled=compiled,
|
|
|
- size_m=size_m,
|
|
|
- hidden_size=hidden_size,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- num_experts=num_experts,
|
|
|
- top_k=top_k,
|
|
|
- activation=activation,
|
|
|
- apply_router_weight_on_input=bool(apply_router_weight_on_input),
|
|
|
- zero_fc2_output=bool(zero_fc2_output),
|
|
|
- element_dtype=element_dtype,
|
|
|
- fast_math=bool(fast_math),
|
|
|
- swiglu_limit=swiglu_limit,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- fc1_tile_n=fc1_tile_n,
|
|
|
- fc1_tile_k=fc1_tile_k,
|
|
|
- fc2_tile_n=fc2_tile_n,
|
|
|
- fc2_tile_k=fc2_tile_k,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- max_m_blocks=max_m_blocks,
|
|
|
- blocks_per_sm=kernel.blocks_per_sm,
|
|
|
- weight_layout=weight_layout,
|
|
|
- w13_layout=w13_layout,
|
|
|
- direct_topk_routes=kernel.direct_topk_routes,
|
|
|
- scale_format=scale_format,
|
|
|
- tc_decode_fused_sum=bool(tc_decode_fused_sum),
|
|
|
- collect_activation_amax=collect_activation_amax,
|
|
|
- )
|
|
|
- _FUSED_CACHE[cache_key] = result
|
|
|
- return result
|
|
|
-
|
|
|
-
|
|
|
-def clear_w4a16_kernel_cache() -> None:
|
|
|
- _CACHE.clear()
|
|
|
- _FUSED_CACHE.clear()
|
|
|
- _ACTIVATION_CACHE.clear()
|
|
|
- _SUM_CACHE.clear()
|
|
|
- _SMALL_M_DIRECT_CACHE.clear()
|
|
|
-
|
|
|
-
|
|
|
-def compile_w4a16_activation(
|
|
|
- *,
|
|
|
- rows: int,
|
|
|
- intermediate_size: int,
|
|
|
- activation: str,
|
|
|
- element_dtype: str = "bf16",
|
|
|
- fast_math: bool = True,
|
|
|
- swiglu_limit: float | None = None,
|
|
|
- swiglu_alpha: float | None = None,
|
|
|
- swiglu_beta: float | None = None,
|
|
|
-) -> W4A16ActivationCompileResult:
|
|
|
- cutlass_dtype = _cutlass_element_dtype(element_dtype)
|
|
|
- activation = normalize_moe_activation(activation)
|
|
|
- is_gated = validate_activation(activation)
|
|
|
- swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params(
|
|
|
- activation,
|
|
|
- swiglu_limit,
|
|
|
- swiglu_alpha,
|
|
|
- swiglu_beta,
|
|
|
- )
|
|
|
- kernel = W4A16ActivationKernel(
|
|
|
- rows=rows,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- activation=activation,
|
|
|
- element_dtype=element_dtype,
|
|
|
- fast_math=fast_math,
|
|
|
- swiglu_limit=swiglu_limit,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- )
|
|
|
- cache_key = (
|
|
|
- "w4a16_activation",
|
|
|
- kernel.__cache_key__,
|
|
|
- )
|
|
|
- cached = _ACTIVATION_CACHE.get(cache_key)
|
|
|
- if cached is not None:
|
|
|
- return replace(cached, rows=rows)
|
|
|
-
|
|
|
- w13_shards = 2 if is_gated else 1
|
|
|
- compile_rows = _fake_m_for_specialization(rows)
|
|
|
- fc1_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass_dtype,
|
|
|
- (compile_rows * w13_shards * intermediate_size,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- activated_fake = cute.runtime.make_fake_compact_tensor(
|
|
|
- cutlass_dtype,
|
|
|
- (compile_rows * intermediate_size,),
|
|
|
- assumed_align=16,
|
|
|
- )
|
|
|
- raise_if_kernel_resolution_frozen(
|
|
|
- "cute.compile", target=kernel, cache_key=cache_key
|
|
|
- )
|
|
|
- compiled = b12x_compile(
|
|
|
- kernel,
|
|
|
- fc1_fake,
|
|
|
- activated_fake,
|
|
|
- Int32(compile_rows),
|
|
|
- current_cuda_stream(),
|
|
|
- compile_spec=KernelCompileSpec.from_key(
|
|
|
- "moe.w4a16.activation",
|
|
|
- 1,
|
|
|
- cache_key,
|
|
|
- ),
|
|
|
- )
|
|
|
- result = W4A16ActivationCompileResult(
|
|
|
- compiled=compiled,
|
|
|
- rows=rows,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- activation=activation,
|
|
|
- swiglu_limit=swiglu_limit,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- )
|
|
|
- _ACTIVATION_CACHE[cache_key] = result
|
|
|
- return result
|
|
|
-
|
|
|
-
|
|
|
-def compile_w4a16_topk_sum(
|
|
|
- *,
|
|
|
- m: int,
|
|
|
- topk: int,
|
|
|
- hidden_size: int,
|
|
|
- element_dtype: str = "bf16",
|
|
|
-) -> W4A16TopKSumCompileResult:
|
|
|
- cutlass_dtype = _cutlass_element_dtype(element_dtype)
|
|
|
- cache_key = ("w4a16_topk_sum", element_dtype, topk, hidden_size)
|
|
|
- cached = _SUM_CACHE.get(cache_key)
|
|
|
- if cached is not None:
|
|
|
- return cached
|
|
|
-
|
|
|
- fc2_fake = make_ptr(cutlass_dtype, 16, cute.AddressSpace.gmem, assumed_align=16)
|
|
|
- output_fake = make_ptr(cutlass_dtype, 16, cute.AddressSpace.gmem, assumed_align=16)
|
|
|
- kernel = W4A16TopKSumKernel(
|
|
|
- topk=topk,
|
|
|
- hidden_size=hidden_size,
|
|
|
- element_dtype=element_dtype,
|
|
|
- )
|
|
|
- raise_if_kernel_resolution_frozen(
|
|
|
- "cute.compile", target=kernel, cache_key=cache_key
|
|
|
- )
|
|
|
- compiled = b12x_compile(
|
|
|
- kernel,
|
|
|
- fc2_fake,
|
|
|
- output_fake,
|
|
|
- 1,
|
|
|
- current_cuda_stream(),
|
|
|
- compile_spec=KernelCompileSpec.from_key(
|
|
|
- "moe.w4a16.topk_sum",
|
|
|
- 1,
|
|
|
- cache_key,
|
|
|
- ),
|
|
|
- )
|
|
|
- result = W4A16TopKSumCompileResult(
|
|
|
- compiled=compiled,
|
|
|
- m=0,
|
|
|
- topk=topk,
|
|
|
- hidden_size=hidden_size,
|
|
|
- )
|
|
|
- _SUM_CACHE[cache_key] = result
|
|
|
- return result
|
|
|
-
|
|
|
-
|
|
|
-def _w4a16_small_m_direct_launch_flat(
|
|
|
- a_input: torch.Tensor,
|
|
|
- w13_u8: torch.Tensor,
|
|
|
- w13_scale_u8: torch.Tensor,
|
|
|
- w13_global_scale: torch.Tensor,
|
|
|
- w2_global_scale: torch.Tensor,
|
|
|
- inter_u32: torch.Tensor,
|
|
|
- w2_u8: torch.Tensor,
|
|
|
- w2_scale_u8: torch.Tensor,
|
|
|
- topk_ids: torch.Tensor,
|
|
|
- topk_weights: torch.Tensor,
|
|
|
- output: torch.Tensor,
|
|
|
- barrier_count: torch.Tensor,
|
|
|
- barrier_epoch: torch.Tensor,
|
|
|
- m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- num_experts: int,
|
|
|
- topk: int,
|
|
|
- activation: str,
|
|
|
- fast_math: bool,
|
|
|
- scale_format: str,
|
|
|
- has_swiglu_limit: bool,
|
|
|
- swiglu_limit_value: float,
|
|
|
- swiglu_alpha: float,
|
|
|
- swiglu_beta: float,
|
|
|
- w13_layout: str,
|
|
|
- stream_int: int,
|
|
|
-) -> None:
|
|
|
- swiglu_limit = float(swiglu_limit_value) if has_swiglu_limit else None
|
|
|
- direct_launch = _compile_w4a16_small_m_direct(
|
|
|
- m=m,
|
|
|
- hidden_size=hidden_size,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- num_experts=num_experts,
|
|
|
- topk=topk,
|
|
|
- activation=activation,
|
|
|
- fast_math=bool(fast_math),
|
|
|
- topk_ids_dtype=topk_ids.dtype,
|
|
|
- scale_format=scale_format,
|
|
|
- swiglu_limit=swiglu_limit,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- w13_layout=w13_layout,
|
|
|
- device=a_input.device,
|
|
|
- )
|
|
|
-
|
|
|
- def ptr(dt, tensor: torch.Tensor):
|
|
|
- return make_ptr(dt, tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=16)
|
|
|
-
|
|
|
- ids_dtype = cutlass.Int64 if topk_ids.dtype == torch.int64 else cutlass.Int32
|
|
|
- direct_launch.compiled(
|
|
|
- ptr(cutlass.BFloat16, a_input),
|
|
|
- ptr(cutlass.Uint8, w13_u8),
|
|
|
- ptr(cutlass.Uint8, w13_scale_u8.view(torch.uint8)),
|
|
|
- ptr(cutlass.Float32, w13_global_scale),
|
|
|
- ptr(cutlass.Float32, w13_global_scale),
|
|
|
- ptr(cutlass.Float32, w2_global_scale),
|
|
|
- ptr(cutlass.Uint32, inter_u32.view(torch.uint32)),
|
|
|
- ptr(cutlass.Uint8, w2_u8),
|
|
|
- ptr(cutlass.Uint8, w2_scale_u8.view(torch.uint8)),
|
|
|
- ptr(cutlass.Float32, w2_global_scale),
|
|
|
- ptr(ids_dtype, topk_ids),
|
|
|
- ptr(cutlass.Float32, topk_weights),
|
|
|
- ptr(cutlass.BFloat16, output),
|
|
|
- barrier_count,
|
|
|
- barrier_epoch,
|
|
|
- Int32(m),
|
|
|
- Int32(direct_launch.grid_x),
|
|
|
- cuda.CUstream(stream_int),
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-@torch.library.custom_op(
|
|
|
- "b12x::w4a16_small_m_direct_launch",
|
|
|
- mutates_args="unknown",
|
|
|
-)
|
|
|
-def _w4a16_small_m_direct_launch_op(
|
|
|
- a_input: torch.Tensor,
|
|
|
- w13_u8: torch.Tensor,
|
|
|
- w13_scale_u8: torch.Tensor,
|
|
|
- w13_global_scale: torch.Tensor,
|
|
|
- w2_global_scale: torch.Tensor,
|
|
|
- inter_u32: torch.Tensor,
|
|
|
- w2_u8: torch.Tensor,
|
|
|
- w2_scale_u8: torch.Tensor,
|
|
|
- topk_ids: torch.Tensor,
|
|
|
- topk_weights: torch.Tensor,
|
|
|
- output: torch.Tensor,
|
|
|
- barrier_count: torch.Tensor,
|
|
|
- barrier_epoch: torch.Tensor,
|
|
|
- m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- num_experts: int,
|
|
|
- topk: int,
|
|
|
- activation: str,
|
|
|
- fast_math: bool,
|
|
|
- scale_format: str,
|
|
|
- has_swiglu_limit: bool,
|
|
|
- swiglu_limit_value: float,
|
|
|
- swiglu_alpha: float,
|
|
|
- swiglu_beta: float,
|
|
|
- w13_layout: str,
|
|
|
- stream_int: int,
|
|
|
-) -> None:
|
|
|
- _w4a16_small_m_direct_launch_flat(
|
|
|
- a_input=a_input,
|
|
|
- w13_u8=w13_u8,
|
|
|
- w13_scale_u8=w13_scale_u8,
|
|
|
- w13_global_scale=w13_global_scale,
|
|
|
- w2_global_scale=w2_global_scale,
|
|
|
- inter_u32=inter_u32,
|
|
|
- w2_u8=w2_u8,
|
|
|
- w2_scale_u8=w2_scale_u8,
|
|
|
- topk_ids=topk_ids,
|
|
|
- topk_weights=topk_weights,
|
|
|
- output=output,
|
|
|
- barrier_count=barrier_count,
|
|
|
- barrier_epoch=barrier_epoch,
|
|
|
- m=m,
|
|
|
- hidden_size=hidden_size,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- num_experts=num_experts,
|
|
|
- topk=topk,
|
|
|
- activation=activation,
|
|
|
- fast_math=fast_math,
|
|
|
- scale_format=scale_format,
|
|
|
- has_swiglu_limit=has_swiglu_limit,
|
|
|
- swiglu_limit_value=swiglu_limit_value,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- w13_layout=w13_layout,
|
|
|
- stream_int=stream_int,
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-@_w4a16_small_m_direct_launch_op.register_fake
|
|
|
-def _w4a16_small_m_direct_launch_fake(
|
|
|
- a_input: torch.Tensor,
|
|
|
- w13_u8: torch.Tensor,
|
|
|
- w13_scale_u8: torch.Tensor,
|
|
|
- w13_global_scale: torch.Tensor,
|
|
|
- w2_global_scale: torch.Tensor,
|
|
|
- inter_u32: torch.Tensor,
|
|
|
- w2_u8: torch.Tensor,
|
|
|
- w2_scale_u8: torch.Tensor,
|
|
|
- topk_ids: torch.Tensor,
|
|
|
- topk_weights: torch.Tensor,
|
|
|
- output: torch.Tensor,
|
|
|
- barrier_count: torch.Tensor,
|
|
|
- barrier_epoch: torch.Tensor,
|
|
|
- m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- num_experts: int,
|
|
|
- topk: int,
|
|
|
- activation: str,
|
|
|
- fast_math: bool,
|
|
|
- scale_format: str,
|
|
|
- has_swiglu_limit: bool,
|
|
|
- swiglu_limit_value: float,
|
|
|
- swiglu_alpha: float,
|
|
|
- swiglu_beta: float,
|
|
|
- w13_layout: str,
|
|
|
- stream_int: int,
|
|
|
-) -> None:
|
|
|
- return None
|
|
|
-
|
|
|
-
|
|
|
-def _w4a16_fused_moe_launch_flat(
|
|
|
- a_input: torch.Tensor,
|
|
|
- w13_arg: torch.Tensor,
|
|
|
- w2_arg: torch.Tensor,
|
|
|
- fc1_out: torch.Tensor,
|
|
|
- activated: torch.Tensor,
|
|
|
- fc2_out: torch.Tensor,
|
|
|
- w13_scale_i32: torch.Tensor,
|
|
|
- w2_scale_i32: torch.Tensor,
|
|
|
- w13_global_scale: torch.Tensor,
|
|
|
- w2_global_scale: torch.Tensor,
|
|
|
- packed_route_indices: torch.Tensor,
|
|
|
- block_expert_ids: torch.Tensor,
|
|
|
- packed_route_count: torch.Tensor,
|
|
|
- activation_amax: torch.Tensor | None,
|
|
|
- layer_idx: int,
|
|
|
- topk_weights: torch.Tensor,
|
|
|
- fc1_scratch: torch.Tensor,
|
|
|
- fc2_scratch: torch.Tensor,
|
|
|
- workspace: torch.Tensor,
|
|
|
- m: int,
|
|
|
- size_m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- num_experts: int,
|
|
|
- topk: int,
|
|
|
- activation: str,
|
|
|
- apply_router_weight_on_input: bool,
|
|
|
- zero_fc2_output: bool,
|
|
|
- moe_block_size: int,
|
|
|
- max_m_blocks: int,
|
|
|
- element_dtype: str,
|
|
|
- fast_math: bool,
|
|
|
- sms: int,
|
|
|
- max_shared_mem: int,
|
|
|
- has_swiglu_limit: bool,
|
|
|
- swiglu_limit_value: float,
|
|
|
- swiglu_alpha: float,
|
|
|
- swiglu_beta: float,
|
|
|
- weight_layout: str,
|
|
|
- scale_format: str,
|
|
|
- w13_layout: str,
|
|
|
- direct_topk_routes: bool,
|
|
|
- tc_decode_fused_sum: bool,
|
|
|
- collect_activation_amax: bool,
|
|
|
- stream_int: int,
|
|
|
-) -> None:
|
|
|
- swiglu_limit = float(swiglu_limit_value) if has_swiglu_limit else None
|
|
|
- collect_activation_amax = bool(collect_activation_amax)
|
|
|
- if collect_activation_amax and activation_amax is None:
|
|
|
- raise ValueError("activation_amax is required for calibrated W4A16 launch")
|
|
|
- activation_amax_arg = (
|
|
|
- activation_amax.view(-1) if activation_amax is not None else w13_global_scale
|
|
|
- )
|
|
|
- fused = compile_w4a16_fused_moe(
|
|
|
- size_m=size_m,
|
|
|
- hidden_size=hidden_size,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- num_experts=num_experts,
|
|
|
- top_k=topk,
|
|
|
- activation=activation,
|
|
|
- apply_router_weight_on_input=bool(apply_router_weight_on_input),
|
|
|
- zero_fc2_output=bool(zero_fc2_output),
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- max_m_blocks=max_m_blocks,
|
|
|
- element_dtype=element_dtype,
|
|
|
- fast_math=bool(fast_math),
|
|
|
- sms=sms,
|
|
|
- max_shared_mem=max_shared_mem,
|
|
|
- swiglu_limit=swiglu_limit,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- weight_layout=weight_layout,
|
|
|
- scale_format=scale_format,
|
|
|
- w13_layout=w13_layout,
|
|
|
- direct_topk_routes=bool(direct_topk_routes),
|
|
|
- tc_decode_fused_sum=bool(tc_decode_fused_sum),
|
|
|
- collect_activation_amax=collect_activation_amax,
|
|
|
- )
|
|
|
- fused.compiled(
|
|
|
- make_ptr(
|
|
|
- _cutlass_element_dtype(element_dtype),
|
|
|
- a_input.data_ptr(),
|
|
|
- cute.AddressSpace.gmem,
|
|
|
- assumed_align=16,
|
|
|
- ),
|
|
|
- w13_arg,
|
|
|
- w2_arg,
|
|
|
- fc1_out,
|
|
|
- activated,
|
|
|
- fc2_out,
|
|
|
- w13_scale_i32,
|
|
|
- w2_scale_i32,
|
|
|
- w13_global_scale,
|
|
|
- w2_global_scale,
|
|
|
- packed_route_indices,
|
|
|
- block_expert_ids,
|
|
|
- packed_route_count,
|
|
|
- activation_amax_arg,
|
|
|
- int(layer_idx),
|
|
|
- make_ptr(
|
|
|
- cutlass.Float32,
|
|
|
- topk_weights.data_ptr(),
|
|
|
- cute.AddressSpace.gmem,
|
|
|
- assumed_align=4,
|
|
|
- ),
|
|
|
- fc1_scratch,
|
|
|
- fc2_scratch,
|
|
|
- workspace,
|
|
|
- m,
|
|
|
- _w4a16_fused_persistent_grid_x(
|
|
|
- fused=fused,
|
|
|
- m=m,
|
|
|
- topk=topk,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- activation=activation,
|
|
|
- direct_topk_routes=bool(direct_topk_routes),
|
|
|
- sms=sms,
|
|
|
- ),
|
|
|
- cuda.CUstream(stream_int),
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-def _w4a16_fused_persistent_grid_x(
|
|
|
- *,
|
|
|
- fused: W4A16FusedMoeCompileResult,
|
|
|
- m: int,
|
|
|
- topk: int,
|
|
|
- intermediate_size: int,
|
|
|
- activation: str,
|
|
|
- direct_topk_routes: bool,
|
|
|
- sms: int,
|
|
|
-) -> int:
|
|
|
- """Right-size the persistent grid for the fused FC1+FC2 launch.
|
|
|
-
|
|
|
- The default over-subscribes the cooperative grid (sms*blocks_per_sm CTAs)
|
|
|
- and forces the larger GEMM into the cross-CTA split-K tail (lock-serialized
|
|
|
- finalize + c_tmp gmem round-trip) whenever its mn-tile count exceeds the
|
|
|
- grid. For the small-M direct-topk decode the host knows the exact FC1
|
|
|
- mn-tile count, so pick the fewest full persistent waves that fit the
|
|
|
- co-residency cap and set grid_x to that wave's tile count. Then every CTA
|
|
|
- owns whole FC1 (and FC2, whose tile count is an integer multiple) mn-tiles
|
|
|
- over the full K -- no split-K reduction, no lock traffic -- with fewer
|
|
|
- grid-barrier participants, while staying <= the cap so the cooperative
|
|
|
- barrier never deadlocks. Falls back to the default for the route-pack path
|
|
|
- where the host cannot know route_blocks ahead of the launch.
|
|
|
- """
|
|
|
- cap = int(sms) * int(fused.blocks_per_sm)
|
|
|
- if not direct_topk_routes or m <= 0:
|
|
|
- return max(cap, 1)
|
|
|
- is_gated = is_gated_moe_activation(activation)
|
|
|
- fc1_cols = int(intermediate_size) * (2 if is_gated else 1)
|
|
|
- fc1_tile_n = int(getattr(fused, "fc1_tile_n", 0))
|
|
|
- if fc1_tile_n <= 0 or fc1_cols % fc1_tile_n != 0:
|
|
|
- return max(cap, 1)
|
|
|
- n_tiles = fc1_cols // fc1_tile_n
|
|
|
- route_blocks = int(m) * int(topk)
|
|
|
- fc1_mn_tiles = route_blocks * n_tiles
|
|
|
- if fc1_mn_tiles <= 0 or cap <= 0:
|
|
|
- return max(cap, 1)
|
|
|
- waves = (fc1_mn_tiles + cap - 1) // cap
|
|
|
- if waves <= 0:
|
|
|
- return max(cap, 1)
|
|
|
- # Only commit to the right-sized grid when every wave is a whole tile cover
|
|
|
- # (no remainder => no split-K tail); otherwise keep the safe default.
|
|
|
- if fc1_mn_tiles % waves != 0:
|
|
|
- return max(cap, 1)
|
|
|
- grid_x = fc1_mn_tiles // waves
|
|
|
- if grid_x < 1 or grid_x > cap:
|
|
|
- return max(cap, 1)
|
|
|
- return grid_x
|
|
|
-
|
|
|
-
|
|
|
-@torch.library.custom_op(
|
|
|
- "b12x::w4a16_fused_moe_launch",
|
|
|
- mutates_args="unknown",
|
|
|
-)
|
|
|
-def _w4a16_fused_moe_launch_op(
|
|
|
- a_input: torch.Tensor,
|
|
|
- w13_arg: torch.Tensor,
|
|
|
- w2_arg: torch.Tensor,
|
|
|
- fc1_out: torch.Tensor,
|
|
|
- activated: torch.Tensor,
|
|
|
- fc2_out: torch.Tensor,
|
|
|
- w13_scale_i32: torch.Tensor,
|
|
|
- w2_scale_i32: torch.Tensor,
|
|
|
- w13_global_scale: torch.Tensor,
|
|
|
- w2_global_scale: torch.Tensor,
|
|
|
- packed_route_indices: torch.Tensor,
|
|
|
- block_expert_ids: torch.Tensor,
|
|
|
- packed_route_count: torch.Tensor,
|
|
|
- topk_weights: torch.Tensor,
|
|
|
- fc1_scratch: torch.Tensor,
|
|
|
- fc2_scratch: torch.Tensor,
|
|
|
- workspace: torch.Tensor,
|
|
|
- m: int,
|
|
|
- size_m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- num_experts: int,
|
|
|
- topk: int,
|
|
|
- activation: str,
|
|
|
- apply_router_weight_on_input: bool,
|
|
|
- zero_fc2_output: bool,
|
|
|
- moe_block_size: int,
|
|
|
- max_m_blocks: int,
|
|
|
- element_dtype: str,
|
|
|
- fast_math: bool,
|
|
|
- sms: int,
|
|
|
- max_shared_mem: int,
|
|
|
- has_swiglu_limit: bool,
|
|
|
- swiglu_limit_value: float,
|
|
|
- swiglu_alpha: float,
|
|
|
- swiglu_beta: float,
|
|
|
- weight_layout: str,
|
|
|
- scale_format: str,
|
|
|
- w13_layout: str,
|
|
|
- direct_topk_routes: bool,
|
|
|
- tc_decode_fused_sum: bool,
|
|
|
- stream_int: int,
|
|
|
-) -> None:
|
|
|
- _w4a16_fused_moe_launch_flat(
|
|
|
- a_input=a_input,
|
|
|
- w13_arg=w13_arg,
|
|
|
- w2_arg=w2_arg,
|
|
|
- fc1_out=fc1_out,
|
|
|
- activated=activated,
|
|
|
- fc2_out=fc2_out,
|
|
|
- w13_scale_i32=w13_scale_i32,
|
|
|
- w2_scale_i32=w2_scale_i32,
|
|
|
- w13_global_scale=w13_global_scale,
|
|
|
- w2_global_scale=w2_global_scale,
|
|
|
- packed_route_indices=packed_route_indices,
|
|
|
- block_expert_ids=block_expert_ids,
|
|
|
- packed_route_count=packed_route_count,
|
|
|
- activation_amax=None,
|
|
|
- layer_idx=0,
|
|
|
- topk_weights=topk_weights,
|
|
|
- fc1_scratch=fc1_scratch,
|
|
|
- fc2_scratch=fc2_scratch,
|
|
|
- workspace=workspace,
|
|
|
- m=m,
|
|
|
- size_m=size_m,
|
|
|
- hidden_size=hidden_size,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- num_experts=num_experts,
|
|
|
- topk=topk,
|
|
|
- activation=activation,
|
|
|
- apply_router_weight_on_input=apply_router_weight_on_input,
|
|
|
- zero_fc2_output=zero_fc2_output,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- max_m_blocks=max_m_blocks,
|
|
|
- element_dtype=element_dtype,
|
|
|
- fast_math=fast_math,
|
|
|
- sms=sms,
|
|
|
- max_shared_mem=max_shared_mem,
|
|
|
- has_swiglu_limit=has_swiglu_limit,
|
|
|
- swiglu_limit_value=swiglu_limit_value,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- weight_layout=weight_layout,
|
|
|
- scale_format=scale_format,
|
|
|
- w13_layout=w13_layout,
|
|
|
- direct_topk_routes=direct_topk_routes,
|
|
|
- tc_decode_fused_sum=tc_decode_fused_sum,
|
|
|
- collect_activation_amax=False,
|
|
|
- stream_int=stream_int,
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-@_w4a16_fused_moe_launch_op.register_fake
|
|
|
-def _w4a16_fused_moe_launch_fake(
|
|
|
- a_input: torch.Tensor,
|
|
|
- w13_arg: torch.Tensor,
|
|
|
- w2_arg: torch.Tensor,
|
|
|
- fc1_out: torch.Tensor,
|
|
|
- activated: torch.Tensor,
|
|
|
- fc2_out: torch.Tensor,
|
|
|
- w13_scale_i32: torch.Tensor,
|
|
|
- w2_scale_i32: torch.Tensor,
|
|
|
- w13_global_scale: torch.Tensor,
|
|
|
- w2_global_scale: torch.Tensor,
|
|
|
- packed_route_indices: torch.Tensor,
|
|
|
- block_expert_ids: torch.Tensor,
|
|
|
- packed_route_count: torch.Tensor,
|
|
|
- topk_weights: torch.Tensor,
|
|
|
- fc1_scratch: torch.Tensor,
|
|
|
- fc2_scratch: torch.Tensor,
|
|
|
- workspace: torch.Tensor,
|
|
|
- m: int,
|
|
|
- size_m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- num_experts: int,
|
|
|
- topk: int,
|
|
|
- activation: str,
|
|
|
- apply_router_weight_on_input: bool,
|
|
|
- zero_fc2_output: bool,
|
|
|
- moe_block_size: int,
|
|
|
- max_m_blocks: int,
|
|
|
- element_dtype: str,
|
|
|
- fast_math: bool,
|
|
|
- sms: int,
|
|
|
- max_shared_mem: int,
|
|
|
- has_swiglu_limit: bool,
|
|
|
- swiglu_limit_value: float,
|
|
|
- swiglu_alpha: float,
|
|
|
- swiglu_beta: float,
|
|
|
- weight_layout: str,
|
|
|
- scale_format: str,
|
|
|
- w13_layout: str,
|
|
|
- direct_topk_routes: bool,
|
|
|
- tc_decode_fused_sum: bool,
|
|
|
- stream_int: int,
|
|
|
-) -> None:
|
|
|
- return None
|
|
|
-
|
|
|
-
|
|
|
-@torch.library.custom_op(
|
|
|
- "b12x::w4a16_fused_moe_calibrated_launch",
|
|
|
- mutates_args="unknown",
|
|
|
-)
|
|
|
-def _w4a16_fused_moe_calibrated_launch_op(
|
|
|
- a_input: torch.Tensor,
|
|
|
- w13_arg: torch.Tensor,
|
|
|
- w2_arg: torch.Tensor,
|
|
|
- fc1_out: torch.Tensor,
|
|
|
- activated: torch.Tensor,
|
|
|
- fc2_out: torch.Tensor,
|
|
|
- w13_scale_i32: torch.Tensor,
|
|
|
- w2_scale_i32: torch.Tensor,
|
|
|
- w13_global_scale: torch.Tensor,
|
|
|
- w2_global_scale: torch.Tensor,
|
|
|
- packed_route_indices: torch.Tensor,
|
|
|
- block_expert_ids: torch.Tensor,
|
|
|
- packed_route_count: torch.Tensor,
|
|
|
- activation_amax: torch.Tensor,
|
|
|
- layer_idx: int,
|
|
|
- topk_weights: torch.Tensor,
|
|
|
- fc1_scratch: torch.Tensor,
|
|
|
- fc2_scratch: torch.Tensor,
|
|
|
- workspace: torch.Tensor,
|
|
|
- m: int,
|
|
|
- size_m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- num_experts: int,
|
|
|
- topk: int,
|
|
|
- activation: str,
|
|
|
- apply_router_weight_on_input: bool,
|
|
|
- zero_fc2_output: bool,
|
|
|
- moe_block_size: int,
|
|
|
- max_m_blocks: int,
|
|
|
- element_dtype: str,
|
|
|
- fast_math: bool,
|
|
|
- sms: int,
|
|
|
- max_shared_mem: int,
|
|
|
- has_swiglu_limit: bool,
|
|
|
- swiglu_limit_value: float,
|
|
|
- swiglu_alpha: float,
|
|
|
- swiglu_beta: float,
|
|
|
- weight_layout: str,
|
|
|
- scale_format: str,
|
|
|
- w13_layout: str,
|
|
|
- stream_int: int,
|
|
|
-) -> None:
|
|
|
- _w4a16_fused_moe_launch_flat(
|
|
|
- a_input=a_input,
|
|
|
- w13_arg=w13_arg,
|
|
|
- w2_arg=w2_arg,
|
|
|
- fc1_out=fc1_out,
|
|
|
- activated=activated,
|
|
|
- fc2_out=fc2_out,
|
|
|
- w13_scale_i32=w13_scale_i32,
|
|
|
- w2_scale_i32=w2_scale_i32,
|
|
|
- w13_global_scale=w13_global_scale,
|
|
|
- w2_global_scale=w2_global_scale,
|
|
|
- packed_route_indices=packed_route_indices,
|
|
|
- block_expert_ids=block_expert_ids,
|
|
|
- packed_route_count=packed_route_count,
|
|
|
- activation_amax=activation_amax,
|
|
|
- layer_idx=layer_idx,
|
|
|
- topk_weights=topk_weights,
|
|
|
- fc1_scratch=fc1_scratch,
|
|
|
- fc2_scratch=fc2_scratch,
|
|
|
- workspace=workspace,
|
|
|
- m=m,
|
|
|
- size_m=size_m,
|
|
|
- hidden_size=hidden_size,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- num_experts=num_experts,
|
|
|
- topk=topk,
|
|
|
- activation=activation,
|
|
|
- apply_router_weight_on_input=apply_router_weight_on_input,
|
|
|
- zero_fc2_output=zero_fc2_output,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- max_m_blocks=max_m_blocks,
|
|
|
- element_dtype=element_dtype,
|
|
|
- fast_math=fast_math,
|
|
|
- sms=sms,
|
|
|
- max_shared_mem=max_shared_mem,
|
|
|
- has_swiglu_limit=has_swiglu_limit,
|
|
|
- swiglu_limit_value=swiglu_limit_value,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- weight_layout=weight_layout,
|
|
|
- scale_format=scale_format,
|
|
|
- w13_layout=w13_layout,
|
|
|
- direct_topk_routes=False,
|
|
|
- tc_decode_fused_sum=False,
|
|
|
- collect_activation_amax=True,
|
|
|
- stream_int=stream_int,
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-@_w4a16_fused_moe_calibrated_launch_op.register_fake
|
|
|
-def _w4a16_fused_moe_calibrated_launch_fake(
|
|
|
- a_input: torch.Tensor,
|
|
|
- w13_arg: torch.Tensor,
|
|
|
- w2_arg: torch.Tensor,
|
|
|
- fc1_out: torch.Tensor,
|
|
|
- activated: torch.Tensor,
|
|
|
- fc2_out: torch.Tensor,
|
|
|
- w13_scale_i32: torch.Tensor,
|
|
|
- w2_scale_i32: torch.Tensor,
|
|
|
- w13_global_scale: torch.Tensor,
|
|
|
- w2_global_scale: torch.Tensor,
|
|
|
- packed_route_indices: torch.Tensor,
|
|
|
- block_expert_ids: torch.Tensor,
|
|
|
- packed_route_count: torch.Tensor,
|
|
|
- activation_amax: torch.Tensor,
|
|
|
- layer_idx: int,
|
|
|
- topk_weights: torch.Tensor,
|
|
|
- fc1_scratch: torch.Tensor,
|
|
|
- fc2_scratch: torch.Tensor,
|
|
|
- workspace: torch.Tensor,
|
|
|
- m: int,
|
|
|
- size_m: int,
|
|
|
- hidden_size: int,
|
|
|
- intermediate_size: int,
|
|
|
- num_experts: int,
|
|
|
- topk: int,
|
|
|
- activation: str,
|
|
|
- apply_router_weight_on_input: bool,
|
|
|
- zero_fc2_output: bool,
|
|
|
- moe_block_size: int,
|
|
|
- max_m_blocks: int,
|
|
|
- element_dtype: str,
|
|
|
- fast_math: bool,
|
|
|
- sms: int,
|
|
|
- max_shared_mem: int,
|
|
|
- has_swiglu_limit: bool,
|
|
|
- swiglu_limit_value: float,
|
|
|
- swiglu_alpha: float,
|
|
|
- swiglu_beta: float,
|
|
|
- weight_layout: str,
|
|
|
- scale_format: str,
|
|
|
- w13_layout: str,
|
|
|
- stream_int: int,
|
|
|
-) -> None:
|
|
|
- return None
|
|
|
-
|
|
|
-
|
|
|
-def _w4a16_topk_sum_launch_flat(
|
|
|
- fc2_out: torch.Tensor,
|
|
|
- output: torch.Tensor,
|
|
|
- m: int,
|
|
|
- topk: int,
|
|
|
- hidden_size: int,
|
|
|
- element_dtype: str,
|
|
|
- stream_int: int,
|
|
|
-) -> None:
|
|
|
- sum_kernel = compile_w4a16_topk_sum(
|
|
|
- m=m,
|
|
|
- topk=topk,
|
|
|
- hidden_size=hidden_size,
|
|
|
- element_dtype=element_dtype,
|
|
|
- )
|
|
|
- sum_kernel.compiled(
|
|
|
- make_ptr(
|
|
|
- _cutlass_element_dtype(element_dtype),
|
|
|
- fc2_out.data_ptr(),
|
|
|
- cute.AddressSpace.gmem,
|
|
|
- assumed_align=16,
|
|
|
- ),
|
|
|
- make_ptr(
|
|
|
- _cutlass_element_dtype(element_dtype),
|
|
|
- output.data_ptr(),
|
|
|
- cute.AddressSpace.gmem,
|
|
|
- assumed_align=16,
|
|
|
- ),
|
|
|
- m,
|
|
|
- cuda.CUstream(stream_int),
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-@torch.library.custom_op(
|
|
|
- "b12x::w4a16_topk_sum_launch",
|
|
|
- mutates_args=("output",),
|
|
|
-)
|
|
|
-def _w4a16_topk_sum_launch_op(
|
|
|
- fc2_out: torch.Tensor,
|
|
|
- output: torch.Tensor,
|
|
|
- m: int,
|
|
|
- topk: int,
|
|
|
- hidden_size: int,
|
|
|
- element_dtype: str,
|
|
|
- stream_int: int,
|
|
|
-) -> None:
|
|
|
- _w4a16_topk_sum_launch_flat(
|
|
|
- fc2_out=fc2_out,
|
|
|
- output=output,
|
|
|
- m=m,
|
|
|
- topk=topk,
|
|
|
- hidden_size=hidden_size,
|
|
|
- element_dtype=element_dtype,
|
|
|
- stream_int=stream_int,
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-@_w4a16_topk_sum_launch_op.register_fake
|
|
|
-def _w4a16_topk_sum_launch_fake(
|
|
|
- fc2_out: torch.Tensor,
|
|
|
- output: torch.Tensor,
|
|
|
- m: int,
|
|
|
- topk: int,
|
|
|
- hidden_size: int,
|
|
|
- element_dtype: str,
|
|
|
- stream_int: int,
|
|
|
-) -> None:
|
|
|
- return None
|
|
|
-
|
|
|
-
|
|
|
-def _get_c_tmp(
|
|
|
- elements: int,
|
|
|
- *,
|
|
|
- device: torch.device,
|
|
|
- scratch: torch.Tensor | None = None,
|
|
|
-) -> torch.Tensor:
|
|
|
- if scratch is not None:
|
|
|
- if scratch.dtype != torch.float32:
|
|
|
- raise TypeError("W4A16 c_tmp scratch buffers must be torch.float32")
|
|
|
- if scratch.device != device:
|
|
|
- raise ValueError(f"W4A16 c_tmp scratch buffers must be on {device}")
|
|
|
- if not scratch.is_contiguous():
|
|
|
- raise ValueError("W4A16 c_tmp scratch buffers must be contiguous")
|
|
|
- if int(scratch.numel()) >= int(elements):
|
|
|
- return scratch[: int(elements)]
|
|
|
- if torch.cuda.is_current_stream_capturing():
|
|
|
- raise RuntimeError(
|
|
|
- "W4A16 GEMM scratch is not initialized for CUDA graph capture; "
|
|
|
- "provide a preallocated fc*_c_tmp workspace with sufficient capacity"
|
|
|
- )
|
|
|
- return torch.empty((elements,), dtype=torch.float32, device=device)
|
|
|
-
|
|
|
-
|
|
|
-def _validate_topk_ids(
|
|
|
- topk_ids: torch.Tensor,
|
|
|
- *,
|
|
|
- require_cuda: bool,
|
|
|
- require_contiguous: bool = True,
|
|
|
-) -> None:
|
|
|
- if topk_ids.dtype not in (torch.int32, torch.int64):
|
|
|
- raise TypeError("topk_ids must be torch.int32 or torch.int64")
|
|
|
- if require_cuda and not topk_ids.is_cuda:
|
|
|
- raise ValueError("topk_ids must be a CUDA tensor")
|
|
|
- if require_contiguous and not topk_ids.is_contiguous():
|
|
|
- raise ValueError("topk_ids must be contiguous")
|
|
|
-
|
|
|
-
|
|
|
-def _validate_expert_map(
|
|
|
- expert_map: torch.Tensor | None,
|
|
|
- *,
|
|
|
- device: torch.device | None = None,
|
|
|
- exact_num_experts: int | None = None,
|
|
|
-) -> None:
|
|
|
- if expert_map is None:
|
|
|
- return
|
|
|
- if expert_map.dtype != torch.int32:
|
|
|
- raise TypeError("expert_map must be torch.int32")
|
|
|
- if device is None:
|
|
|
- if not expert_map.is_cuda:
|
|
|
- raise ValueError("expert_map must be a CUDA tensor")
|
|
|
- elif expert_map.device != device:
|
|
|
- raise ValueError("expert_map must be on the same device as a_input")
|
|
|
- if exact_num_experts is None:
|
|
|
- if expert_map.ndim != 1 or not expert_map.is_contiguous():
|
|
|
- raise ValueError("expert_map must be a contiguous rank-1 tensor")
|
|
|
- return
|
|
|
- if not expert_map.is_contiguous():
|
|
|
- raise ValueError("expert_map must be contiguous")
|
|
|
- if expert_map.ndim != 1 or int(expert_map.numel()) != int(exact_num_experts):
|
|
|
- raise ValueError(
|
|
|
- f"expert_map must have shape {(int(exact_num_experts),)}, got {tuple(expert_map.shape)}"
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-def _validate_activation_amax(
|
|
|
- activation_amax: torch.Tensor | None,
|
|
|
- *,
|
|
|
- layer_idx: int | None,
|
|
|
- num_experts: int,
|
|
|
- device: torch.device,
|
|
|
-) -> int | None:
|
|
|
- if activation_amax is None:
|
|
|
- if layer_idx is not None:
|
|
|
- raise ValueError("layer_idx requires activation_amax")
|
|
|
- return None
|
|
|
- if layer_idx is None:
|
|
|
- raise ValueError("layer_idx is required when activation_amax is provided")
|
|
|
- if activation_amax.dtype != torch.float32:
|
|
|
- raise TypeError("activation_amax must be torch.float32")
|
|
|
- if activation_amax.device != device:
|
|
|
- raise ValueError("activation_amax must be on the same device as a_input")
|
|
|
- if not activation_amax.is_cuda:
|
|
|
- raise ValueError("activation_amax must be a CUDA tensor")
|
|
|
- if not activation_amax.is_contiguous():
|
|
|
- raise ValueError("activation_amax must be contiguous")
|
|
|
- if activation_amax.ndim != 3 or int(activation_amax.shape[2]) != 2:
|
|
|
- raise ValueError(
|
|
|
- "activation_amax must have shape [num_layers, num_experts, 2]"
|
|
|
- )
|
|
|
- if int(activation_amax.shape[1]) < int(num_experts):
|
|
|
- raise ValueError(
|
|
|
- "activation_amax expert dimension is smaller than the local expert count"
|
|
|
- )
|
|
|
- layer = int(layer_idx)
|
|
|
- if layer < 0 or layer >= int(activation_amax.shape[0]):
|
|
|
- raise ValueError(
|
|
|
- f"layer_idx {layer} is out of bounds for activation_amax with "
|
|
|
- f"{int(activation_amax.shape[0])} layers"
|
|
|
- )
|
|
|
- return layer
|
|
|
-
|
|
|
-
|
|
|
-def _compile_w4a16_gemm_launch(
|
|
|
- *,
|
|
|
- size_m: int,
|
|
|
- size_n: int,
|
|
|
- size_k: int,
|
|
|
- num_experts: int,
|
|
|
- top_k: int,
|
|
|
- mul_topk_weights: bool,
|
|
|
- moe_block_size: int,
|
|
|
- max_m_blocks: int,
|
|
|
- element_dtype: str,
|
|
|
- packed_route_indices: torch.Tensor,
|
|
|
- sms: int,
|
|
|
- max_shared_mem: int,
|
|
|
- device: torch.device,
|
|
|
- c_tmp: torch.Tensor | None = None,
|
|
|
- scale_format: str = "e4m3_k16",
|
|
|
-) -> _W4A16GemmLaunch:
|
|
|
- tile_k, tile_n, _, _ = _select_tile_config(
|
|
|
- problem_m=size_m,
|
|
|
- problem_n=size_n,
|
|
|
- problem_k=size_k,
|
|
|
- top_k=top_k,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- sms=sms,
|
|
|
- max_shared_mem=max_shared_mem,
|
|
|
- scale_format=scale_format,
|
|
|
- )
|
|
|
- kernel = compile_w4a16_gemm(
|
|
|
- size_m=size_m,
|
|
|
- size_n=size_n,
|
|
|
- size_k=size_k,
|
|
|
- num_experts=num_experts,
|
|
|
- top_k=top_k,
|
|
|
- mul_topk_weights=mul_topk_weights,
|
|
|
- tile_n=tile_n,
|
|
|
- tile_k=tile_k,
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- max_m_blocks=max_m_blocks,
|
|
|
- element_dtype=element_dtype,
|
|
|
- scale_format=scale_format,
|
|
|
- )
|
|
|
- c_tmp = _get_c_tmp(
|
|
|
- packed_gemm_scratch_elements(
|
|
|
- size_n=size_n,
|
|
|
- route_slots=int(packed_route_indices.numel()),
|
|
|
- moe_block_size=moe_block_size,
|
|
|
- sms=sms,
|
|
|
- ),
|
|
|
- device=device,
|
|
|
- scratch=c_tmp,
|
|
|
- )
|
|
|
- return _W4A16GemmLaunch(kernel=kernel, c_tmp=c_tmp)
|
|
|
-
|
|
|
-
|
|
|
-def pack_topk_routes_by_expert(
|
|
|
- topk_ids: torch.Tensor,
|
|
|
- block_size: int,
|
|
|
- num_experts: int,
|
|
|
- *,
|
|
|
- expert_map: torch.Tensor | None = None,
|
|
|
- packed_route_indices: torch.Tensor | None = None,
|
|
|
- block_expert_ids: torch.Tensor | None = None,
|
|
|
- packed_route_count: torch.Tensor | None = None,
|
|
|
- expert_offsets: torch.Tensor | None = None,
|
|
|
- stream: cuda.CUstream | None = None,
|
|
|
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
|
- """Group top-k routes by expert and pad each group to the GEMM M-block size."""
|
|
|
- _validate_topk_ids(topk_ids, require_cuda=True)
|
|
|
- if block_size <= 0:
|
|
|
- raise ValueError(f"block_size must be positive, got {block_size}")
|
|
|
- if num_experts <= 0:
|
|
|
- raise ValueError(f"num_experts must be positive, got {num_experts}")
|
|
|
- _validate_expert_map(expert_map, exact_num_experts=int(num_experts))
|
|
|
- del stream
|
|
|
- return _pack_topk_routes_by_expert(
|
|
|
- topk_ids,
|
|
|
- int(block_size),
|
|
|
- int(num_experts),
|
|
|
- expert_map=expert_map,
|
|
|
- packed_route_indices=packed_route_indices,
|
|
|
- block_expert_ids=block_expert_ids,
|
|
|
- packed_route_count=packed_route_count,
|
|
|
- expert_offsets=expert_offsets,
|
|
|
- )
|
|
|
-
|
|
|
-
|
|
|
-def run_w4a16_moe(
|
|
|
- a_input: torch.Tensor,
|
|
|
- prepared,
|
|
|
- topk_weights: torch.Tensor,
|
|
|
- topk_ids: torch.Tensor,
|
|
|
- *,
|
|
|
- activation: str,
|
|
|
- intermediate_cache13: torch.Tensor,
|
|
|
- intermediate_cache2: torch.Tensor,
|
|
|
- output: torch.Tensor,
|
|
|
- fc1_c_tmp: torch.Tensor | None = None,
|
|
|
- fc2_c_tmp: torch.Tensor | None = None,
|
|
|
- packed_route_indices: torch.Tensor | None = None,
|
|
|
- block_expert_ids: torch.Tensor | None = None,
|
|
|
- packed_route_count: torch.Tensor | None = None,
|
|
|
- expert_offsets: torch.Tensor | None = None,
|
|
|
- expert_map: torch.Tensor | None = None,
|
|
|
- activation_amax: torch.Tensor | None = None,
|
|
|
- layer_idx: int | None = None,
|
|
|
- apply_router_weight_on_input: bool = False,
|
|
|
- fast_math: bool = True,
|
|
|
- swiglu_limit: float | None = None,
|
|
|
- swiglu_alpha: float | None = None,
|
|
|
- swiglu_beta: float | None = None,
|
|
|
- fused_launch: W4A16FusedMoeCompileResult | None = None,
|
|
|
- topk_sum_launch: W4A16TopKSumCompileResult | None = None,
|
|
|
- stream: cuda.CUstream | None = None,
|
|
|
-) -> torch.Tensor:
|
|
|
- activation = normalize_moe_activation(activation)
|
|
|
- is_gated = validate_activation(activation)
|
|
|
- swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params(
|
|
|
- activation,
|
|
|
- swiglu_limit,
|
|
|
- swiglu_alpha,
|
|
|
- swiglu_beta,
|
|
|
- )
|
|
|
- element_dtype = _normalize_element_dtype(a_input.dtype)
|
|
|
- if output.dtype != a_input.dtype:
|
|
|
- raise TypeError(f"output must have dtype {a_input.dtype}, got {output.dtype}")
|
|
|
- prepared_dtype = getattr(prepared, "params_dtype", a_input.dtype)
|
|
|
- if prepared_dtype != a_input.dtype:
|
|
|
- raise TypeError(
|
|
|
- f"prepared weights were built for {prepared_dtype}, but a_input has dtype {a_input.dtype}"
|
|
|
- )
|
|
|
- weight_layout = getattr(prepared, "weight_layout", "packed")
|
|
|
- if weight_layout not in _WEIGHT_LAYOUTS:
|
|
|
- raise ValueError(f"unsupported W4A16 weight_layout {weight_layout!r}")
|
|
|
- scale_format = _normalize_scale_format(
|
|
|
- getattr(prepared, "scale_format", None)
|
|
|
- or (
|
|
|
- "e8m0_k32"
|
|
|
- if getattr(prepared, "source_format", "") == "fp4_e8m0_k32"
|
|
|
- else "e4m3_k16"
|
|
|
- )
|
|
|
- )
|
|
|
- w13_layout = getattr(
|
|
|
- prepared,
|
|
|
- "w13_layout",
|
|
|
- "w13" if weight_layout == "modelopt" else "packed",
|
|
|
- )
|
|
|
- if weight_layout == "modelopt":
|
|
|
- if w13_layout not in _MODEL_OPT_W13_LAYOUTS:
|
|
|
- raise ValueError(f"unsupported W4A16 w13_layout {w13_layout!r}")
|
|
|
- else:
|
|
|
- w13_layout = "packed"
|
|
|
- if topk_weights.dtype != torch.float32:
|
|
|
- raise TypeError("topk_weights must be torch.float32")
|
|
|
- _validate_topk_ids(topk_ids, require_cuda=False, require_contiguous=False)
|
|
|
- if (
|
|
|
- not a_input.is_contiguous()
|
|
|
- or not topk_weights.is_contiguous()
|
|
|
- or not topk_ids.is_contiguous()
|
|
|
- ):
|
|
|
- raise ValueError("a_input, topk_weights, and topk_ids must be contiguous")
|
|
|
- _validate_expert_map(expert_map, device=a_input.device)
|
|
|
-
|
|
|
- m, hidden_size = a_input.shape
|
|
|
- topk = int(topk_ids.shape[1])
|
|
|
- if tuple(topk_weights.shape) != (m, topk):
|
|
|
- raise ValueError(f"topk_weights must have shape {(m, topk)}")
|
|
|
- if int(prepared.hidden_size) != hidden_size:
|
|
|
- raise ValueError("prepared hidden_size does not match a_input")
|
|
|
- if bool(prepared.is_gated) != is_gated:
|
|
|
- raise ValueError("prepared weights do not match activation")
|
|
|
- if tuple(output.shape) != (m, hidden_size):
|
|
|
- raise ValueError(f"output must have shape {(m, hidden_size)}")
|
|
|
- if expert_map is not None and int(expert_map.numel()) < int(prepared.num_experts):
|
|
|
- raise ValueError("expert_map cannot be shorter than the local expert count")
|
|
|
- layer_idx_int = _validate_activation_amax(
|
|
|
- activation_amax,
|
|
|
- layer_idx=layer_idx,
|
|
|
- num_experts=int(prepared.num_experts),
|
|
|
- device=a_input.device,
|
|
|
- )
|
|
|
- collect_activation_amax = activation_amax is not None
|
|
|
-
|
|
|
- route_num_experts = (
|
|
|
- int(expert_map.numel()) if expert_map is not None else int(prepared.num_experts)
|
|
|
- )
|
|
|
- if fused_launch is None:
|
|
|
- block_size_m = select_route_block_size_m(m, topk, route_num_experts)
|
|
|
- else:
|
|
|
- block_size_m = int(fused_launch.moe_block_size)
|
|
|
- if block_size_m not in _ALLOWED_ROUTED_SIZES:
|
|
|
- raise ValueError(f"unsupported W4A16 moe_block_size={block_size_m}")
|
|
|
-
|
|
|
- stream = current_cuda_stream() if stream is None else stream
|
|
|
- if (not collect_activation_amax) and _small_m_direct_supported(
|
|
|
- m=m,
|
|
|
- hidden_size=hidden_size,
|
|
|
- intermediate_size=int(prepared.intermediate_size),
|
|
|
- num_experts=int(prepared.num_experts),
|
|
|
- topk=topk,
|
|
|
- activation=activation,
|
|
|
- apply_router_weight_on_input=bool(apply_router_weight_on_input),
|
|
|
- swiglu_limit=swiglu_limit,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- element_dtype=element_dtype,
|
|
|
- weight_layout=weight_layout,
|
|
|
- w13_layout=w13_layout,
|
|
|
- scale_format=scale_format,
|
|
|
- expert_map=expert_map,
|
|
|
- ):
|
|
|
- if topk_ids.dtype not in (torch.int32, torch.int64):
|
|
|
- raise TypeError("W4A16 small-M direct path requires int32/int64 topk_ids")
|
|
|
- if not topk_ids.is_cuda:
|
|
|
- raise ValueError("W4A16 small-M direct path requires CUDA topk_ids")
|
|
|
- if not intermediate_cache2.is_contiguous() or not output.is_contiguous():
|
|
|
- raise ValueError(
|
|
|
- "W4A16 small-M direct path requires contiguous intermediate_cache2 and output"
|
|
|
- )
|
|
|
- if intermediate_cache2.dtype != a_input.dtype:
|
|
|
- raise TypeError(f"intermediate_cache2 must be {a_input.dtype}")
|
|
|
- if int(prepared.workspace.numel()) < 2:
|
|
|
- raise ValueError("prepared W4A16 workspace is too small for small-M direct")
|
|
|
- intermediate_size = int(prepared.intermediate_size)
|
|
|
- fc2_n_chunks = ((intermediate_size // 2) + 127) // 128
|
|
|
- inter_u32_per_m = fc2_n_chunks * 128 * topk
|
|
|
- inter_u32 = intermediate_cache2.view(-1).view(torch.uint32)
|
|
|
- if int(inter_u32.numel()) < m * inter_u32_per_m:
|
|
|
- raise ValueError(
|
|
|
- "intermediate_cache2 is smaller than the W4A16 small-M direct scratch "
|
|
|
- f"requirement: have_u32={int(inter_u32.numel())}, "
|
|
|
- f"need_u32={m * inter_u32_per_m}"
|
|
|
- )
|
|
|
- micro_w13_scale = getattr(prepared, "micro_w13_scale", None)
|
|
|
- micro_w2_scale = getattr(prepared, "micro_w2_scale", None)
|
|
|
- micro_w13_global = getattr(prepared, "micro_w13_global_scale", None)
|
|
|
- micro_w2_global = getattr(prepared, "micro_w2_global_scale", None)
|
|
|
- if (
|
|
|
- micro_w13_scale is None
|
|
|
- or micro_w2_scale is None
|
|
|
- or micro_w13_global is None
|
|
|
- or micro_w2_global is None
|
|
|
- ):
|
|
|
- raise RuntimeError(
|
|
|
- "W4A16 small-M direct path requires prepared micro scale metadata"
|
|
|
- )
|
|
|
- barrier_count = prepared.workspace[-2:-1]
|
|
|
- barrier_epoch = prepared.workspace[-1:]
|
|
|
- barrier_count.zero_()
|
|
|
- barrier_epoch.zero_()
|
|
|
- torch.ops.b12x.w4a16_small_m_direct_launch(
|
|
|
- a_input,
|
|
|
- prepared.w13.view(torch.uint8),
|
|
|
- micro_w13_scale,
|
|
|
- micro_w13_global,
|
|
|
- micro_w2_global,
|
|
|
- inter_u32[: m * inter_u32_per_m],
|
|
|
- prepared.w2.view(torch.uint8),
|
|
|
- micro_w2_scale,
|
|
|
- topk_ids,
|
|
|
- topk_weights,
|
|
|
- output,
|
|
|
- barrier_count,
|
|
|
- barrier_epoch,
|
|
|
- m,
|
|
|
- hidden_size,
|
|
|
- intermediate_size,
|
|
|
- int(prepared.num_experts),
|
|
|
- topk,
|
|
|
- activation,
|
|
|
- bool(fast_math),
|
|
|
- scale_format,
|
|
|
- swiglu_limit is not None,
|
|
|
- float(swiglu_limit or 0.0),
|
|
|
- float(swiglu_alpha),
|
|
|
- float(swiglu_beta),
|
|
|
- w13_layout,
|
|
|
- int(stream),
|
|
|
- )
|
|
|
- return output
|
|
|
-
|
|
|
- # TC-decode: small-M packed decode that folds the top-k sum into the FC2
|
|
|
- # store epilogue. Reuses the packed tensor-core MMA path; only the launch
|
|
|
- # scheduling/epilogue changes. Requires the packed object, bf16 gated
|
|
|
- # activation, int32 routes, no expert_map, and a runtime-compiled launch.
|
|
|
- # A preplanned launch built with the TC-decode fused-sum epilogue carries
|
|
|
- # ``tc_decode_fused_sum``; accept it through the binding path. A runtime
|
|
|
- # ``fused_launch is None`` (e.g. the standalone benchmark) compiles its own.
|
|
|
- preplanned_tc_decode = bool(getattr(fused_launch, "tc_decode_fused_sum", False))
|
|
|
- use_tc_decode = bool(
|
|
|
- (not collect_activation_amax)
|
|
|
- and (fused_launch is None or preplanned_tc_decode)
|
|
|
- and weight_layout in ("packed", "nf3_2p1")
|
|
|
- and expert_map is None
|
|
|
- and is_gated
|
|
|
- and element_dtype == "bf16"
|
|
|
- and topk_ids.dtype in (torch.int32, torch.int64)
|
|
|
- and topk_ids.is_cuda
|
|
|
- and int(m) <= _TC_DECODE_MAX_M
|
|
|
- )
|
|
|
- if use_tc_decode and topk_ids.dtype != torch.int32:
|
|
|
- # The inline direct-topk route path needs int32 route indices.
|
|
|
- topk_ids = topk_ids.to(torch.int32)
|
|
|
-
|
|
|
- direct_topk_eligible = (
|
|
|
- (not collect_activation_amax)
|
|
|
- and (m <= _MAX_DIRECT_TOPK_ROUTE_M or use_tc_decode)
|
|
|
- and weight_layout in ("packed", "nf3_2p1")
|
|
|
- and expert_map is None
|
|
|
- )
|
|
|
- use_direct_topk_routes = bool(
|
|
|
- direct_topk_eligible
|
|
|
- and topk_ids.dtype == torch.int32
|
|
|
- and (
|
|
|
- fused_launch is None
|
|
|
- or bool(getattr(fused_launch, "direct_topk_routes", False))
|
|
|
- )
|
|
|
- )
|
|
|
- if (
|
|
|
- bool(getattr(fused_launch, "direct_topk_routes", False))
|
|
|
- and not use_direct_topk_routes
|
|
|
- ):
|
|
|
- raise RuntimeError(
|
|
|
- "preplanned W4A16 direct top-k routing requires small-M packed "
|
|
|
- "int32 topk_ids without expert_map"
|
|
|
- )
|
|
|
-
|
|
|
- # TC-decode requires the inline direct-topk route path (no route-pack).
|
|
|
- use_tc_decode = bool(use_tc_decode and use_direct_topk_routes)
|
|
|
-
|
|
|
- # A preplanned TC-decode launch atomically accumulates FC2 partials into the
|
|
|
- # (pre-zeroed) output and emits no separate top-k sum. If it was selected but
|
|
|
- # the decode preconditions don't hold, running it would corrupt the output,
|
|
|
- # so fail loudly instead.
|
|
|
- if preplanned_tc_decode and not use_tc_decode:
|
|
|
- raise RuntimeError(
|
|
|
- "preplanned TC-decode W4A16 launch requires small-M packed bf16 "
|
|
|
- f"decode (m <= {_TC_DECODE_MAX_M}, cuda int32/int64 topk_ids, "
|
|
|
- "no expert_map)"
|
|
|
- )
|
|
|
-
|
|
|
- route_slots_for_scratch = int(m) * int(topk) * int(block_size_m)
|
|
|
- required_m_blocks = int(m) * int(topk) if use_direct_topk_routes else 0
|
|
|
- if fused_launch is not None and not use_direct_topk_routes:
|
|
|
- route_slots_capacity = max_packed_route_slots(
|
|
|
- int(fused_launch.size_m) * int(topk),
|
|
|
- int(block_size_m),
|
|
|
- route_num_experts,
|
|
|
- )
|
|
|
- route_blocks_capacity = (route_slots_capacity + int(block_size_m) - 1) // int(
|
|
|
- block_size_m
|
|
|
- )
|
|
|
- if packed_route_indices is not None:
|
|
|
- if int(packed_route_indices.numel()) < route_slots_capacity:
|
|
|
- raise ValueError(
|
|
|
- "packed_route_indices is smaller than the selected W4A16 launch capacity"
|
|
|
- )
|
|
|
- packed_route_indices = packed_route_indices[:route_slots_capacity]
|
|
|
- if block_expert_ids is not None:
|
|
|
- if int(block_expert_ids.numel()) < route_blocks_capacity:
|
|
|
- raise ValueError(
|
|
|
- "block_expert_ids is smaller than the selected W4A16 launch capacity"
|
|
|
- )
|
|
|
- block_expert_ids = block_expert_ids[:route_blocks_capacity]
|
|
|
- if use_direct_topk_routes:
|
|
|
- packed_route_indices = topk_ids.view(-1)
|
|
|
- if block_expert_ids is None:
|
|
|
- block_expert_ids = packed_route_indices
|
|
|
- if packed_route_count is None:
|
|
|
- packed_route_count = packed_route_indices
|
|
|
- else:
|
|
|
- packed_route_indices, block_expert_ids, packed_route_count = (
|
|
|
- pack_topk_routes_by_expert(
|
|
|
- topk_ids,
|
|
|
- block_size_m,
|
|
|
- route_num_experts,
|
|
|
- expert_map=expert_map,
|
|
|
- packed_route_indices=packed_route_indices,
|
|
|
- block_expert_ids=block_expert_ids,
|
|
|
- packed_route_count=packed_route_count,
|
|
|
- expert_offsets=expert_offsets,
|
|
|
- stream=stream,
|
|
|
- )
|
|
|
- )
|
|
|
- route_slots_for_scratch = int(packed_route_indices.numel())
|
|
|
- required_m_blocks = int(block_expert_ids.numel())
|
|
|
-
|
|
|
- props = torch.cuda.get_device_properties(a_input.device)
|
|
|
- sms = int(props.multi_processor_count)
|
|
|
- max_shared_mem = int(
|
|
|
- getattr(props, "shared_memory_per_block_optin", _DEFAULT_MAX_SHARED_MEM)
|
|
|
- )
|
|
|
- buffer_plan = plan_w4a16_buffers(
|
|
|
- prepared,
|
|
|
- m=m,
|
|
|
- topk=topk,
|
|
|
- route_num_experts=route_num_experts,
|
|
|
- sms=sms,
|
|
|
- )
|
|
|
- intermediate_size = int(prepared.intermediate_size)
|
|
|
- fc1_cols = buffer_plan.fc1_cols
|
|
|
-
|
|
|
- if intermediate_cache13.numel() < buffer_plan.intermediate_cache13_elements:
|
|
|
- raise ValueError(
|
|
|
- f"intermediate_cache13 has {intermediate_cache13.numel()} elements; "
|
|
|
- f"need at least {buffer_plan.intermediate_cache13_elements}"
|
|
|
- )
|
|
|
- if intermediate_cache2.numel() < buffer_plan.intermediate_cache2_elements:
|
|
|
- raise ValueError(
|
|
|
- f"intermediate_cache2 has {intermediate_cache2.numel()} elements; "
|
|
|
- f"need at least {buffer_plan.intermediate_cache2_elements}"
|
|
|
- )
|
|
|
- if (
|
|
|
- intermediate_cache13.dtype != a_input.dtype
|
|
|
- or intermediate_cache2.dtype != a_input.dtype
|
|
|
- ):
|
|
|
- raise TypeError(f"intermediate caches must be {a_input.dtype}")
|
|
|
- if (
|
|
|
- not intermediate_cache13.is_contiguous()
|
|
|
- or not intermediate_cache2.is_contiguous()
|
|
|
- or not output.is_contiguous()
|
|
|
- ):
|
|
|
- raise ValueError("intermediate caches and output must be contiguous")
|
|
|
-
|
|
|
- intermediate_cache13_flat = intermediate_cache13.view(-1)
|
|
|
- intermediate_cache2_flat = intermediate_cache2.view(-1)
|
|
|
-
|
|
|
- if int(prepared.workspace.numel()) < sms * 4 + 2:
|
|
|
- raise ValueError("prepared W4A16 workspace is too small for fused FC1+FC2")
|
|
|
- if fused_launch is None:
|
|
|
- fused = compile_w4a16_fused_moe(
|
|
|
- size_m=m,
|
|
|
- hidden_size=hidden_size,
|
|
|
- intermediate_size=intermediate_size,
|
|
|
- num_experts=int(prepared.num_experts),
|
|
|
- top_k=topk,
|
|
|
- activation=activation,
|
|
|
- apply_router_weight_on_input=bool(apply_router_weight_on_input),
|
|
|
- zero_fc2_output=expert_map is not None,
|
|
|
- moe_block_size=block_size_m,
|
|
|
- max_m_blocks=int(required_m_blocks),
|
|
|
- element_dtype=element_dtype,
|
|
|
- fast_math=bool(fast_math),
|
|
|
- sms=sms,
|
|
|
- max_shared_mem=max_shared_mem,
|
|
|
- swiglu_limit=swiglu_limit,
|
|
|
- swiglu_alpha=swiglu_alpha,
|
|
|
- swiglu_beta=swiglu_beta,
|
|
|
- weight_layout=weight_layout,
|
|
|
- scale_format=scale_format,
|
|
|
- w13_layout=w13_layout,
|
|
|
- direct_topk_routes=use_direct_topk_routes,
|
|
|
- tc_decode_fused_sum=use_tc_decode,
|
|
|
- collect_activation_amax=collect_activation_amax,
|
|
|
- )
|
|
|
- else:
|
|
|
- if int(fused_launch.size_m) < m:
|
|
|
- raise RuntimeError(
|
|
|
- "preplanned W4A16 fused MoE launch capacity is smaller than requested rows: "
|
|
|
- f"requested={m}, planned={int(fused_launch.size_m)}"
|
|
|
- )
|
|
|
- expected_fused = (
|
|
|
- hidden_size,
|
|
|
- intermediate_size,
|
|
|
- int(prepared.num_experts),
|
|
|
- topk,
|
|
|
- activation,
|
|
|
- bool(apply_router_weight_on_input),
|
|
|
- expert_map is not None,
|
|
|
- element_dtype,
|
|
|
- bool(fast_math),
|
|
|
- swiglu_limit,
|
|
|
- float(swiglu_alpha),
|
|
|
- float(swiglu_beta),
|
|
|
- weight_layout,
|
|
|
- scale_format,
|
|
|
- w13_layout,
|
|
|
- bool(use_direct_topk_routes),
|
|
|
- bool(collect_activation_amax),
|
|
|
- block_size_m,
|
|
|
- )
|
|
|
- actual_fused = (
|
|
|
- int(fused_launch.hidden_size),
|
|
|
- int(fused_launch.intermediate_size),
|
|
|
- int(fused_launch.num_experts),
|
|
|
- int(fused_launch.top_k),
|
|
|
- fused_launch.activation,
|
|
|
- bool(fused_launch.apply_router_weight_on_input),
|
|
|
- bool(fused_launch.zero_fc2_output),
|
|
|
- fused_launch.element_dtype,
|
|
|
- bool(fused_launch.fast_math),
|
|
|
- fused_launch.swiglu_limit,
|
|
|
- float(fused_launch.swiglu_alpha),
|
|
|
- float(fused_launch.swiglu_beta),
|
|
|
- getattr(fused_launch, "weight_layout", "packed"),
|
|
|
- getattr(fused_launch, "scale_format", "e4m3_k16"),
|
|
|
- getattr(
|
|
|
- fused_launch,
|
|
|
- "w13_layout",
|
|
|
- "w13"
|
|
|
- if getattr(fused_launch, "weight_layout", "packed") == "modelopt"
|
|
|
- else "packed",
|
|
|
- ),
|
|
|
- bool(getattr(fused_launch, "direct_topk_routes", False)),
|
|
|
- bool(getattr(fused_launch, "collect_activation_amax", False)),
|
|
|
- int(fused_launch.moe_block_size),
|
|
|
- )
|
|
|
- if actual_fused != expected_fused or int(fused_launch.max_m_blocks) < int(
|
|
|
- required_m_blocks
|
|
|
- ):
|
|
|
- raise RuntimeError(
|
|
|
- "preplanned W4A16 fused MoE launch does not match requested contract: "
|
|
|
- f"requested={expected_fused + (int(required_m_blocks),)}, "
|
|
|
- f"planned={actual_fused + (int(fused_launch.max_m_blocks),)}"
|
|
|
- )
|
|
|
- fused = fused_launch
|
|
|
- capacity_m = int(fused.size_m)
|
|
|
- capacity_routed_rows = capacity_m * topk
|
|
|
- if intermediate_cache13_flat.numel() < capacity_routed_rows * max(fc1_cols, hidden_size):
|
|
|
- raise ValueError(
|
|
|
- "intermediate_cache13 is smaller than the selected W4A16 launch capacity: "
|
|
|
- f"capacity_rows={capacity_m}, topk={topk}"
|
|
|
- )
|
|
|
- if intermediate_cache2_flat.numel() < capacity_routed_rows * intermediate_size:
|
|
|
- raise ValueError(
|
|
|
- "intermediate_cache2 is smaller than the selected W4A16 launch capacity: "
|
|
|
- f"capacity_rows={capacity_m}, topk={topk}"
|
|
|
- )
|
|
|
- fc1_out = intermediate_cache13_flat[: capacity_routed_rows * fc1_cols]
|
|
|
- activated = intermediate_cache2_flat[: capacity_routed_rows * intermediate_size]
|
|
|
- if use_tc_decode:
|
|
|
- # FC2 atomically accumulates per-route partials directly into the
|
|
|
- # per-token output, so the output is the FC2 store target and must be
|
|
|
- # pre-zeroed. The fused tc_decode kernel now zeroes the output in its
|
|
|
- # own prologue (before FC1, made visible by the existing post-FC1 grid
|
|
|
- # barrier), so the separate host-side output.zero_() launch is removed
|
|
|
- # from the decode critical path here. This drops the separate top-k-sum
|
|
|
- # launch as well.
|
|
|
- fc2_out = output.view(-1)
|
|
|
- else:
|
|
|
- fc2_out = intermediate_cache13_flat[: capacity_routed_rows * hidden_size]
|
|
|
- fc1_scratch = _get_c_tmp(
|
|
|
- packed_gemm_scratch_elements(
|
|
|
- size_n=fc1_cols,
|
|
|
- route_slots=int(route_slots_for_scratch),
|
|
|
- moe_block_size=block_size_m,
|
|
|
- sms=sms,
|
|
|
- ),
|
|
|
- device=a_input.device,
|
|
|
- scratch=fc1_c_tmp,
|
|
|
- )
|
|
|
- fc2_scratch = _get_c_tmp(
|
|
|
- packed_gemm_scratch_elements(
|
|
|
- size_n=hidden_size,
|
|
|
- route_slots=int(route_slots_for_scratch),
|
|
|
- moe_block_size=block_size_m,
|
|
|
- sms=sms,
|
|
|
- ),
|
|
|
- device=a_input.device,
|
|
|
- scratch=fc2_c_tmp,
|
|
|
- )
|
|
|
- if weight_layout == "modelopt":
|
|
|
- w13_arg = prepared.w13.view(torch.uint8).view(-1)
|
|
|
- w2_arg = prepared.w2.view(torch.uint8).view(-1)
|
|
|
- else:
|
|
|
- w13_arg = prepared.w13.view(torch.int32).view(-1)
|
|
|
- w2_arg = prepared.w2.view(torch.int32).view(-1)
|
|
|
- launch_common = (
|
|
|
- a_input,
|
|
|
- w13_arg,
|
|
|
- w2_arg,
|
|
|
- fc1_out,
|
|
|
- activated,
|
|
|
- fc2_out,
|
|
|
- prepared.w13_scale.view(torch.uint8).view(torch.int32).view(-1),
|
|
|
- prepared.w2_scale.view(torch.uint8).view(torch.int32).view(-1),
|
|
|
- prepared.w13_global_scale,
|
|
|
- prepared.w2_global_scale,
|
|
|
- packed_route_indices,
|
|
|
- block_expert_ids,
|
|
|
- packed_route_count,
|
|
|
- )
|
|
|
- launch_tail = (
|
|
|
- topk_weights,
|
|
|
- fc1_scratch,
|
|
|
- fc2_scratch,
|
|
|
- prepared.workspace,
|
|
|
- m,
|
|
|
- capacity_m,
|
|
|
- hidden_size,
|
|
|
- intermediate_size,
|
|
|
- int(prepared.num_experts),
|
|
|
- topk,
|
|
|
- activation,
|
|
|
- bool(apply_router_weight_on_input),
|
|
|
- expert_map is not None,
|
|
|
- block_size_m,
|
|
|
- int(fused.max_m_blocks),
|
|
|
- element_dtype,
|
|
|
- bool(fast_math),
|
|
|
- sms,
|
|
|
- max_shared_mem,
|
|
|
- swiglu_limit is not None,
|
|
|
- float(swiglu_limit or 0.0),
|
|
|
- float(swiglu_alpha),
|
|
|
- float(swiglu_beta),
|
|
|
- weight_layout,
|
|
|
- scale_format,
|
|
|
- w13_layout,
|
|
|
- )
|
|
|
- if collect_activation_amax:
|
|
|
- assert activation_amax is not None
|
|
|
- assert layer_idx_int is not None
|
|
|
- torch.ops.b12x.w4a16_fused_moe_calibrated_launch(
|
|
|
- *launch_common,
|
|
|
- activation_amax,
|
|
|
- int(layer_idx_int),
|
|
|
- *launch_tail,
|
|
|
- int(stream),
|
|
|
- )
|
|
|
- else:
|
|
|
- torch.ops.b12x.w4a16_fused_moe_launch(
|
|
|
- *launch_common,
|
|
|
- *launch_tail,
|
|
|
- bool(use_direct_topk_routes),
|
|
|
- bool(use_tc_decode),
|
|
|
- int(stream),
|
|
|
- )
|
|
|
-
|
|
|
- if use_tc_decode:
|
|
|
- # FC2 already wrote the top-k-summed result into `output`.
|
|
|
- return output
|
|
|
-
|
|
|
- if topk_sum_launch is not None:
|
|
|
- expected_sum = (topk, hidden_size)
|
|
|
- actual_sum = (
|
|
|
- int(topk_sum_launch.topk),
|
|
|
- int(topk_sum_launch.hidden_size),
|
|
|
- )
|
|
|
- if actual_sum != expected_sum:
|
|
|
- raise RuntimeError(
|
|
|
- "preplanned W4A16 top-k sum launch does not match requested contract: "
|
|
|
- f"requested={expected_sum}, planned={actual_sum}"
|
|
|
- )
|
|
|
- torch.ops.b12x.w4a16_topk_sum_launch(
|
|
|
- fc2_out,
|
|
|
- output,
|
|
|
- m,
|
|
|
- topk,
|
|
|
- hidden_size,
|
|
|
- element_dtype,
|
|
|
- int(stream),
|
|
|
- )
|
|
|
- return output
|
|
|
-
|
|
|
-
|
|
|
-__all__ = [
|
|
|
- "W4A16ActivationCompileResult",
|
|
|
- "W4A16FusedMoeCompileResult",
|
|
|
- "W4A16GemmCompileResult",
|
|
|
- "W4A16TopKSumCompileResult",
|
|
|
- "W4A16FusedMoeKernel",
|
|
|
- "W4A16ActivationKernel",
|
|
|
- "W4A16GemmKernel",
|
|
|
- "W4A16TopKSumKernel",
|
|
|
- "clear_w4a16_kernel_cache",
|
|
|
- "compile_w4a16_activation",
|
|
|
- "compile_w4a16_fused_moe",
|
|
|
- "compile_w4a16_gemm",
|
|
|
- "compile_w4a16_topk_sum",
|
|
|
- "pack_topk_routes_by_expert",
|
|
|
- "run_w4a16_moe",
|
|
|
-]
|
