chenchenghao/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add sm90 megamoe phase1 interfaces

Browse Source
This commit is contained in:
Xinyi Liu
2026-06-17 23:54:49 +08:00
parent 062cb160cf
commit 1dafd0f173
12 changed files with 1520 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files
deep_gemm/__init__.py
+1
View File
@@ -85,6 +85,7 @@ from .mega import (
SymmBuffer,
get_symm_buffer_for_mega_moe,
transform_weights_for_mega_moe,
fp8_mega_moe,
fp8_fp4_mega_moe,
)
deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh
+58
View File
@@ -0,0 +1,58 @@
#pragma once
#include <cutlass/cutlass.h>
#include <deep_gemm/common/exception.cuh>
#include <deep_gemm/common/tma_utils.cuh>
#include <deep_gemm/common/types.cuh>
#include <deep_gemm/layout/mega_moe.cuh>
#include <deep_gemm/layout/sym_buffer.cuh>
#include <deep_gemm/scheduler/sm90_mega_moe.cuh>
namespace deep_gemm {
template <
uint32_t kNumMaxTokensPerRank,
uint32_t kHidden, uint32_t kIntermediateHidden,
uint32_t kNumExperts, uint32_t kNumTopk,
uint32_t kNumExpertsPerWave,
uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
uint32_t kStoreBlockM,
uint32_t kNumMaxPoolTokens,
uint32_t kNumPaddedSFPoolTokens,
uint32_t kNumStages,
uint32_t kNumDispatchThreads, uint32_t kNumTMAThreads,
uint32_t kNumMathThreads,
bool kCooperativeMode, bool kUseNMajorL2,
uint32_t kNumSMs, uint32_t kNumRanks,
float kActivationClamp, bool kFastMath,
uint32_t L1_SHAPE_N = kIntermediateHidden * 2,
uint32_t L1_SHAPE_K = kHidden,
uint32_t L2_SHAPE_N = kHidden,
uint32_t L2_SHAPE_K = kIntermediateHidden,
uint32_t kNumThreads = kNumDispatchThreads + kNumTMAThreads + kNumMathThreads,
uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks>
CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void
sm90_fp8_mega_moe_impl(void* y,
int* cumulative_local_expert_recv_stats,
const uint32_t num_tokens,
const __grid_constant__ layout::SymBuffer<kNumRanks> sym_buffer,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights,
const void* l1_weights_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf,
const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights,
const void* l2_weights_sf) {
DG_STATIC_ASSERT(kNumThreads == 384, "SM90 MegaMoE expects 384 threads");
DG_STATIC_ASSERT(BLOCK_N == 128, "SM90 MegaMoE expects BLOCK_N=128");
DG_STATIC_ASSERT(BLOCK_K == 128, "SM90 MegaMoE expects BLOCK_K=128");
DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks");
// Phase 1 only validates the host/JIT/API path and launches an empty kernel.
return;
}
} // namespace deep_gemm
deep_gemm/include/deep_gemm/layout/mega_moe.cuh
+10 -10
View File
@@ -108,46 +108,46 @@ struct Workspace {
static constexpr uint32_t kNumMaxGridSyncCounters = 4;
template <uint32_t kIndex = 0>
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint32_t* get_grid_sync_count_ptr() const {
DG_STATIC_ASSERT(kIndex < kNumMaxGridSyncCounters, "Grid sync index out of bounds");
return static_cast<uint32_t*>(base) + kIndex;
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint32_t* get_nvl_barrier_counter_ptr() const {
return static_cast<uint32_t*>(base) + kNumMaxGridSyncCounters;
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
int* get_nvl_barrier_signal_ptr(const uint32_t& phase) const {
// NOTES: the signal is signed, as we may minus
return math::advance_ptr<int>(base, (kNumMaxGridSyncCounters + 1) * sizeof(uint32_t) + phase * sizeof(int));
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint64_t* get_expert_send_count_ptr(const uint32_t& expert_idx = 0) const {
return math::advance_ptr<uint64_t>(base, kNumBarrierSignalBytes) + expert_idx;
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint64_t* get_expert_recv_count_ptr(
const uint32_t& rank_idx = 0, const uint32_t& expert_idx = 0) const {
return get_expert_send_count_ptr(num_experts) + rank_idx * num_experts_per_rank + expert_idx;
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint64_t* get_expert_recv_count_sum_ptr(const uint32_t& expert_idx = 0) const {
return get_expert_send_count_ptr(num_experts * 2) + expert_idx;
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint32_t* get_l1_arrival_count_ptr(const uint32_t& pool_block_idx = 0) const {
const auto base = get_expert_recv_count_sum_ptr(num_experts_per_rank);
return reinterpret_cast<uint32_t*>(base) + pool_block_idx;
}
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint64_t* get_l2_arrival_mask_ptr(const uint32_t& pool_block_idx = 0) const {
// Pad L1 entry count to even so that the `l2_arrival_mask` is 8-byte aligned
const auto base = get_l1_arrival_count_ptr(math::align(num_max_pool_blocks, 2u));
@@ -155,7 +155,7 @@ struct Workspace {
}
// For dispatch pulling
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
uint32_t* get_src_token_topk_idx_ptr(
const uint32_t& expert_idx = 0, const uint32_t& rank_idx = 0, const uint32_t& token_idx = 0) const {
const auto base = get_l2_arrival_mask_ptr(num_max_pool_blocks);
@@ -165,7 +165,7 @@ struct Workspace {
}
// For combine usages
CUTLASS_DEVICE
CUTLASS_HOST_DEVICE
TokenSrcMetadata* get_token_src_metadata_ptr(const uint32_t& pool_token_idx = 0) const {
const auto base = reinterpret_cast<TokenSrcMetadata*>(get_src_token_topk_idx_ptr(num_experts_per_rank));
return base + pool_token_idx;
deep_gemm/include/deep_gemm/scheduler/sm90_mega_moe.cuh
+199
View File
@@ -0,0 +1,199 @@
#pragma once
#include <deep_gemm/common/cute_tie.cuh>
#include <deep_gemm/common/math.cuh>
#include <deep_gemm/common/types.cuh>
#include <deep_gemm/layout/mega_moe.cuh>
#include <deep_gemm/ptx/ld_st.cuh>
#include <deep_gemm/ptx/utils.cuh>
namespace deep_gemm::sched {
enum class SM90BlockPhase {
None = 0,
Linear1 = 1,
Linear2 = 2
};
template <uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
uint32_t L1_SHAPE_N, uint32_t L1_SHAPE_K,
uint32_t L2_SHAPE_N, uint32_t L2_SHAPE_K,
uint32_t kNumExpertsPerRank,
uint32_t kNumExpertsPerWave,
uint32_t kNumSMs, uint32_t kNumRanks,
bool kUseNMajorL2,
uint32_t kNumExpertsPerLane = math::constexpr_ceil_div(kNumExpertsPerRank, 32u),
uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N,
uint32_t kNumL2BlockNs = L2_SHAPE_N / BLOCK_N,
uint32_t kNumL1BlockKs = L1_SHAPE_K / BLOCK_K,
uint32_t kNumL2BlockKs = L2_SHAPE_K / BLOCK_K>
struct SM90MegaMoEScheduler {
DG_STATIC_ASSERT(L1_SHAPE_N % BLOCK_N == 0, "Invalid shape");
DG_STATIC_ASSERT(L2_SHAPE_N % BLOCK_N == 0, "Invalid shape");
DG_STATIC_ASSERT(L1_SHAPE_K % BLOCK_K == 0, "Invalid shape");
DG_STATIC_ASSERT(L2_SHAPE_K % BLOCK_K == 0, "Invalid shape");
DG_STATIC_ASSERT(kNumExpertsPerRank % kNumExpertsPerWave == 0, "Invalid wave config");
const layout::Workspace& workspace;
SM90BlockPhase next_phase = SM90BlockPhase::Linear1;
uint32_t current_local_expert_idx = 0;
uint32_t current_num_tokens = 0;
uint32_t current_pool_block_offset = 0;
uint32_t block_idx = 0;
uint32_t m_block_idx = 0;
uint32_t n_block_idx = 0;
uint32_t stored_num_tokens_per_expert[kNumExpertsPerLane] = {};
CUTLASS_DEVICE explicit SM90MegaMoEScheduler(const layout::Workspace& workspace): workspace(workspace) {
block_idx = blockIdx.x;
}
CUTLASS_DEVICE uint32_t get_wave_expert_end_idx() const {
return math::align(current_local_expert_idx + 1, kNumExpertsPerWave);
}
CUTLASS_DEVICE uint32_t get_num_tokens(const uint32_t& expert_idx) const {
uint32_t valid_value;
#pragma unroll
for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) {
valid_value = (expert_idx == i * 32 + ptx::get_lane_idx()) ?
stored_num_tokens_per_expert[i] : valid_value;
}
return ptx::exchange(valid_value, expert_idx % 32);
}
CUTLASS_DEVICE uint32_t get_pool_block_offset(const uint32_t& expert_idx) {
uint32_t num_blocks = 0;
#pragma unroll
for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) {
if (i * 32 + ptx::get_lane_idx() < expert_idx)
num_blocks += math::ceil_div(stored_num_tokens_per_expert[i], BLOCK_M);
}
return __reduce_add_sync(0xffffffff, num_blocks);
}
CUTLASS_DEVICE void advance_expert_idx() {
current_pool_block_offset += get_current_num_m_blocks();
current_local_expert_idx += 1;
current_num_tokens = get_num_tokens(current_local_expert_idx);
}
CUTLASS_DEVICE void set_expert_idx(const uint32_t& expert_idx) {
current_local_expert_idx = expert_idx;
current_num_tokens = get_num_tokens(expert_idx);
current_pool_block_offset = get_pool_block_offset(expert_idx);
}
CUTLASS_DEVICE uint32_t get_current_pool_block_offset() const {
return current_pool_block_offset;
}
CUTLASS_DEVICE uint32_t get_current_num_m_blocks() const {
return math::ceil_div(current_num_tokens, BLOCK_M);
}
CUTLASS_DEVICE uint32_t get_valid_m() const {
return cute::min(current_num_tokens - m_block_idx * BLOCK_M, BLOCK_M);
}
CUTLASS_DEVICE bool fetch_next_l1_block() {
const auto wave_end_expert_idx = get_wave_expert_end_idx();
while (current_local_expert_idx < wave_end_expert_idx) {
const auto num_m_blocks = get_current_num_m_blocks();
m_block_idx = block_idx / kNumL1BlockNs;
if (m_block_idx < num_m_blocks)
return true;
block_idx -= num_m_blocks * kNumL1BlockNs;
advance_expert_idx();
}
return false;
}
CUTLASS_DEVICE bool fetch_next_l2_block() {
const auto wave_end_expert_idx = get_wave_expert_end_idx();
while (current_local_expert_idx < wave_end_expert_idx) {
const auto num_m_blocks = get_current_num_m_blocks();
if (block_idx < num_m_blocks * kNumL2BlockNs) {
if constexpr (kUseNMajorL2) {
n_block_idx = block_idx / num_m_blocks;
m_block_idx = block_idx - n_block_idx * num_m_blocks;
} else {
m_block_idx = block_idx / kNumL2BlockNs;
n_block_idx = block_idx - m_block_idx * kNumL2BlockNs;
}
return true;
}
block_idx -= num_m_blocks * kNumL2BlockNs;
advance_expert_idx();
}
return false;
}
CUTLASS_DEVICE cute::tuple<SM90BlockPhase, uint32_t, uint32_t, uint32_t> get_next_block() {
while (true) {
if (current_local_expert_idx >= kNumExpertsPerRank)
break;
if (next_phase == SM90BlockPhase::Linear1) {
if (fetch_next_l1_block()) {
n_block_idx = block_idx - m_block_idx * kNumL1BlockNs;
block_idx += kNumSMs;
return {SM90BlockPhase::Linear1, current_local_expert_idx, m_block_idx, n_block_idx};
} else {
next_phase = SM90BlockPhase::Linear2;
set_expert_idx(math::align<uint32_t, false>(current_local_expert_idx - 1, kNumExpertsPerWave));
}
} else {
if (fetch_next_l2_block()) {
if constexpr (not kUseNMajorL2) {
n_block_idx = block_idx - m_block_idx * kNumL2BlockNs;
}
block_idx += kNumSMs;
return {SM90BlockPhase::Linear2, current_local_expert_idx, m_block_idx, n_block_idx};
} else {
next_phase = SM90BlockPhase::Linear1;
}
}
}
return {SM90BlockPhase::None, 0, 0, 0};
}
CUTLASS_DEVICE void fetch_expert_recv_count() {
#pragma unroll
for (uint32_t i = 0; i < kNumExpertsPerLane; ++ i) {
const auto expert_idx = i * 32 + ptx::get_lane_idx();
uint64_t value = 0;
if (expert_idx < kNumExpertsPerRank) {
do {
value = ptx::ld_volatile(workspace.get_expert_recv_count_sum_ptr(expert_idx));
} while (static_cast<uint32_t>(value >> 32) != kNumSMs * kNumRanks);
}
stored_num_tokens_per_expert[i] = static_cast<uint32_t>(value);
}
__syncwarp();
}
template <typename Func>
CUTLASS_DEVICE void for_each_block(Func&& func) {
fetch_expert_recv_count();
set_expert_idx(0);
while (true) {
CUTE_TIE_DECL(get_next_block(), block_phase, current_local_expert_idx, m_block_idx, n_block_idx);
if (block_phase == SM90BlockPhase::None)
break;
func(block_phase, current_local_expert_idx,
block_phase == SM90BlockPhase::Linear2 ? kNumL2BlockKs : kNumL1BlockKs,
m_block_idx, n_block_idx);
}
}
};
} // namespace deep_gemm::sched
deep_gemm/mega/__init__.py
+116 -31
View File
@@ -13,6 +13,14 @@ except Exception as exception:
from .. import _C
def _is_sm90() -> bool:
return torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9
def _is_sm100() -> bool:
return torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 10
class SymmBuffer:
def __init__(self, group: dist.ProcessGroup,
# MoE arguments
@@ -28,13 +36,23 @@ class SymmBuffer:
self.hidden = hidden
self.intermediate_hidden = intermediate_hidden
# Allocate a symmetric buffer
num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_mega_moe(
group.size(), num_experts,
num_max_tokens_per_rank, num_topk,
hidden, intermediate_hidden,
use_fp8_dispatch, activation
)
# Allocate a symmetric buffer (route by architecture)
if _is_sm90():
num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_sm90_mega_moe(
group.size(), num_experts,
num_max_tokens_per_rank, num_topk,
hidden, intermediate_hidden,
use_fp8_dispatch, activation
)
elif _is_sm100():
num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_mega_moe(
group.size(), num_experts,
num_max_tokens_per_rank, num_topk,
hidden, intermediate_hidden,
use_fp8_dispatch, activation
)
else:
raise RuntimeError('Unsupported architecture for MegaMoE')
self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda')
self.handle = symm_mem.rendezvous(self.buffer, group=group)
self.buffer.zero_()
@@ -46,6 +64,10 @@ class SymmBuffer:
self.topk_idx, self.topk_weights,
self.l1_acts, self.l1_acts_sf,
self.l2_acts, self.l2_acts_sf) = slice_input_buffers(self.buffer)
self.l1_topk_weights = None
self.expert_recv_count_sum = None
self.l1_arrival_count = None
self.token_src_metadata = None
def destroy(self):
self.handle = None
@@ -53,6 +75,16 @@ class SymmBuffer:
self.group = None
self.x = None
self.x_sf = None
self.topk_idx = None
self.topk_weights = None
self.l1_acts = None
self.l1_acts_sf = None
self.l1_topk_weights = None
self.l2_acts = None
self.l2_acts_sf = None
self.expert_recv_count_sum = None
self.l1_arrival_count = None
self.token_src_metadata = None
def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup,
@@ -62,7 +94,13 @@ def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup,
use_fp8_dispatch: bool = True,
activation: str = 'swiglu') -> SymmBuffer:
# Token count must be aligned to block sizes
num_max_tokens_per_rank = align(num_max_tokens_per_rank, _C.get_token_alignment_for_mega_moe())
if _is_sm90():
alignment = _C.get_token_alignment_for_sm90_mega_moe()
elif _is_sm100():
alignment = _C.get_token_alignment_for_mega_moe()
else:
raise RuntimeError('Unsupported architecture for MegaMoE')
num_max_tokens_per_rank = align(num_max_tokens_per_rank, alignment)
return SymmBuffer(
group, num_experts,
@@ -72,16 +110,17 @@ def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup,
)
def _interleave_l1_weights(l1_weights: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
def _interleave_l1_weight_tensor(t: torch.Tensor, gran: int = 8) -> torch.Tensor:
# [gate: 0..7, up: 0..7, gate: 8..15, up: 8..15, ...] instead of [gate | up]
def interleave(t, gran: int = 8) -> torch.Tensor:
g, n, *rest = t.shape
half = n // 2
gate = t[:, :half].reshape(g, half // gran, gran, *rest)
up = t[:, half:].reshape(g, half // gran, gran, *rest)
return torch.empty_like(t).copy_(torch.stack([gate, up], dim=2).reshape(g, n, *rest))
g, n, *rest = t.shape
half = n // 2
gate = t[:, :half].reshape(g, half // gran, gran, *rest)
up = t[:, half:].reshape(g, half // gran, gran, *rest)
return torch.empty_like(t).copy_(torch.stack([gate, up], dim=2).reshape(g, n, *rest))
return interleave(l1_weights[0]), interleave(l1_weights[1])
def _interleave_l1_weights(l1_weights: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
return _interleave_l1_weight_tensor(l1_weights[0]), _interleave_l1_weight_tensor(l1_weights[1])
def _transpose_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor:
@@ -93,36 +132,82 @@ def _transpose_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor:
return torch.empty_like(sf).copy_(result)
def transform_weights_for_mega_moe_sm90(
l1_weights: Tuple[torch.Tensor, torch.Tensor],
l2_weights: Tuple[torch.Tensor, torch.Tensor]
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
# L1: interleave FP8 gate/up weights only; SM90 float weight SF stays natural MN-major.
l1_weights = (_interleave_l1_weight_tensor(l1_weights[0]), l1_weights[1])
# L2: no transform
return l1_weights, l2_weights
def transform_weights_for_mega_moe(
l1_weights: Tuple[torch.Tensor, torch.Tensor],
l2_weights: Tuple[torch.Tensor, torch.Tensor]
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
# L1: interleave gate/up, then transpose SF for UTCCP
if _is_sm90():
return transform_weights_for_mega_moe_sm90(l1_weights, l2_weights)
# SM100: L1 interleave gate/up + UTCCP SF transpose, L2 UTCCP SF transpose
l1_interleaved = _interleave_l1_weights(l1_weights)
l1_weights = (l1_interleaved[0], _transpose_sf_for_utccp(l1_interleaved[1]))
# L2: only transpose SF for UTCCP
l2_weights = (l2_weights[0], _transpose_sf_for_utccp(l2_weights[1]))
return l1_weights, l2_weights
def fp8_mega_moe(y: torch.Tensor,
l1_weights: Tuple[torch.Tensor, torch.Tensor],
l2_weights: Tuple[torch.Tensor, torch.Tensor],
sym_buffer: SymmBuffer,
cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None,
recipe: Optional[Tuple[int, int, int]] = None,
activation: str = 'swiglu',
activation_clamp: Optional[float] = None,
fast_math: bool = True):
if _is_sm90():
if recipe is None:
recipe = (1, 128, 128)
_C.fp8_mega_moe(
y,
l1_weights, l2_weights,
cumulative_local_expert_recv_stats,
sym_buffer.buffer,
sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(),
sym_buffer.num_max_tokens_per_rank,
sym_buffer.num_experts, sym_buffer.num_topk,
recipe,
activation, activation_clamp,
fast_math
)
elif _is_sm100():
if recipe is None:
recipe = (1, 1, 32)
_C.fp8_fp4_mega_moe(
y,
l1_weights, l2_weights,
cumulative_local_expert_recv_stats,
sym_buffer.buffer,
sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(),
sym_buffer.num_max_tokens_per_rank,
sym_buffer.num_experts, sym_buffer.num_topk,
recipe,
activation, activation_clamp,
fast_math
)
else:
raise RuntimeError('Unsupported architecture for MegaMoE')
# Backward-compatible alias
def fp8_fp4_mega_moe(y: torch.Tensor,
l1_weights: Tuple[torch.Tensor, torch.Tensor],
l2_weights: Tuple[torch.Tensor, torch.Tensor],
sym_buffer: SymmBuffer,
cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None,
recipe: Tuple[int, int, int] = (1, 1, 32),
recipe: Optional[Tuple[int, int, int]] = None,
activation: str = 'swiglu',
activation_clamp: Optional[float] = None,
fast_math: bool = True):
_C.fp8_fp4_mega_moe(
y,
l1_weights, l2_weights,
cumulative_local_expert_recv_stats,
sym_buffer.buffer,
sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(),
sym_buffer.num_max_tokens_per_rank,
sym_buffer.num_experts, sym_buffer.num_topk,
recipe,
activation, activation_clamp,
fast_math
)
fp8_mega_moe(y, l1_weights, l2_weights, sym_buffer,
cumulative_local_expert_recv_stats, recipe,
activation, activation_clamp, fast_math)