diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9f113345e..a0a85849c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -15,7 +15,7 @@ repos: minimum_pre_commit_version: 2.9.2 require_serial: true types_or: [python, pyi] - args: ["--line-length", "160"] + args: ["--line-length", "160", "--fast"] - id: black-jupyter name: black-jupyter description: diff --git a/README.md b/README.md index 53765a3fb..444643c36 100644 --- a/README.md +++ b/README.md @@ -23,9 +23,10 @@ We are now shipping **OSS kernels**, allowing you to inspect, modify, and contri * **[GEMM + SwiGLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/gemm_swiglu):** High-performance implementation of the SwiGLU activation fused with GEMM. * **[GEMM + sReLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/gemm_srelu):** High-performance implementation of squared-ReLU fused with GEMM. * **[GEMM + dsReLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/gemm_dsrelu):** High-performance implementation of dsquared-ReLU fused with GEMM. -* **[Grouped GEMM + GLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_glu):** Unified grouped GEMM GLU API supporting dense and discrete MoE weight layouts. +* **[Grouped GEMM (BF16)](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_unfused):** Unfused BF16 grouped GEMM with dense and discrete MoE weight layouts. +* **[Grouped GEMM + GLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_glu):** Unified BF16 and legacy block-scaled grouped GEMM GLU API supporting dense and discrete MoE weight layouts. * **[Grouped GEMM + GLU + Hadamard](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_glu_hadamard):** Dense grouped GEMM GLU forward fusion with a fused Hadamard transform and per-expert AMAX reduction. -* **[Grouped GEMM + dGLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_dglu):** Unified grouped GEMM dGLU backward API supporting dense and discrete MoE weight layouts. +* **[Grouped GEMM + dGLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_dglu):** Unified BF16 and legacy block-scaled grouped GEMM dGLU backward API supporting dense and discrete MoE weight layouts. * **[Grouped GEMM + SwiGLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_swiglu):** SwiGLU activation fused with Grouped GEMM. * **[Grouped GEMM + dSwiglu](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_dswiglu):** dSwiglu activation fused with Grouped GEMM. * **[Grouped GEMM + sReLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_srelu):** Contiguous grouped squared-ReLU GEMM for MoE workloads. @@ -34,7 +35,7 @@ We are now shipping **OSS kernels**, allowing you to inspect, modify, and contri * **[Discrete Grouped GEMM + dSwiGLU](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/discrete_grouped_gemm/discrete_grouped_gemm_dswiglu):** Per-expert-pointer dSwiGLU backward grouped GEMM for MoE workloads without weight packing. * **[Grouped GEMM + Quant](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_quant):** Legacy dense-only grouped GEMM quant API for MoE FC2/dFC1 workloads. * **[Grouped GEMM + Quant (Unified)](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_quant):** Unified grouped GEMM quant API with per-row gating for MoE FC2/dFC1 workloads. -* **[Grouped GEMM + Wgrad](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_wgrad):** Unified grouped GEMM weight-gradient API supporting dense and discrete output layouts for MoE workloads. +* **[Grouped GEMM + Wgrad](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/grouped_gemm/grouped_gemm_wgrad):** Unified BF16 and legacy block-scaled grouped GEMM weight-gradient API supporting dense and discrete output layouts for MoE workloads. * **[BSA](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/block_sparse_attention/):** Block-sparse attention forward and backward CuTe DSL kernels for block-level routing metadata. * **[NSA](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/native_sparse_attention/):** Native Sparse attention as described in the Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention. * **[SDPA Backward: SM100, D=256](https://github.com/NVIDIA/cudnn-frontend/tree/main/python/cudnn/sdpa):** SDPA Backward pass for D=256 on SM100. diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm.md new file mode 100644 index 000000000..ec932bae3 --- /dev/null +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm.md @@ -0,0 +1,153 @@ +# Grouped GEMM (SM100 BF16) + +**This is an experimental API and subject to change.** It requires an NVIDIA +SM100-or-newer GPU and the optional CuTe DSL dependencies: + +```bash +pip install nvidia-cudnn-frontend[cutedsl] +``` + +`GroupedGemmSm100` and `grouped_gemm_wrapper_sm100` implement the neutral, +unfused BF16 MoE grouped GEMM. They support dense stacked expert weights or a +device pointer per expert, optional bias, optional materialization of the +intermediate `C`, and static or dynamic tile scheduling. + +## Operation + +Let `padded_offsets[g]` be the exclusive end of expert `g`'s contiguous row +range and let `begin_g` be zero for the first expert or +`padded_offsets[g - 1]` otherwise. For rows in `[begin_g, padded_offsets[g])`: + +```text +G_g = alpha[g] * A_g @ B_g.T +C_g = G_g + prob_g * bias[:, g] # when bias is present +D_g = C_g +``` + +Without bias, `C_g = G_g` and `D_g = prob_g * G_g`. Accumulation is FP32; +`C` and `D` may independently use BF16, FP16, or FP32. + +## Tensors and layouts + +For total padded rows `M`, inner dimension `K`, output dimension `N`, and `L` +experts: + +| Tensor | Shape | Required stride / dtype | +| --- | --- | --- | +| `A` | `(M, K, 1)` | `(K, 1, M*K)`, BF16 | +| dense `B` | `(N, K, L)` | `(K, 1, N*K)`, BF16 | +| discrete `b_ptrs` | `(L,)` | contiguous CUDA int64 pointers to `(N, K)` BF16 matrices | +| `padded_offsets` | `(L,)` | `(1,)`, CUDA int32 cumulative ends | +| `alpha` | `(L,)` | `(1,)`, CUDA FP32 | +| `prob` | `(M, 1, 1)` | `(1, 1, 1)`, CUDA FP32 | +| optional `bias` | `(N, L)` | `(1, N)`, BF16/FP16/FP32 | +| `C`, `D` | `(M, N, 1)` | `(N, 1, M*N)`, BF16/FP16/FP32 | + +`M` and every cumulative offset must be 256-aligned. Dense weights are +K-major. Discrete mode uses `b_major="k"`, `n=N`, and +`b_dtype=torch.bfloat16`. + +The pointer-array tensor must be contiguous, non-null, eight-byte aligned, and +on the same device as `A`; each target pointer must satisfy the kernel's +alignment contract. The API records the pointer-array tensor on the launch +stream. The caller must keep every pointed-to expert allocation alive and must +not modify or free it until that stream completes. + +## Wrapper API + +Dense mode: + +```python +import cudnn +import torch + +result = cudnn.grouped_gemm_wrapper_sm100( + a_tensor=a, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + b_tensor=b, + bias_tensor=bias, + prob_tensor=prob, + c_dtype=torch.float32, + d_dtype=torch.bfloat16, + generate_c=True, + use_dynamic_sched=True, +) +d, c = result +assert d is result["d_tensor"] +assert c is result["c_tensor"] +``` + +Discrete mode changes only the weight arguments: + +```python +result = cudnn.grouped_gemm_wrapper_sm100( + a_tensor=a, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + b_ptrs=b_ptrs, + n=N, + b_dtype=torch.bfloat16, + b_major="k", + bias_tensor=bias, + prob_tensor=prob, +) +``` + +The `TupleDict` order is exactly `d_tensor`, then `c_tensor`. `c_tensor` is +`None` unless `generate_c=True`; the kernel still uses an internal C buffer +when it is needed for execution. + +## Class API + +The class API takes representative sample tensors, then follows +`check_support()` -> `compile()` -> `execute()`: + +```python +op = cudnn.GroupedGemmSm100( + sample_a=a, + sample_c=c, + sample_d=d, + sample_padded_offsets=padded_offsets, + sample_alpha=alpha, + sample_b=b, + sample_bias=bias, + sample_prob=prob, + acc_dtype=torch.float32, + generate_c=True, + use_dynamic_sched=False, +) +assert op.check_support() +op.compile() +op.execute( + a_tensor=a, + c_tensor=c, + d_tensor=d, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + b_tensor=b, + bias_tensor=bias, + prob_tensor=prob, +) +``` + +For a discrete class instance, replace `sample_b` with +`num_experts=L`, `b_shape=(N, K)`, `b_dtype=torch.bfloat16`, and pass +`b_ptrs` to `execute()`. + +## Scheduling, caching, and errors + +- `use_dynamic_sched=False` uses the static scheduler. +- `use_dynamic_sched=True` compiles a dynamic-M callable and reuses it for + compatible M values; discrete mode also allocates the per-expert tensor-map + workspace. Wrapper cache keys retain dtype, layout, expert count, optional + features, scheduler choice, output policy, tile/cluster shape, and overlap + margin. +- Dense and discrete weight arguments are mutually exclusive. Invalid shapes, + strides, dtypes, devices, alignment, offsets, pointer entries, output + descriptors, tiles/clusters, or a target below SM100 raise `ValueError` or + `RuntimeError` before launch. +- Fused GLU, dGLU, and WGrad APIs select BF16 from BF16 operands while keeping + their existing FP4/FP8 block-scaled backends. For BF16 on those fused APIs, + scale-factor controls are `None`; see their operation pages for the exact + dispatch and return contracts. diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md index c2a0f9bdd..ba3cf90a0 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_dglu.md @@ -4,7 +4,10 @@ ## Overview -**Unified Grouped GEMM + dGLU fusion**: A block-scaled grouped GEMM fused with a dGLU backward epilogue (dSwiGLU or dGeGLU) on NVIDIA Blackwell GPUs (SM100+), designed for MoE (Mixture of Experts) workloads. Implemented with CUTLASS/CUTE. +**Unified Grouped GEMM + dGLU fusion**: one public class and wrapper select a +plain BF16 or legacy block-scaled grouped GEMM fused with a dGLU backward +epilogue (dSwiGLU or dGeGLU) on NVIDIA Blackwell GPUs (SM100+). The operation +is implemented with CUTLASS/CuTe DSL. This is a **unified API** that supports both weight layout modes: - **Dense mode**: All expert weights packed into a single contiguous `(N, K, L)` tensor @@ -16,7 +19,98 @@ And both backward activation functions: Groups are contiguous in the M dimension and described by `padded_offsets` (cumulative aligned end offsets). -This kernel performs: +### Backend dispatch + +| Operand contract | Selected backend | +| --- | --- | +| `A` and `B` are BF16 | BF16 | +| matching supported FP4/FP8 `A` and `B` plus scale descriptors | block-scaled | + +Mixed families and unsupported pairs are rejected before allocation or +compilation. Each backend's argument contract is described below. + +## BF16 contract + +Pass `sfa_tensor=None`, `sfb_tensor=None` (or `sfb_ptrs=None`), and +`norm_const_tensor=None`; keep `sf_vec_size=16`, `discrete_col_sfd=False`, and +`epilogue_op=None`. BF16 also uses the source GeGLU constants +`geglu_alpha=1.702`, `glu_clamp_max=7`, and `glu_clamp_min=-7`. + +### Tensors, layouts, and equation + +For padded rows `M`, reduction dimension `K`, compact gradient width `N`, and +`L` experts, BF16 uses: + +- `A`: `(M, K, 1)`, stride `(K, 1, M*K)`, BF16; +- dense `B`: `(N, K, L)`, K-major stride `(K, 1, N*K)`, BF16; +- discrete `b_ptrs`: contiguous CUDA int64 pointers to expert `(N, K)` BF16 + matrices, with `n=N`, `b_dtype=torch.bfloat16`, and `b_major="k"` or `"n"`; +- `C`: `(M, 2N, 1)`, stride `(2N, 1, 2M*N)`, BF16/FP16/FP32; +- `padded_offsets`: `(L,)` int32 cumulative 256-aligned ends; +- `alpha`, `beta`: `(L,)` FP32; `prob`: `(M, 1, 1)` FP32; +- caller-zeroed `dprob`: `(M, 1, 1)`, stride `(1, 1, 1)`, FP32; +- `D_row`: `(M, 2N, 1)`, stride `(2N, 1, 2M*N)`, BF16/FP16/FP32; +- caller-zeroed optional `dbias`: `(L, 2N, 1)`, stride `(2N, 1, 1)`, BF16. + +For expert `g`, the compact GEMM gradient and scaled forward activation are + +$$ +R_g = \alpha_g^2 A_g B_g^T, \qquad X_g = \beta_g C_g. +$$ + +Split `X` into alternating 32-wide gate/input blocks. For dSwiGLU, with +`s = sigmoid(gate)`: + +$$ +d\mathrm{input} = R\,\mathrm{prob}\,(\mathrm{gate}\,s), +$$ + +$$ +d\mathrm{gate} = R\,\mathrm{prob}\,\mathrm{input}\,s + (1 + \mathrm{gate}(1-s)). +$$ + +For dGeGLU, distinguish raw values, clamped activation values, and the source's +value-bearing filters: + +```text +raw_gate = gate(X) +raw_input = input(X) +clamped_gate = min(raw_gate, 7) +clamped_input = clamp(raw_input, -7, 7) +gate_filter = raw_gate if raw_gate <= 7 else 0 +input_filter = raw_input if -7 <= raw_input <= 7 else 0 +s = sigmoid(1.702 * clamped_gate) +``` + +With the default `linear_offset=1`, the kernel computes + +$$ +d\mathrm{gate} = R\,\mathrm{prob}\, + (\mathrm{clamped\_input}+\mathrm{linear\_offset})\,s + (1 + 1.702\,\mathrm{clamped\_gate}(1-s))\, + \mathrm{gate\_filter}, +$$ + +$$ +d\mathrm{input} = R\,\mathrm{prob}\,\mathrm{clamped\_gate}\,s\, + \mathrm{input\_filter}. +$$ + +`dprob` accumulates the row sum of the matching unscaled activation times `R`; +`dbias` is the per-expert row reduction of interleaved `D_row`. The +pointer-array tensor is stream-recorded, while every pointed allocation must +remain alive and unchanged until that stream completes. + +The wrapper return order is exactly `d_row_tensor`, `d_col_tensor`, +`dprob_tensor`, `dbias_tensor`, `amax_tensor`, `sfd_row_tensor`, +`sfd_col_tensor`. On BF16, `d_col_tensor`, `amax_tensor`, `sfd_row_tensor`, and +`sfd_col_tensor` are `None`; `dbias_tensor` is `None` unless +`generate_dbias=True`. + +## Block-scaled contract + +The block-scaled backend performs: 1. **Block-scaled grouped GEMM**: Low-precision GEMM (FP4, FP8) with per-block scale factors across multiple expert groups 2. **dGLU backward epilogue**: Fused backward computation using the forward `C` tensor (input/gate interleaved) 3. **Optional quantized output**: Produces row and column scale factors for downstream quantization @@ -110,9 +204,74 @@ $$ --- -## API Usage +## API usage + +### BF16 + +#### High-level wrapper + +```python +import cudnn +import torch + +dprob.zero_() +out = cudnn.grouped_gemm_dglu_wrapper_sm100( + a_tensor=a, + c_tensor=c, + sfa_tensor=None, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + beta_tensor=beta, + prob_tensor=prob, + dprob_tensor=dprob, + b_tensor=b, + sfb_tensor=None, + act_func="dswiglu", + generate_dbias=True, + use_dynamic_sched=True, +) +d_row, d_col, dprob, dbias, amax, sfd_row, sfd_col = out + +# Discrete mode replaces the dense weight arguments. +out = cudnn.grouped_gemm_dglu_wrapper_sm100( + a_tensor=a, c_tensor=c, sfa_tensor=None, + padded_offsets=padded_offsets, alpha_tensor=alpha, beta_tensor=beta, + prob_tensor=prob, dprob_tensor=dprob, + b_ptrs=b_ptrs, sfb_ptrs=None, n=N, b_dtype=torch.bfloat16, + act_func="dgeglu", +) +``` + +#### Class API + +```python +op = cudnn.GroupedGemmDgluSm100( + sample_a=a, sample_c=c, sample_d_row=d_row, sample_d_col=None, + sample_sfa=None, sample_padded_offsets=padded_offsets, + sample_alpha=alpha, sample_beta=beta, sample_prob=prob, + sample_dprob=dprob, sample_dbias=dbias, + sample_b=b, sample_sfb=None, act_func="dswiglu", +) +assert op.check_support() +op.compile() +dprob.zero_() +dbias.zero_() +op.execute( + a_tensor=a, c_tensor=c, d_row_tensor=d_row, d_col_tensor=None, + sfa_tensor=None, padded_offsets=padded_offsets, alpha_tensor=alpha, + beta_tensor=beta, prob_tensor=prob, dprob_tensor=dprob, + dbias_tensor=dbias, b_tensor=b, sfb_tensor=None, +) +``` + +`use_dynamic_sched=False` selects static scheduling; `True` caches a dynamic-M +callable for compatible shapes. Cache keys retain compile-sensitive layouts, +dtypes, activation, dbias policy, scheduler, tiles/clusters, features, and +overlap margin. + +### Block-scaled -### High-level Wrapper +#### High-level wrapper **Dense mode:** @@ -185,7 +344,7 @@ outputs = grouped_gemm_dglu_wrapper_sm100( ) ``` -### Class API +#### Class API **Dense mode:** diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md index 320912606..0d3ebe85e 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_glu.md @@ -4,7 +4,10 @@ ## Overview -**Unified Grouped GEMM + GLU fusion**: A block-scaled grouped GEMM fused with a GLU epilogue (SwiGLU or GeGLU) on NVIDIA Blackwell GPUs (SM100+), designed for MoE (Mixture of Experts) workloads. Implemented with CUTLASS/CUTE. +**Unified Grouped GEMM + GLU fusion**: one public class and wrapper select a +plain BF16 or legacy block-scaled grouped GEMM fused with a GLU epilogue +(SwiGLU or GeGLU) on NVIDIA Blackwell GPUs (SM100+). The operation is +implemented with CUTLASS/CuTe DSL. This is a **unified API** that supports both weight layout modes: - **Dense mode**: All expert weights packed into a single contiguous `(N, K, L)` tensor @@ -16,7 +19,71 @@ And both activation functions: Groups are contiguous in the M dimension and described by `padded_offsets` (cumulative aligned end offsets). -This kernel performs: +### Backend dispatch + +| Operand contract | Selected backend | +| --- | --- | +| `A` and `B` are BF16 | BF16 | +| matching supported FP4/FP8 `A` and `B` plus scale descriptors | block-scaled | + +Mixed families and unsupported pairs are rejected before allocation or +compilation. Each backend's argument contract is described below. + +## BF16 contract + +Pass `sfa_tensor=None`, `sfb_tensor=None` (or `sfb_ptrs=None`), and +`norm_const_tensor=None`; keep `sf_vec_size=16` and `discrete_col_sfd=False`. +Non-`None` scale controls are an error. + +### Tensors, layouts, and equation + +For padded rows `M`, reduction dimension `K`, pre-GLU width `N`, and `L` +experts, BF16 uses: + +- `A`: `(M, K, 1)`, stride `(K, 1, M*K)`, BF16; +- dense `B`: `(N, K, L)`, K-major stride `(K, 1, N*K)`, BF16; +- discrete `b_ptrs`: contiguous CUDA int64 pointers to expert `(N, K)` BF16 + matrices, with `n=N`, `b_dtype=torch.bfloat16`, and `b_major="k"` or `"n"`; +- `padded_offsets`: `(L,)`, stride `(1,)`, int32 cumulative 256-aligned ends; +- `alpha`: `(L,)`, FP32; `prob`: `(M, 1, 1)`, stride `(1, 1, 1)`, FP32; +- optional bias: `(N, L)`, stride `(1, N)`, BF16/FP16/FP32; +- `C`: `(M, N, 1)`, stride `(N, 1, M*N)`; +- `D`: `(M, N/2, 1)`, stride `(N/2, 1, M*N/2)`. + +For expert `g`, first compute + +$$ +C_g = \alpha_g A_g B_g^T + \mathrm{bias}_g. +$$ + +Columns are paired as alternating 32-wide gate/up blocks. For SwiGLU, + +$$ +D_g = \mathrm{prob}_g \cdot \mathrm{up}(C_g) \cdot + \mathrm{silu}(\mathrm{gate}(C_g)). +$$ + +For GeGLU, let `gate = min(gate(C), 7)`, +`up = clamp(up(C), -7, 7)`, `geglu_alpha=1.702`, and the default +`linear_offset=1`: + +$$ +D_g = \mathrm{prob}_g \cdot (\mathrm{up}+\mathrm{linear\_offset}) + \cdot \mathrm{gate} \cdot \sigma(1.702\,\mathrm{gate}). +$$ + +`C`/`D` may be BF16, FP16, or FP32. `N` is divisible by 64. The pointer-array +tensor is stream-recorded; every pointed allocation must remain alive and +unchanged until the launch stream completes. + +The wrapper return order is exactly `c_tensor`, `d_tensor`, `d_col_tensor`, +`amax_tensor`, `sfd_row_tensor`, `sfd_col_tensor`. On BF16, +`d_col_tensor`, `amax_tensor`, `sfd_row_tensor`, and `sfd_col_tensor` are +always `None`; `c_tensor` is `None` unless `generate_c=True`. + +## Block-scaled contract + +The block-scaled backend performs: 1. **Block-scaled grouped GEMM**: Low-precision GEMM (FP4, FP8) with per-block scale factors across multiple expert groups 2. **GLU activation**: Fused SwiGLU or GeGLU activation applied to the GEMM output 3. **Optional quantized output**: Produces row and column scale factors for downstream quantization @@ -113,9 +180,81 @@ $$ --- -## API Usage +## API usage + +### BF16 + +#### High-level wrapper + +```python +import cudnn +import torch + +# Dense wrapper. The required scale positions are explicitly None for BF16. +out = cudnn.grouped_gemm_glu_wrapper_sm100( + a_tensor=a, + sfa_tensor=None, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + b_tensor=b, + sfb_tensor=None, + bias_tensor=bias, + prob_tensor=prob, + act_func="swiglu", + generate_c=True, + use_dynamic_sched=True, +) +c, d, d_col, amax, sfd_row, sfd_col = out + +# Discrete wrapper. +out = cudnn.grouped_gemm_glu_wrapper_sm100( + a_tensor=a, + sfa_tensor=None, + padded_offsets=padded_offsets, + alpha_tensor=alpha, + b_ptrs=b_ptrs, + sfb_ptrs=None, + n=N, + b_dtype=torch.bfloat16, + prob_tensor=prob, + act_func="geglu", +) +``` + +#### Class API + +```python +op = cudnn.GroupedGemmGluSm100( + sample_a=a, + sample_c=c, + sample_d=d, + sample_d_col=None, + sample_sfa=None, + sample_padded_offsets=padded_offsets, + sample_alpha=alpha, + sample_b=b, + sample_sfb=None, + sample_prob=prob, + act_func="swiglu", + generate_c=True, +) +assert op.check_support() +op.compile() +op.execute( + a_tensor=a, c_tensor=c, d_tensor=d, sfa_tensor=None, + padded_offsets=padded_offsets, alpha_tensor=alpha, + b_tensor=b, sfb_tensor=None, prob_tensor=prob, +) +``` + +`use_dynamic_sched=False` uses static scheduling; `True` caches a dynamic-M +callable for compatible shapes. Cache keys include compile-sensitive layouts, +dtypes, features, activation, scheduler, tile/cluster, output policy, and +overlap margin, but not the runtime GeGLU `linear_offset`. + +### Block-scaled -### High-level Wrapper +#### High-level wrapper **Dense mode:** @@ -187,7 +326,7 @@ outputs = grouped_gemm_glu_wrapper_sm100( `bias_tensor` must use the kernel layout expected by the fused bias path: shape `(N, L)` and stride `(1, N)`. -### Class API +#### Class API **Dense mode:** diff --git a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md index a71de62ef..59acdf82d 100644 --- a/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md +++ b/docs/fe-oss-apis/gemm_fusions/grouped_gemm_wgrad.md @@ -1,47 +1,98 @@ -# Grouped GEMM + Wgrad +# Grouped GEMM + WGrad (Unified) -`GroupedGemmWgradSm100` and `grouped_gemm_wgrad_wrapper_sm100` expose the grouped GEMM weight-gradient kernel integrated from the Cute DSL kernel library. +`GroupedGemmWgradSm100` and `grouped_gemm_wgrad_wrapper_sm100` are experimental +SM100+ APIs for grouped MoE weight gradients. The same public surface dispatches +BF16 inputs to the BF16 kernel and preserves the legacy FP4/FP8 block-scaled +backend. + +Install the optional CuTe DSL dependencies before importing either API: + +```bash +pip install nvidia-cudnn-frontend[cutedsl] +``` ## Operation -The API computes grouped weight gradients in 2Dx2D form: +For expert `e`, let `begin = 0` for the first expert and +`begin = offsets_tensor[e - 1]` otherwise, and let +`end = offsets_tensor[e]`. The API computes ```text -A(hidden, tokens_sum) x B(tokens_sum, intermediate) -> Wgrad(experts, hidden, intermediate) +Wgrad[e] = A[:, begin:end] @ B[begin:end, :] ``` -The grouped dimension is the token axis, segmented by `offsets_tensor`. +When `accumulate_on_output=True`, that result is accumulated into the existing +output. The caller must therefore initialize every output allocation. When it +is false, the kernel overwrites the output; an empty expert produces zero. + +## BF16 contract + +The BF16 backend accepts: + +| Argument | Shape | Supported stride/major | Dtype | +| --- | --- | --- | --- | +| `a_tensor` | `(hidden, tokens_sum)` | `(tokens_sum, 1)` K-major or `(1, hidden)` M-major | `torch.bfloat16` | +| `b_tensor` | `(tokens_sum, intermediate)` | `(1, tokens_sum)` K-major or `(intermediate, 1)` N-major | `torch.bfloat16` | +| `offsets_tensor` | `(num_experts,)` | contiguous `(1,)` | `torch.int32` | +| dense `wgrad_tensor` | `(num_experts, hidden, intermediate)` | `(hidden * intermediate, intermediate, 1)` | BF16, FP16, or FP32 | +| one discrete output | `(hidden, intermediate)` | `(intermediate, 1)` | BF16, FP16, or FP32 | +| `wgrad_ptrs` | `(num_experts,)` | contiguous `(1,)` | `torch.int64` | + +`offsets_tensor` is a non-decreasing cumulative sum. Every expert token count +(`offsets[e] - offsets[e - 1]`) must be a multiple of 256, and the final offset +must equal `tokens_sum`. Inputs, metadata, and outputs must reside on the same +CUDA device and satisfy the API's alignment checks. + +BF16 uses FP32 accumulation and requires `sf_vec_size=16`. Pass `None` for +`sfa_tensor`, `sfb_tensor`, `global_scale_a`, and `global_scale_b`. BF16 rejects +every non-`None` scale or global-scale control with `ValueError`; it never falls +through to another backend. Only a supported FP4/FP8 operand pair selects the +legacy block-scaled backend, which continues to support its existing scale +tensors and global scales. + +`input_order` describes how the token dimension is stored: -## Public APIs +- `"tensor2d"` (default) uses one global 2-D tensor and its declared strides. +- `"tensor_ragged"` uses per-expert K-contiguous blocks concatenated in memory. + In this mode only each input's unit-stride axis is meaningful; non-unit host + strides are ignored when per-expert TMA descriptors are built. -- Class API: `cudnn.GroupedGemmWgradSm100` -- Wrapper API: `cudnn.grouped_gemm_wgrad_wrapper_sm100` +### Output modes -## Inputs +With `output_mode="dense"`, provide or let the wrapper allocate the contiguous +stacked `wgrad_tensor`. `wgrad_ptrs` is forbidden. -- `a_tensor`: input tensor with logical shape `(hidden, tokens_sum)` -- `b_tensor`: input tensor with logical shape `(tokens_sum, intermediate)` -- `sfa_tensor`: assembled scale-factor tensor for `a_tensor` -- `sfb_tensor`: assembled scale-factor tensor for `b_tensor` -- `offsets_tensor`: cumulative end offsets per expert, shape `(num_experts,)`, dtype `torch.int32` -- `global_scale_a` / `global_scale_b`: optional per-expert global scales. These are required for NVFP4 (`sf_vec_size == 16` with FP4 inputs). +With `output_mode="discrete"`, either: -`input_order` selects how `a_tensor` and `b_tensor` are interpreted by the kernel: +- omit both output arguments and let the wrapper allocate a stacked tensor and + construct an internal pointer array; or +- provide a CUDA `torch.int64` `wgrad_ptrs` array containing one non-null, + 16-byte-aligned output address per expert. -- `"tensor2d"` (default): the token axis is one contiguous 2-D tensor. -- `"tensor_ragged"`: each expert's token block is laid out independently in memory and concatenated expert-major; the kernel builds per-expert TMA descriptors for A and B. +For explicit pointer-only output, `result["wgrad_tensor"]` is `None`. The caller +owns all pointed-to output allocations and must keep both those allocations and +the pointer tensor alive until work on `current_stream` completes. The API +records the pointer tensor on the launch stream, but it cannot manage the +lifetime of allocations represented only by integer addresses. -## Output Modes +The wrapper always returns `TupleDict(wgrad_tensor=...)`; it contains exactly +one item and supports either keyed access or tuple unpacking. -The API supports two output modes through one public surface: +## Block-scaled contract -- Dense: - - output tensor is a contiguous stacked tensor with shape `(num_experts, hidden, intermediate)` -- Discrete: - - the kernel uses per-expert output pointers internally - - the convenience wrapper still returns a stacked tensor while exercising the discrete-output path +The legacy block-scaled backend is selected only by a supported matching FP4/FP8 +operand pair. It preserves the pre-existing scale-factor contract: provide +`sfa_tensor` and `sfb_tensor`, and provide `global_scale_a` and +`global_scale_b` where the selected low-precision format requires them. BF16 +does not reinterpret these controls; it rejects them instead. -## Wrapper Example +## API usage + +### BF16 + +#### Wrapper + +Dense BF16 output: ```python import cudnn @@ -50,23 +101,99 @@ import torch result = cudnn.grouped_gemm_wgrad_wrapper_sm100( a_tensor=a_tensor, b_tensor=b_tensor, - sfa_tensor=sfa_tensor, - sfb_tensor=sfb_tensor, + sfa_tensor=None, + sfb_tensor=None, offsets_tensor=offsets_tensor, output_mode="dense", wgrad_dtype=torch.bfloat16, + input_order="tensor2d", +) +wgrad_tensor = result["wgrad_tensor"] +``` + +Discrete BF16 outputs owned by the caller: + +```python +expert_outputs = [ + torch.empty( + (hidden, intermediate), dtype=torch.bfloat16, device="cuda" + ) + for _ in range(offsets_tensor.numel()) +] +wgrad_ptrs = torch.tensor( + [output.data_ptr() for output in expert_outputs], + dtype=torch.int64, + device="cuda", +) + +result = cudnn.grouped_gemm_wgrad_wrapper_sm100( + a_tensor=a_tensor, + b_tensor=b_tensor, + sfa_tensor=None, + sfb_tensor=None, + offsets_tensor=offsets_tensor, + output_mode="discrete", + wgrad_ptrs=wgrad_ptrs, + wgrad_dtype=torch.bfloat16, input_order="tensor_ragged", ) +assert result["wgrad_tensor"] is None +``` -wgrad_tensor = result["wgrad_tensor"] +#### Reusable class lifecycle + +The class API requires output descriptors at construction and output storage at +execution. This dense BF16 example compiles once and accepts later calls with a +different `tokens_sum` when static dimensions, dtypes, majors, and configuration +remain compatible: + +```python +op = cudnn.GroupedGemmWgradSm100( + sample_a=a_tensor, + sample_b=b_tensor, + sample_sfa=None, + sample_sfb=None, + sample_offsets=offsets_tensor, + sample_wgrad=wgrad_tensor, + acc_dtype=torch.float32, + input_order="tensor2d", +) +op.check_support() +op.compile() +op.execute( + a_tensor=a_tensor, + b_tensor=b_tensor, + sfa_tensor=None, + sfb_tensor=None, + offsets_tensor=offsets_tensor, + wgrad_tensor=wgrad_tensor, +) ``` -## Class API Example +For a discrete class instance, replace `sample_wgrad` with +`sample_wgrad_expert=expert_outputs[0]`, `num_experts`, `wgrad_shape`, and +`wgrad_dtype`, then pass `wgrad_ptrs` to `execute`. + +### Block-scaled + +#### Wrapper ```python -import cudnn -import torch +result = cudnn.grouped_gemm_wgrad_wrapper_sm100( + a_tensor=a_tensor, + b_tensor=b_tensor, + sfa_tensor=sfa_tensor, + sfb_tensor=sfb_tensor, + offsets_tensor=offsets_tensor, + output_mode="dense", + wgrad_dtype=torch.bfloat16, + input_order="tensor_ragged", +) +``` +#### Reusable class lifecycle + +```python op = cudnn.GroupedGemmWgradSm100( sample_a=a_tensor, sample_b=b_tensor, @@ -76,7 +203,7 @@ op = cudnn.GroupedGemmWgradSm100( sample_wgrad=sample_wgrad_tensor, acc_dtype=torch.float32, ) -op.check_support() +assert op.check_support() op.compile() op.execute( a_tensor=a_tensor, @@ -88,8 +215,16 @@ op.execute( ) ``` -## Notes +## Scheduling, cache, and errors + +The BF16 kernel uses dynamic persistent scheduling. The token dimension is +compiled dynamically, and the wrapper cache abstracts the token-sized axes of A +and B while retaining static dimensions, layouts, dtypes, output descriptors, +tiling, cluster shape, input order, and accumulation mode in its key. A changed +static contract creates a different cached operator or fails validation. -- Requires SM100+ GPUs. -- `output_mode="discrete"` is available in the wrapper for parity with the underlying kernel mode. -- `accumulate_on_output=True` expects the output tensor to be initialized by the caller; the wrapper zero-initializes it automatically. +The APIs reject unsupported dtypes or layouts, malformed/unaligned offsets or +pointers, mixed devices, forbidden BF16 scale controls, unsupported tiling, use +before `compile()`, unavailable CUDA, and devices below SM100. Support and +validation errors are reported as `ValueError` or `RuntimeError`; callers should +not rely on this experimental API remaining source-compatible across releases. diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 2bc3ccfb2..ed3b8bfcf 100644 --- a/docs/fe-oss-apis/overview.md +++ b/docs/fe-oss-apis/overview.md @@ -8,6 +8,7 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d - [GEMM + SwiGLU](gemm_fusions/gemm_swiglu.md) - [GEMM + sReLU](gemm_fusions/gemm_srelu.md) - [GEMM + dsReLU](gemm_fusions/gemm_dsrelu.md) +- [Grouped GEMM (BF16)](gemm_fusions/grouped_gemm.md) - [Grouped GEMM + GLU (Unified)](gemm_fusions/grouped_gemm_glu.md) - [Grouped GEMM + GLU + Hadamard](gemm_fusions/grouped_gemm_glu_hadamard.md) - [Grouped GEMM + dGLU (Unified)](gemm_fusions/grouped_gemm_dglu.md) @@ -29,7 +30,7 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d ## Installation and setup -All Frontend OSS APIs come installed with the `nvidia-cudnn-frontend` package. However, each API may require additional optional dependencies defined in the `pyproject.toml` file. For instance, GEMM + Amax and GEMM + SwiGLU require the `cute-dsl` optional dependency, which can be installed via: +All Frontend OSS APIs come installed with the `nvidia-cudnn-frontend` package. However, each API may require additional optional dependencies defined in the `pyproject.toml` file. For instance, GEMM + Amax, GEMM + SwiGLU, and the grouped GEMM APIs require the `cutedsl` optional dependency, which can be installed via: ```bash pip install nvidia-cudnn-frontend[cutedsl] ``` diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 6e8b99829..1a4d380dc 100644 --- a/python/cudnn/__init__.py +++ b/python/cudnn/__init__.py @@ -271,10 +271,35 @@ def _dlopen_cudnn(): from .nodes import Node from .graph import graph, jit, graph_cache -from .wrapper import Graph from typing import Any +_EAGER_PUBLIC_NAMES = ( + *symbols_to_import, + *( + symbol + for symbol in ( + "causal_conv1d_forward", + "causal_conv1d_backward", + "causal_conv1d_nwh_forward", + "causal_conv1d_nwh_backward", + "b2b_causal_conv1d_forward", + "b2b_causal_conv1d_backward", + ) + if symbol in globals() + ), + "__version__", + "NodeType", + "Tensor", + "pygraph", + "GraphContext", + "Node", + "graph", + "jit", + "graph_cache", +) +__all__ = [*_EAGER_PUBLIC_NAMES, "Graph", "wrapper"] + _OPTIONAL_DEPENDENCY_INSTALL_HINT = "Install with 'pip install nvidia-cudnn-frontend[cutedsl]'" _LAZY_OPTIONAL_IMPORTS = { @@ -296,6 +321,8 @@ def _dlopen_cudnn(): "RmsNormRhtAmaxSm100": (".rmsnorm_rht_amax", "RmsNormRhtAmaxSm100"), "rmsnorm_rht_amax_wrapper_sm100": (".rmsnorm_rht_amax", "rmsnorm_rht_amax_wrapper_sm100"), "grouped_gemm": (".grouped_gemm", None), + "GroupedGemmSm100": (".grouped_gemm", "GroupedGemmSm100"), + "grouped_gemm_wrapper_sm100": (".grouped_gemm", "grouped_gemm_wrapper_sm100"), "GroupedGemmSwigluSm100": (".grouped_gemm", "GroupedGemmSwigluSm100"), "grouped_gemm_swiglu_wrapper_sm100": (".grouped_gemm", "grouped_gemm_swiglu_wrapper_sm100"), "GroupedGemmDswigluSm100": (".grouped_gemm", "GroupedGemmDswigluSm100"), @@ -339,6 +366,12 @@ def _load_optional_symbol(name: str) -> Any: def __getattr__(name: str) -> Any: + if name in ("Graph", "wrapper"): + _wrapper = importlib.import_module(".wrapper", __name__) + globals()["wrapper"] = _wrapper + globals()["Graph"] = _wrapper.Graph + return globals()[name] + if name == "ops": # Use importlib rather than "from . import ops" to avoid infinite # recursion. The cycle: @@ -361,3 +394,7 @@ def __getattr__(name: str) -> Any: return _load_optional_symbol(name) raise AttributeError(name) + + +def __dir__(): + return sorted(set(globals()) | set(__all__)) diff --git a/python/cudnn/grouped_gemm/__init__.py b/python/cudnn/grouped_gemm/__init__.py index 818a1ab99..592047ce6 100644 --- a/python/cudnn/grouped_gemm/__init__.py +++ b/python/cudnn/grouped_gemm/__init__.py @@ -46,6 +46,11 @@ grouped_gemm_wgrad_wrapper_sm100, ) +from .grouped_gemm_unfused.api import ( + GroupedGemmSm100, + grouped_gemm_wrapper_sm100, +) + __all__ = [ "GroupedGemmSwigluSm100", "grouped_gemm_swiglu_wrapper_sm100", @@ -65,4 +70,6 @@ "grouped_gemm_dglu_wrapper_sm100", "GroupedGemmWgradSm100", "grouped_gemm_wgrad_wrapper_sm100", + "GroupedGemmSm100", + "grouped_gemm_wrapper_sm100", ] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dglu/_bf16_api.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/_bf16_api.py new file mode 100644 index 000000000..5848039d9 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/_bf16_api.py @@ -0,0 +1,572 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Private descriptor-first BF16 API for SM100 grouped GEMM dGLU.""" + +import os +import weakref +from typing import Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cuda.bindings import driver as cuda +from cutlass.cute.nvgpu import OperandMajorMode +from cutlass.cute.runtime import from_dlpack, make_fake_stream +import torch + +from cudnn.api_base import APIBase, TensorDesc +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor + +from ..moe_utils import MoEWeightMode +from .moe_grouped_gemm_dglu_dbias import MoEGroupedGemmDgluDbiasBf16Kernel + +_OUTPUT_DTYPES = [torch.bfloat16, torch.float16, torch.float32] + + +class GroupedGemmDgluBf16API(APIBase): + """Descriptor-first lifecycle API for the source BF16 dGLU kernel.""" + + def __init__( + self, + sample_a: torch.Tensor, + sample_c: torch.Tensor, + sample_d_row: torch.Tensor, + sample_padded_offsets: torch.Tensor, + sample_alpha: torch.Tensor, + sample_beta: torch.Tensor, + sample_prob: torch.Tensor, + sample_dprob: torch.Tensor, + sample_b: Optional[torch.Tensor] = None, + sample_dbias: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + b_shape: Optional[Tuple[int, ...]] = None, + b_dtype: Optional[torch.dtype] = None, + acc_dtype: torch.dtype = torch.float32, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + vector_f32: bool = False, + m_aligned: int = 256, + act_func: str = "dswiglu", + b_major: str = "k", + use_dynamic_sched: bool = False, + ) -> None: + super().__init__() + self._warn_experimental_api() + + if sample_b is not None and num_experts is None: + self.weight_mode = MoEWeightMode.DENSE + elif sample_b is None and num_experts is not None: + self.weight_mode = MoEWeightMode.DISCRETE + if b_shape is None or b_dtype is None: + raise ValueError("b_shape and b_dtype are required in discrete mode") + else: + raise ValueError("Provide sample_b for dense mode or (num_experts, b_shape, b_dtype) " "for discrete mode, but not both") + + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") + self.d_row_desc = self._make_tensor_desc(sample_d_row, name="sample_d_row") + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") + self.beta_desc = self._make_tensor_desc(sample_beta, name="sample_beta") + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") + self.dprob_desc = self._make_tensor_desc(sample_dprob, name="sample_dprob") + self.dbias_desc = self._make_tensor_desc(sample_dbias, name="sample_dbias") + + self._sample_offset_values = self._copy_values_to_host(sample_padded_offsets) + self._sample_offsets_ref = weakref.ref(sample_padded_offsets) + self._sample_offsets_version = int(sample_padded_offsets._version) + self._sample_data_ptrs = { + name: tensor.data_ptr() + for name, tensor in ( + ("sample_a", sample_a), + ("sample_b", sample_b), + ("sample_c", sample_c), + ("sample_d_row", sample_d_row), + ("sample_padded_offsets", sample_padded_offsets), + ("sample_alpha", sample_alpha), + ("sample_beta", sample_beta), + ("sample_prob", sample_prob), + ("sample_dprob", sample_dprob), + ("sample_dbias", sample_dbias), + ) + if tensor is not None + } + + self.expert_cnt = self.b_desc.shape[2] if self.weight_mode == MoEWeightMode.DENSE and self.b_desc.ndim == 3 else int(num_experts or 0) + self.b_shape = tuple(b_shape) if b_shape is not None else None + self.b_dtype = b_dtype if b_dtype is not None else self.b_desc.dtype + self.acc_dtype = acc_dtype + self.mma_tiler_mn = tuple(mma_tiler_mn) + self.use_2cta_instrs = self.mma_tiler_mn[0] == 256 + self.cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if self.use_2cta_instrs else (1, 1))) + self.vector_f32 = vector_f32 + self.m_aligned = m_aligned + self.act_func = act_func + self.b_major = b_major + self.use_dynamic_sched = use_dynamic_sched + self._has_dbias = self.dbias_desc is not None + self._kernel = MoEGroupedGemmDgluDbiasBf16Kernel + self._workspace: Optional[torch.Tensor] = None + self._compile_b_ptrs: Optional[torch.Tensor] = None + self._validated_offsets: dict[int, tuple] = {} + self._validated_pointer_values: dict[int, tuple] = {} + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + @staticmethod + def _expect_shape(desc: TensorDesc, expected: Tuple[int, ...], name: str) -> None: + if desc.shape != expected: + raise ValueError(f"{name} shape mismatch: expected {expected}, got {desc.shape}") + + @staticmethod + def _expect_stride(desc: TensorDesc, expected: Tuple[int, ...], name: str) -> None: + if desc.stride != expected: + raise ValueError(f"{name} must use the source-compatible layout with stride " f"{expected}, got {desc.stride}") + + @staticmethod + def _expect_device(desc: TensorDesc, device: torch.device, name: str) -> None: + if desc.device != device: + raise ValueError(f"{name} must be on {device}, got {desc.device}") + + @staticmethod + def _copy_values_to_host(tensor: torch.Tensor) -> Tuple[int, ...]: + return tuple(int(value) for value in tensor.detach().cpu().tolist()) + + @staticmethod + def _is_validation_cached(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> bool: + cached = cache.get(id(tensor)) + return bool(cached and cached[0]() is tensor and cached[1] == int(tensor._version) and cached[2] == extra) + + @staticmethod + def _remember_validation(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> None: + key = id(tensor) + + def discard(_reference, *, cache=cache, key=key): + cache.pop(key, None) + + cache[key] = (weakref.ref(tensor, discard), int(tensor._version), extra) + + @staticmethod + def _validate_offset_sequence(values: Tuple[int, ...], *, expert_cnt: int, tensor_m: int) -> None: + if len(values) != expert_cnt: + raise ValueError(f"padded_offsets length mismatch: expected {expert_cnt}, got {len(values)}") + previous = 0 + for index, value in enumerate(values): + if value < previous: + raise ValueError("padded_offsets must be a non-decreasing cumulative sum; " f"index {index} is {value} after {previous}") + if value % MoEGroupedGemmDgluDbiasBf16Kernel.FIX_PAD_SIZE != 0: + raise ValueError(f"padded_offsets[{index}] must be 256-aligned, got {value}") + previous = value + if not values or values[-1] <= 0 or values[-1] > tensor_m: + raise ValueError(f"padded_offsets last value must be in [1, {tensor_m}], got " f"{values[-1] if values else None}") + + def _validate_offsets_once(self, offsets: torch.Tensor, *, tensor_m: int) -> None: + extra = (self.expert_cnt, tensor_m) + if self._is_validation_cached(self._validated_offsets, offsets, extra): + return + values = self._copy_values_to_host(offsets) + self._validate_offset_sequence(values, expert_cnt=self.expert_cnt, tensor_m=tensor_m) + self._remember_validation(self._validated_offsets, offsets, extra) + + def _validate_pointer_values_once(self, b_ptrs: torch.Tensor) -> None: + if self._is_validation_cached(self._validated_pointer_values, b_ptrs, self.expert_cnt): + return + values = self._copy_values_to_host(b_ptrs) + if any(value == 0 or value % 16 != 0 for value in values): + raise ValueError("b_ptrs entries must be non-null and 16-byte aligned") + self._remember_validation(self._validated_pointer_values, b_ptrs, self.expert_cnt) + + @staticmethod + def _validate_data_alignment(tensor: torch.Tensor, name: str) -> None: + if tensor.data_ptr() % 16 != 0: + raise ValueError(f"{name} data pointer must be 16-byte aligned") + + @staticmethod + def _validate_pointer_array_alignment(tensor: torch.Tensor) -> None: + if tensor.data_ptr() % 8 != 0: + raise ValueError("b_ptrs data pointer must be 8-byte aligned") + + @staticmethod + def _record_pointer_stream(b_ptrs: torch.Tensor, stream: cuda.CUstream) -> None: + handle = int(stream) + torch_current = torch.cuda.current_stream(b_ptrs.device) + torch_default = torch.cuda.default_stream(b_ptrs.device) + if handle == torch_current.cuda_stream: + launch_stream = torch_current + elif handle == torch_default.cuda_stream: + launch_stream = torch_default + else: + launch_stream = torch.cuda.ExternalStream(handle, device=b_ptrs.device) + b_ptrs.record_stream(launch_stream) + + def check_support(self) -> bool: + if self.a_desc.ndim != 3: + raise ValueError(f"sample_a must be rank-3, got {self.a_desc.shape}") + tensor_m, k, one = self.a_desc.shape + if one != 1: + raise ValueError(f"sample_a trailing dimension must be 1, got {one}") + + if self.weight_mode == MoEWeightMode.DENSE: + if self.b_desc.ndim != 3: + raise ValueError(f"sample_b must be rank-3, got {self.b_desc.shape}") + n, b_k, experts = self.b_desc.shape + if b_k != k: + raise ValueError(f"sample_b K dimension ({b_k}) must match sample_a ({k})") + if experts != self.expert_cnt: + raise ValueError("sample_b expert dimension is inconsistent") + self._expect_stride(self.b_desc, (k, 1, n * k), "sample_b") + else: + if len(self.b_shape) not in (2, 3): + raise ValueError(f"b_shape must be rank-2 or rank-3, got {self.b_shape}") + n, b_k = self.b_shape[:2] + if len(self.b_shape) == 3 and self.b_shape[2] != 1: + raise ValueError(f"b_shape trailing dimension must be 1, got {self.b_shape}") + if b_k != k: + raise ValueError(f"b_shape K dimension ({b_k}) must match sample_a ({k})") + if n <= 0 or n % 32 != 0: + raise ValueError(f"N must be positive and divisible by 32, got {n}") + + two_n = 2 * n + self._expect_shape(self.c_desc, (tensor_m, two_n, 1), "sample_c") + self._expect_shape(self.d_row_desc, (tensor_m, two_n, 1), "sample_d_row") + self._expect_shape(self.padded_offsets_desc, (self.expert_cnt,), "sample_padded_offsets") + self._expect_shape(self.alpha_desc, (self.expert_cnt,), "sample_alpha") + self._expect_shape(self.beta_desc, (self.expert_cnt,), "sample_beta") + self._expect_shape(self.prob_desc, (tensor_m, 1, 1), "sample_prob") + self._expect_shape(self.dprob_desc, (tensor_m, 1, 1), "sample_dprob") + + self._expect_stride(self.a_desc, (k, 1, tensor_m * k), "A tensor") + self._expect_stride(self.c_desc, (two_n, 1, tensor_m * two_n), "sample_c") + self._expect_stride(self.d_row_desc, (two_n, 1, tensor_m * two_n), "sample_d_row") + self._expect_stride(self.padded_offsets_desc, (1,), "sample_padded_offsets") + self._expect_stride(self.alpha_desc, (1,), "sample_alpha") + self._expect_stride(self.beta_desc, (1,), "sample_beta") + self._expect_stride(self.prob_desc, (1, 1, 1), "sample_prob") + self._expect_stride(self.dprob_desc, (1, 1, 1), "sample_dprob") + + self._check_dtype(self.a_desc, torch.bfloat16, "sample_a") + if self.b_desc is not None: + self._check_dtype(self.b_desc, torch.bfloat16, "sample_b") + self._check_dtype(self.b_dtype, torch.bfloat16, "b_dtype") + self._check_dtype(self.c_desc, _OUTPUT_DTYPES, "sample_c") + self._check_dtype(self.d_row_desc, _OUTPUT_DTYPES, "sample_d_row") + self._check_dtype(self.padded_offsets_desc, torch.int32, "sample_padded_offsets") + self._check_dtype(self.alpha_desc, torch.float32, "sample_alpha") + self._check_dtype(self.beta_desc, torch.float32, "sample_beta") + self._check_dtype(self.prob_desc, torch.float32, "sample_prob") + self._check_dtype(self.dprob_desc, torch.float32, "sample_dprob") + + device = self.a_desc.device + for desc, name in ( + (self.c_desc, "sample_c"), + (self.d_row_desc, "sample_d_row"), + (self.padded_offsets_desc, "sample_padded_offsets"), + (self.alpha_desc, "sample_alpha"), + (self.beta_desc, "sample_beta"), + (self.prob_desc, "sample_prob"), + (self.dprob_desc, "sample_dprob"), + ): + self._expect_device(desc, device, name) + if self.b_desc is not None: + self._expect_device(self.b_desc, device, "sample_b") + + if self.dbias_desc is not None: + self._expect_shape(self.dbias_desc, (self.expert_cnt, two_n, 1), "sample_dbias") + self._expect_stride(self.dbias_desc, (two_n, 1, 1), "sample_dbias") + self._check_dtype(self.dbias_desc, torch.bfloat16, "sample_dbias") + self._expect_device(self.dbias_desc, device, "sample_dbias") + + for name, pointer in self._sample_data_ptrs.items(): + if pointer % 16 != 0: + raise ValueError(f"{name} data pointer must be 16-byte aligned") + + if self.acc_dtype != torch.float32: + raise ValueError(f"acc_dtype must be torch.float32, got {self.acc_dtype}") + if self.m_aligned != 256: + raise ValueError(f"m_aligned must be 256, got {self.m_aligned}") + if tensor_m % 256 != 0: + raise ValueError(f"sample_a M dimension must be 256-aligned, got {tensor_m}") + if self.act_func not in ("dswiglu", "dgeglu"): + raise ValueError(f"act_func must be 'dswiglu' or 'dgeglu', got {self.act_func}") + if self.b_major not in ("k", "n"): + raise ValueError(f"b_major must be 'k' or 'n', got {self.b_major}") + if self.expert_cnt <= 0 or self.expert_cnt > 1024: + raise ValueError(f"expert count must be in [1, 1024], got {self.expert_cnt}") + + self._validate_offset_sequence(self._sample_offset_values, expert_cnt=self.expert_cnt, tensor_m=tensor_m) + sample_offsets = self._sample_offsets_ref() + if sample_offsets is not None and int(sample_offsets._version) == self._sample_offsets_version: + self._remember_validation(self._validated_offsets, sample_offsets, (self.expert_cnt, tensor_m)) + elif sample_offsets is not None: + self._validate_offsets_once(sample_offsets, tensor_m=tensor_m) + + if not self._kernel.can_implement( + _convert_to_cutlass_data_type(torch.bfloat16), + _convert_to_cutlass_data_type(self.c_desc.dtype), + _convert_to_cutlass_data_type(self.d_row_desc.dtype), + _convert_to_cutlass_data_type(self.acc_dtype), + self.use_2cta_instrs, + self.mma_tiler_mn, + self.cluster_shape_mn, + tensor_m, + n, + k, + self.expert_cnt, + "k", + self.b_major, + "n", + self.m_aligned, + self.act_func, + ): + raise ValueError("Unsupported BF16 grouped GEMM dGLU configuration") + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + major, minor = torch.cuda.get_device_capability(self.a_desc.device) + capability = major * 10 + minor + if capability < 100: + raise RuntimeError(f"GroupedGemmDgluSm100 requires SM100+, found SM{capability} on {self.a_desc.device}") + self._is_supported = True + return True + + def compile(self) -> None: + self._ensure_support_checked() + if self._compiled_kernel is not None: + return + + kernel = self._kernel( + acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), + use_2cta_instrs=self.use_2cta_instrs, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + vectorized_f32=self.vector_f32, + expert_cnt=self.expert_cnt, + weight_mode=self.weight_mode, + use_dynamic_sched=self.use_dynamic_sched, + act_func=self.act_func, + ) + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - self.num_cluster_overlap_margin + if max_active_clusters <= 0: + raise ValueError("max_active_clusters must be > 0 after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + + workspace_bytes = kernel.get_workspace_bytes() + self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device=self.a_desc.device) + if self._workspace.data_ptr() % 128 != 0: + raise RuntimeError("workspace allocation must be 128-byte aligned") + workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + + a_fake = self._make_fake_cute_compact_tensor( + self.a_desc.dtype, + self.a_desc.shape, + self.a_desc.stride_order, + dynamic_mode=0, + divisibility=256, + ) + c_fake = self._make_fake_cute_compact_tensor( + self.c_desc.dtype, + self.c_desc.shape, + self.c_desc.stride_order, + dynamic_mode=0, + divisibility=256, + ) + d_fake = self._make_fake_cute_compact_tensor( + self.d_row_desc.dtype, + self.d_row_desc.shape, + self.d_row_desc.stride_order, + dynamic_mode=0, + divisibility=256, + ) + valid_m = cute.sym_int(divisibility=256) + prob_fake = self._make_fake_cute_tensor(self.prob_desc.dtype, (valid_m, 1, 1), self.prob_desc.stride) + dprob_fake = self._make_fake_cute_tensor(self.dprob_desc.dtype, (valid_m, 1, 1), self.dprob_desc.stride) + + if self.weight_mode == MoEWeightMode.DENSE: + b_fake = self._make_fake_cute_tensor_from_desc(self.b_desc) + n_value = cutlass.Int32(0) + k_value = cutlass.Int32(0) + b_stride = cutlass.Int64(0) + b_major_mode = OperandMajorMode.K + else: + self._compile_b_ptrs = torch.empty((self.expert_cnt,), dtype=torch.int64, device=self.a_desc.device) + self._validate_pointer_array_alignment(self._compile_b_ptrs) + b_fake = from_dlpack(self._compile_b_ptrs, assumed_align=8).iterator + n, k = self.b_shape[:2] + n_value = cutlass.Int32(n) + k_value = cutlass.Int32(k) + b_stride = cutlass.Int64(k if self.b_major == "k" else n) + b_major_mode = OperandMajorMode.K if self.b_major == "k" else OperandMajorMode.MN + + raw_compiled = cute.compile( + kernel, + a=a_fake, + b=b_fake, + n=n_value, + k=k_value, + b_stride_size=b_stride, + b_major_mode=b_major_mode, + workspace_ptr=workspace_ptr, + c=c_fake, + d=d_fake, + padded_offsets=self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc), + alpha=self._make_fake_cute_tensor_from_desc(self.alpha_desc), + beta=self._make_fake_cute_tensor_from_desc(self.beta_desc), + prob=prob_fake, + dprob=dprob_fake, + linear_offset=cutlass.Float32(0.0), + dbias_tensor=self._make_fake_cute_tensor_from_desc(self.dbias_desc), + max_active_clusters=max_active_clusters, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + cached_n, cached_k, cached_b_stride = n_value, k_value, b_stride + + def tensor_api( + a_tensor, + c_tensor, + d_row_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + dprob_tensor, + b_tensor, + b_ptrs, + dbias_tensor, + stream, + linear_offset, + ) -> None: + b_arg = b_tensor if self.weight_mode == MoEWeightMode.DENSE else int(b_ptrs.data_ptr()) + raw_compiled( + a_tensor, + b_arg, + cached_n, + cached_k, + cached_b_stride, + workspace_ptr, + c_tensor, + d_row_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + dprob_tensor, + cutlass.Float32(linear_offset), + dbias_tensor, + stream, + ) + + self._compiled_kernel = tensor_api + + def _validate_live_tensor(self, tensor: torch.Tensor, sample: TensorDesc, name: str, *, dynamic_m: bool = False) -> TensorDesc: + desc = self._make_tensor_desc(tensor, name=name) + if desc.dtype != sample.dtype: + raise ValueError(f"{name} dtype mismatch: expected {sample.dtype}, got {desc.dtype}") + if desc.device != sample.device: + raise ValueError(f"{name} device mismatch: expected {sample.device}, got {desc.device}") + if dynamic_m: + if desc.shape[1:] != sample.shape[1:]: + raise ValueError(f"{name} shape suffix mismatch: expected {sample.shape[1:]}, got {desc.shape[1:]}") + if desc.stride_order != sample.stride_order: + raise ValueError(f"{name} layout mismatch: expected stride order {sample.stride_order}, got {desc.stride_order}") + elif desc.shape != sample.shape or desc.stride != sample.stride: + raise ValueError(f"{name} descriptor mismatch: expected shape/stride {sample.shape}/{sample.stride}, " f"got {desc.shape}/{desc.stride}") + return desc + + def execute( + self, + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_row_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + beta_tensor: torch.Tensor, + prob_tensor: torch.Tensor, + dprob_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor] = None, + b_ptrs: Optional[torch.Tensor] = None, + dbias_tensor: Optional[torch.Tensor] = None, + linear_offset: float = 0.0, + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + current_stream = self._get_default_stream(current_stream) + if self._compiled_kernel is None: + raise RuntimeError("Kernel not compiled; call compile() first") + + a_desc = self._validate_live_tensor(a_tensor, self.a_desc, "a_tensor", dynamic_m=True) + tensor_m, k, _ = a_desc.shape + if tensor_m % 256 != 0: + raise ValueError(f"a_tensor M dimension must be 256-aligned, got {tensor_m}") + c_desc = self._validate_live_tensor(c_tensor, self.c_desc, "c_tensor", dynamic_m=True) + d_desc = self._validate_live_tensor(d_row_tensor, self.d_row_desc, "d_row_tensor", dynamic_m=True) + prob_desc = self._validate_live_tensor(prob_tensor, self.prob_desc, "prob_tensor", dynamic_m=True) + dprob_desc = self._validate_live_tensor(dprob_tensor, self.dprob_desc, "dprob_tensor", dynamic_m=True) + self._validate_live_tensor(padded_offsets, self.padded_offsets_desc, "padded_offsets") + self._validate_live_tensor(alpha_tensor, self.alpha_desc, "alpha_tensor") + self._validate_live_tensor(beta_tensor, self.beta_desc, "beta_tensor") + + two_n = c_desc.shape[1] + self._expect_shape(c_desc, (tensor_m, two_n, 1), "c_tensor") + self._expect_shape(d_desc, (tensor_m, two_n, 1), "d_row_tensor") + self._expect_shape(prob_desc, (tensor_m, 1, 1), "prob_tensor") + self._expect_shape(dprob_desc, (tensor_m, 1, 1), "dprob_tensor") + self._expect_stride(a_desc, (k, 1, tensor_m * k), "a_tensor") + self._expect_stride(c_desc, (two_n, 1, tensor_m * two_n), "c_tensor") + self._expect_stride(d_desc, (two_n, 1, tensor_m * two_n), "d_row_tensor") + self._expect_stride(prob_desc, (1, 1, 1), "prob_tensor") + self._expect_stride(dprob_desc, (1, 1, 1), "dprob_tensor") + self._validate_offsets_once(padded_offsets, tensor_m=tensor_m) + + for tensor, name in ( + (a_tensor, "a_tensor"), + (c_tensor, "c_tensor"), + (d_row_tensor, "d_row_tensor"), + (padded_offsets, "padded_offsets"), + (alpha_tensor, "alpha_tensor"), + (beta_tensor, "beta_tensor"), + (prob_tensor, "prob_tensor"), + (dprob_tensor, "dprob_tensor"), + ): + self._validate_data_alignment(tensor, name) + + if self._has_dbias: + if dbias_tensor is None: + raise ValueError("dbias_tensor is required because the API was compiled with sample_dbias") + self._validate_live_tensor(dbias_tensor, self.dbias_desc, "dbias_tensor") + self._validate_data_alignment(dbias_tensor, "dbias_tensor") + elif dbias_tensor is not None: + raise ValueError("dbias_tensor must be omitted because the API was compiled without sample_dbias") + + if self.weight_mode == MoEWeightMode.DENSE: + if b_tensor is None or b_ptrs is not None: + raise ValueError("Dense execution requires b_tensor and forbids b_ptrs") + self._validate_live_tensor(b_tensor, self.b_desc, "b_tensor") + self._validate_data_alignment(b_tensor, "b_tensor") + else: + if b_tensor is not None or b_ptrs is None: + raise ValueError("Discrete execution requires b_ptrs and forbids b_tensor") + _require_pointer_tensor(b_ptrs, "b_ptrs", self.expert_cnt) + if b_ptrs.device != self.a_desc.device: + raise ValueError(f"b_ptrs must be on the same device as a_tensor ({self.a_desc.device}), got {b_ptrs.device}") + if b_ptrs.data_ptr() % 8 != 0: + raise ValueError("b_ptrs data pointer must be 8-byte aligned") + self._validate_pointer_values_once(b_ptrs) + self._record_pointer_stream(b_ptrs, current_stream) + + self._compiled_kernel( + a_tensor, + c_tensor, + d_row_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + dprob_tensor, + b_tensor, + b_ptrs, + dbias_tensor, + current_stream, + linear_offset, + ) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dglu/_blockscaled_api.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/_blockscaled_api.py new file mode 100644 index 000000000..3fc5e4ed0 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/_blockscaled_api.py @@ -0,0 +1,1294 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Unified API for Grouped GEMM dGLU Backward Kernel (SM100+) + +This module provides a single API class that supports both contiguous (dense) +and discrete weight modes for block-scaled grouped GEMM with dGLU activation +gradient (dSwiGLU / dGeGLU) in MoE (Mixture of Experts) workloads. + +Dense mode + All expert weights are packed contiguously in a 3-D tensor (N, K, L). + Callers supply ``sample_b`` and ``sample_sfb``. + +Discrete mode + Each expert has its own memory allocation. Callers supply + ``num_experts``, ``b_shape``, ``b_dtype``, and per-expert pointer arrays + at execution time. +""" + +from .moe_blockscaled_grouped_gemm_dglu_dbias import BlockScaledMoEGroupedGemmDgluDbiasKernel +from ..moe_utils import MoEWeightMode +from cuda.bindings import driver as cuda +import os +import torch +from typing import Tuple, Optional + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import OperandMajorMode +from cutlass.cute.runtime import from_dlpack, make_fake_stream + +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.api_base import APIBase, ceil_div, is_power_of_2 + + +class GroupedGemmDgluBlockScaledAPI(APIBase): + """Unified API for grouped GEMM dGLU backward operation on SM100+ GPUs. + + This kernel performs block-scaled grouped GEMM with dGLU activation + gradient (dSwiGLU or dGeGLU), designed for MoE workloads. It supports + both dense (contiguous) and discrete (per-expert pointer) weight layouts + through the ``BlockScaledMoEGroupedGemmDgluDbiasKernel``. + + Weight mode is auto-detected from the constructor arguments: + + - **Dense**: provide ``sample_b`` and ``sample_sfb``. + - **Discrete**: provide ``num_experts``, ``b_shape``, and ``b_dtype``. + + Example:: + + # Dense mode + api = GroupedGemmDgluSm100( + sample_a=a, sample_c=c, + sample_d_row=d_row, sample_d_col=d_col, + sample_sfa=sfa, sample_padded_offsets=offsets, + sample_alpha=alpha, sample_beta=beta, + sample_prob=prob, sample_dprob=dprob, + sample_b=b, sample_sfb=sfb, + ) + + # Discrete mode + api = GroupedGemmDgluSm100( + sample_a=a, sample_c=c, + sample_d_row=d_row, sample_d_col=d_col, + sample_sfa=sfa, sample_padded_offsets=offsets, + sample_alpha=alpha, sample_beta=beta, + sample_prob=prob, sample_dprob=dprob, + num_experts=8, b_shape=(n, k), b_dtype=torch.uint8, + ) + + api.check_support() + api.compile() + api.execute(...) + """ + + def __init__( + self, + sample_a: torch.Tensor, + sample_c: torch.Tensor, + sample_d_row: torch.Tensor, + sample_d_col: torch.Tensor, + sample_sfa: torch.Tensor, + sample_padded_offsets: torch.Tensor, + sample_alpha: torch.Tensor, + sample_beta: torch.Tensor, + sample_prob: torch.Tensor, + sample_dprob: torch.Tensor, + # Dense mode (contiguous) -- provide these. sample_dbias is optional: + sample_b: Optional[torch.Tensor] = None, + sample_sfb: Optional[torch.Tensor] = None, + sample_dbias: Optional[torch.Tensor] = None, + # Discrete mode -- provide these instead: + num_experts: Optional[int] = None, + b_shape: Optional[Tuple[int, ...]] = None, + b_dtype: Optional[torch.dtype] = None, + # Optional quantization output arguments + sample_sfd_row: Optional[torch.Tensor] = None, + sample_sfd_col: Optional[torch.Tensor] = None, + sample_amax: Optional[torch.Tensor] = None, + sample_norm_const: Optional[torch.Tensor] = None, + # Configuration + acc_dtype: torch.dtype = torch.float32, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "dswiglu", + b_major: str = "k", + epilogue_op: Optional[str] = None, + use_dynamic_sched: bool = False, + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + ): + """Initialize the GroupedGemmDgluSm100 API. + + :param sample_a: Sample A tensor (valid_m, k, 1) + :param sample_c: Sample C tensor -- forward activations (valid_m, n, 1) + :param sample_d_row: Sample D row output tensor (valid_m, n*2, 1) + :param sample_d_col: Sample D col output tensor (valid_m, n*2, 1) + :param sample_sfa: Sample scale factor A tensor + :param sample_padded_offsets: End offset for each expert after padding + :param sample_alpha: Per-group alpha scaling factors + :param sample_beta: Per-group beta scaling factors + :param sample_prob: Per-row probability tensor (valid_m, 1, 1) + :param sample_dprob: Gradient of probability tensor (valid_m, 1, 1), must be zero-initialized + :param sample_b: (Dense) Sample B tensor (n, k, l) + :param sample_sfb: (Dense) Sample scale factor B tensor + :param sample_dbias: Optional dbias output tensor (expert_cnt, 2*n, 1) + :param num_experts: (Discrete) Number of experts + :param b_shape: (Discrete) Shape of a single expert B tensor, e.g. (n, k) + :param b_dtype: (Discrete) Data type of B tensors + :param sample_sfd_row: Optional row scale factor for D + :param sample_sfd_col: Optional column scale factor for D + :param sample_amax: Optional amax tensor for quantization, shape (expert_cnt, 2, 1) + :param sample_norm_const: Optional normalization constant + :param acc_dtype: Accumulator data type + :param mma_tiler_mn: MMA tiler shape (M, N) + :param cluster_shape_mn: Cluster shape (M, N) + :param sf_vec_size: Scale factor vector size + :param vector_f32: Use vectorized f32 operations + :param m_aligned: Alignment for group M dimension + :param discrete_col_sfd: Generate discrete col-major scale factor tensor + :param act_func: Activation function, one of "dswiglu" or "dgeglu" + :param b_major: Major dimension for B tensor, one of "k" or "n" + :param epilogue_op: Optional epilogue operation. Valid: None, "none", "identity", "relu", "srelu" + :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing + :param linear_offset: Compile-time linear offset for dGeGLU. When None, + defaults to 1.0 for dGeGLU and 0.0 for dSwiGLU. + Ignored when ``act_func == "dswiglu"``. + :param geglu_alpha: Compile-time dGeGLU GeGLU alpha. Ignored when + ``act_func == "dswiglu"``. + :param glu_clamp_max: Compile-time dGeGLU upper clamp. Ignored when + ``act_func == "dswiglu"``. + :param glu_clamp_min: Compile-time dGeGLU lower clamp. Ignored when + ``act_func == "dswiglu"``. + """ + super().__init__() + + self._warn_experimental_api() + self._logger.debug("Entering __init__") + + # ---- Weight mode auto-detection ---- + if sample_b is not None and num_experts is None: + self.weight_mode = MoEWeightMode.DENSE + if sample_sfb is None: + raise ValueError("sample_sfb is required when sample_b is provided (dense mode)") + elif num_experts is not None and sample_b is None: + self.weight_mode = MoEWeightMode.DISCRETE + if b_shape is None or b_dtype is None: + raise ValueError("b_shape and b_dtype are required in discrete mode") + else: + raise ValueError("Provide either (sample_b, sample_sfb) for dense mode " "or (num_experts, b_shape, b_dtype) for discrete mode, but not both.") + + # ---- Common tensor descriptors ---- + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") + self.d_row_desc = self._make_tensor_desc(sample_d_row, name="sample_d_row") + self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") + self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") + self.beta_desc = self._make_tensor_desc(sample_beta, name="sample_beta") + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") + self.dprob_desc = self._make_tensor_desc(sample_dprob, name="sample_dprob") + self.dbias_desc = self._make_tensor_desc(sample_dbias, name="sample_dbias") + + self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row") + self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col") + self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") + self.norm_const_desc = self._unpad_tensor_to_ndim( + self._make_tensor_desc(sample_norm_const, name="sample_norm_const"), + 1, + "norm_const", + ) + + # ---- Mode-specific state ---- + if self.weight_mode == MoEWeightMode.DENSE: + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") + self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb") + self.expert_cnt = self.padded_offsets_desc.shape[0] + else: + self._value_error_if(num_experts == 0, "num_experts must be > 0") + self.expert_cnt = num_experts + self.b_shape = b_shape + self.b_dtype = b_dtype + self.b_major = b_major + self._value_error_if( + self.padded_offsets_desc.shape[0] != self.expert_cnt, + f"padded_offsets length ({self.padded_offsets_desc.shape[0]}) " f"must equal num_experts ({self.expert_cnt})", + ) + + # ---- Configuration ---- + self.acc_dtype = acc_dtype + self.mma_tiler_mn = mma_tiler_mn + self.use_2cta_instrs = mma_tiler_mn[0] == 256 + if cluster_shape_mn is None: + self.cluster_shape_mn = (2, 1) if self.use_2cta_instrs else (1, 1) + else: + self.cluster_shape_mn = cluster_shape_mn + self.sf_vec_size = sf_vec_size + self.vector_f32 = vector_f32 + self.m_aligned = m_aligned + self.discrete_col_sfd = discrete_col_sfd + self.act_func = act_func + if self.weight_mode == MoEWeightMode.DENSE: + self.b_major = b_major # stored for both modes + + # Epilogue operation + if epilogue_op in [None, "none", "identity"]: + self.epilogue_op = lambda x: x + elif epilogue_op == "relu": + self.epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) + elif epilogue_op == "srelu": + self.epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 + else: + raise ValueError(f"Invalid epilogue operation: {epilogue_op}. " f"Valid values: None, 'none', 'identity', 'relu', 'srelu'") + + self.use_dynamic_sched = use_dynamic_sched + if linear_offset is None: + self.linear_offset = 1.0 if self.act_func == "dgeglu" else 0.0 + else: + self.linear_offset = float(linear_offset) + self.geglu_alpha = geglu_alpha + self.glu_clamp_max = glu_clamp_max + self.glu_clamp_min = glu_clamp_min + + self._interpret_uint8_as_fp4x2 = True + self._has_dbias = self.dbias_desc is not None + self._kernel = BlockScaledMoEGroupedGemmDgluDbiasKernel + + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + self._logger.debug(f"setting num_cluster_overlap_margin: {self.num_cluster_overlap_margin}") + + self._workspace = None + + self._logger.debug("__init__ completed") + + @staticmethod + def _record_pointer_stream(pointers: torch.Tensor, current_stream: cuda.CUstream) -> None: + handle = int(current_stream) + torch_current = torch.cuda.current_stream(pointers.device) + torch_default = torch.cuda.default_stream(pointers.device) + if handle == torch_current.cuda_stream: + launch_stream = torch_current + elif handle == torch_default.cuda_stream: + launch_stream = torch_default + else: + launch_stream = torch.cuda.ExternalStream(handle, device=pointers.device) + pointers.record_stream(launch_stream) + + # --------------------------------------------------------------------- # + # check_support + # --------------------------------------------------------------------- # + + def check_support(self) -> bool: + """Check if the kernel configuration is supported. + + :return: True if supported, raises exception otherwise + """ + self._logger.debug("Entering check_support") + + # ---- SFD group validation ---- + all_none = all(x is None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) + all_provided = all(x is not None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) + self._value_error_if( + not (all_none or all_provided), + "sfd_row_desc, sfd_col_desc, and norm_const_desc must be all None or all not None", + ) + self._user_requested_sfd = all_provided + + # ---- Shapes and strides ---- + self._logger.debug("Checking tensor shapes and strides") + tensor_m, k, _one = self._tensor_shape(self.a_desc, name="sample_a") + + if self.weight_mode == MoEWeightMode.DENSE: + n, _, l = self._tensor_shape(self.b_desc, name="sample_b") + else: + # Discrete: extract n, k from b_shape + if len(self.b_shape) == 2: + n, b_k = self.b_shape + else: + n, b_k, _ = self.b_shape + self._value_error_if(b_k != k, f"B K dimension ({b_k}) must match A K dimension ({k})") + l = self.expert_cnt # for shape checks that use l + + n_out = 2 * n + + self._value_error_if( + n % 32 != 0, + f"N must be divisible by 32 for dGLU (32-column input/gate interleaving), got N={n}", + ) + + self._check_tensor_shape(self.a_desc, (tensor_m, k, 1), "A") + if self.weight_mode == MoEWeightMode.DENSE: + self._check_tensor_shape(self.b_desc, (n, k, l), "B") + self._check_tensor_shape(self.c_desc, (tensor_m, n_out, 1), "C") + self._check_tensor_shape(self.d_row_desc, (tensor_m, n_out, 1), "D_row") + self._check_tensor_shape(self.d_col_desc, (tensor_m, n_out, 1), "D_col") + + rest_k = ceil_div(ceil_div(k, self.sf_vec_size), 4) + self._check_tensor_shape(self.sfa_desc, (32, 4, ceil_div(tensor_m, 128), 4, rest_k, 1), "SFA") + if self.weight_mode == MoEWeightMode.DENSE: + self._check_tensor_shape(self.sfb_desc, (32, 4, ceil_div(n, 128), 4, rest_k, l), "SFB") + + # SFD uses n_out dimension since D has n_out columns + rest_n_out = ceil_div(ceil_div(n_out, self.sf_vec_size), 4) + self._check_tensor_shape( + self.sfd_row_desc, + (32, 4, ceil_div(tensor_m, 128), 4, rest_n_out, 1), + "SFD_row", + ) + rest_m = ceil_div(ceil_div(tensor_m, self.sf_vec_size), 4) + self._check_tensor_shape(self.sfd_col_desc, (32, 4, ceil_div(n_out, 128), 4, rest_m, 1), "SFD_col") + + self._check_tensor_shape(self.alpha_desc, (self.expert_cnt,), "alpha") + self._check_tensor_shape(self.beta_desc, (self.expert_cnt,), "beta") + self._check_tensor_shape(self.prob_desc, (tensor_m, 1, 1), "prob") + self._check_tensor_shape(self.dprob_desc, (tensor_m, 1, 1), "dprob") + self._check_tensor_shape(self.dbias_desc, (self.expert_cnt, n_out, 1), "dbias") + self._check_tensor_shape(self.amax_desc, (self.expert_cnt, 2, 1), "amax") + self._check_tensor_shape(self.norm_const_desc, (1,), "norm_const") + self._check_tensor_shape(self.padded_offsets_desc, (self.expert_cnt,), "padded_offsets") + + # Strides + _ = self._check_tensor_stride( + self.a_desc, + stride=[(k, 1, tensor_m * k)], + extra_error_msg="A must have k-major layout", + ) + if self.weight_mode == MoEWeightMode.DENSE: + if self._is_fp8(self.a_desc): + _ = self._check_tensor_stride( + self.b_desc, + stride=[(k, 1, n * k), (1, n, n * k)], + extra_error_msg="For fp8 ab_dtype, B must have k- or n-major layout", + ) + else: + _ = self._check_tensor_stride( + self.b_desc, + stride=[(k, 1, n * k)], + extra_error_msg="For fp4 ab_dtype, B must have k-major layout", + ) + _ = self._check_tensor_stride( + self.c_desc, + stride=[(n_out, 1, tensor_m * n_out)], + extra_error_msg="C must have n-major layout", + ) + _ = self._check_tensor_stride( + self.d_row_desc, + stride=[(n_out, 1, tensor_m * n_out)], + extra_error_msg="D_row must have n-major layout", + ) + _ = self._check_tensor_stride( + self.d_col_desc, + stride=[(n_out, 1, tensor_m * n_out)], + extra_error_msg="D_col must have n-major layout", + ) + + # ---- Data types ---- + self._logger.debug("Checking data types") + self.ab_dtype = self._check_dtype( + self.a_desc, + dtype=[ + torch.float4_e2m1fn_x2, + torch.uint8, + torch.float8_e5m2, + torch.float8_e4m3fn, + ], + name="A/B", + ) + if self.weight_mode == MoEWeightMode.DENSE: + self._check_dtype( + self.b_desc, + dtype=self.ab_dtype, + name="B", + extra_error_msg="B must have the same dtype as A", + ) + else: + self._value_error_if( + self.b_dtype != self.ab_dtype, + f"b_dtype ({self.b_dtype}) must match A dtype ({self.ab_dtype})", + ) + + self.sf_dtype = self._check_dtype( + self.sfa_desc, + dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], + name="SFA/SFB/SFD", + ) + if self.weight_mode == MoEWeightMode.DENSE: + self._check_dtype( + self.sfb_desc, + dtype=self.sf_dtype, + name="SFB", + extra_error_msg="SFB must have the same dtype as SFA", + ) + self._check_dtype( + self.sfd_row_desc, + dtype=self.sf_dtype, + name="SFD_row", + extra_error_msg="SFD_row must have the same dtype as SFA", + ) + self._check_dtype( + self.sfd_col_desc, + dtype=self.sf_dtype, + name="SFD_col", + extra_error_msg="SFD_col must have the same dtype as SFA", + ) + + self._value_error_if( + self.sf_vec_size not in [16, 32], + f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}", + ) + self._value_error_if( + self.sf_dtype in [torch.float8_e4m3fn] and self.sf_vec_size == 32, + f"sf_dtype {self.sf_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", + ) + self._value_error_if( + self._is_fp8(self.ab_dtype) and self.sf_vec_size == 16, + f"ab_dtype {self.ab_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", + ) + + self._check_dtype( + self.acc_dtype, + dtype=torch.float32, + name="Accumulator", + extra_error_msg="Accumulator must be float32", + ) + self._check_dtype( + self.prob_desc, + dtype=torch.float32, + name="Prob", + extra_error_msg="Prob must be float32", + ) + self._check_dtype( + self.dprob_desc, + dtype=torch.float32, + name="Dprob", + extra_error_msg="Dprob must be float32", + ) + self._check_dtype( + self.dbias_desc, + dtype=torch.bfloat16, + name="Dbias", + extra_error_msg="dbias must be bfloat16", + ) + self.c_dtype = self._check_dtype( + self.c_desc, + dtype=[torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2], + name="C", + ) + if self._is_fp8(self.c_dtype) and self.vector_f32: + raise ValueError("Invalid configuration: fp8 c_dtype and vector_f32 is not supported. " "Please use vector_f32=False or c_dtype=bfloat16 instead") + + if self._is_fp4x2(self.ab_dtype): + self.d_dtype = self._check_dtype( + self.d_row_desc, + dtype=[torch.float16, torch.bfloat16, torch.float32], + name="D_row", + extra_error_msg="D_row must be fp16, bf16, or float32 when ab_dtype is fp4", + ) + elif self._is_fp8(self.ab_dtype): + self.d_dtype = self._check_dtype( + self.d_row_desc, + dtype=[ + torch.float8_e4m3fn, + torch.float8_e5m2, + ], + name="D_row", + extra_error_msg="D_row must be fp8 dtype when ab_dtype is fp8", + ) + else: + raise NotImplementedError(f"Invalid ab_dtype: {self.ab_dtype}, expected fp4 or fp8") + self._check_dtype( + self.d_col_desc, + dtype=self.d_dtype, + name="D_col", + extra_error_msg="D_col must have the same dtype as D_row", + ) + + # ---- SFD generation logic ---- + kernel_generate_sfd = self._is_fp8(self.ab_dtype) and self.sf_dtype == torch.float8_e8m0fnu and self._is_fp8(self.d_dtype) + self._value_error_if( + kernel_generate_sfd and not self._user_requested_sfd, + "sfd_row, sfd_col, and norm_const are required for FP8 input/FP8 output with sf_dtype=torch.float8_e8m0fnu", + ) + if not kernel_generate_sfd and self._user_requested_sfd: + self._logger.warning( + "sfd_row/sfd_col/norm_const were provided, but this configuration does not generate SFD outputs; " "the tensors will be ignored by the kernel", + ) + self.generate_sfd = kernel_generate_sfd + if self.discrete_col_sfd and not self.generate_sfd: + self._logger.warning("discrete_col_sfd is True but generate_sfd is False, discrete_col_sfd will be ignored") + self.discrete_col_sfd = False + + # ---- Activation function validation ---- + self._value_error_if( + self.act_func not in ["dswiglu", "dgeglu"], + f"act_func must be 'dswiglu' or 'dgeglu', got {self.act_func}", + ) + + # ---- Discrete-mode-specific validation ---- + if self.weight_mode == MoEWeightMode.DISCRETE: + self._value_error_if( + self.b_major not in ["k", "n"], + f"b_major must be 'k' or 'n', got {self.b_major}", + ) + self._value_error_if( + self._is_fp4x2(self.ab_dtype) and self.b_major != "k", + "b_major must be 'k' when ab_dtype is fp4", + ) + + # ---- MMA tile / cluster shape ---- + self._logger.debug("Checking MMA tile shape and cluster shape") + self._value_error_if( + not self.use_2cta_instrs and self.mma_tiler_mn[0] != 128, + f"MMA tiler M must be 128 when use_2cta_instrs=False, got {self.mma_tiler_mn[0]}", + ) + self._value_error_if( + self.use_2cta_instrs and self.mma_tiler_mn[0] != 256, + f"MMA tiler M must be 256 when use_2cta_instrs=True, got {self.mma_tiler_mn[0]}", + ) + self._value_error_if( + self.mma_tiler_mn[1] != 256, + f"MMA tiler N must be 256, got {self.mma_tiler_mn[1]}", + ) + self._value_error_if( + self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0, + f"cluster_shape_mn[0] must be divisible by 2 when use_2cta_instrs=True, got {self.cluster_shape_mn[0]}", + ) + self._value_error_if( + not ( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1] <= 16 + and self.cluster_shape_mn[0] > 0 + and self.cluster_shape_mn[1] > 0 + and self.cluster_shape_mn[0] <= 4 + and self.cluster_shape_mn[1] <= 4 + and is_power_of_2(self.cluster_shape_mn[0]) + and is_power_of_2(self.cluster_shape_mn[1]) + ), + f"Invalid cluster shape: expected values to be powers of 2 and product <= 16, got {self.cluster_shape_mn}", + ) + cluster_tiler_m = (self.cluster_shape_mn[0] // (2 if self.use_2cta_instrs else 1)) * self.mma_tiler_mn[0] + self._value_error_if( + cluster_tiler_m not in [128, 256], + f"Invalid cluster tiler shape: expected cluster_tiler_m in {{128, 256}}, got {cluster_tiler_m}", + ) + self._value_error_if( + self.m_aligned % self.mma_tiler_mn[0] != 0, + f"m_aligned must be divisible by mma_tiler_mn[0], got {self.m_aligned} % {self.mma_tiler_mn[0]} != 0", + ) + self._value_error_if( + self.m_aligned != BlockScaledMoEGroupedGemmDgluDbiasKernel.FIX_PAD_SIZE, + f"m_aligned must be {BlockScaledMoEGroupedGemmDgluDbiasKernel.FIX_PAD_SIZE} (FIX_PAD_SIZE), got {self.m_aligned}", + ) + + # ---- Tensor alignment ---- + self._logger.debug("Checking tensor alignment") + + def check_contiguous_16B_alignment(dtype, stride_order, tensor_shape): + is_mode0_major = stride_order == (0, 1, 2) + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // (_convert_to_cutlass_data_type(dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2).width) + return num_major_elements % num_contiguous_elements == 0 + + if self.weight_mode == MoEWeightMode.DENSE: + b_stride_order_for_check = self.b_desc.stride_order + b_shape_for_check = (n, k, l) + else: + b_stride_order_for_check = (0, 1, 2) if self.b_major == "n" else (1, 0, 2) + b_shape_for_check = (n, k, 1) + + self._value_error_if( + not ( + check_contiguous_16B_alignment(self.ab_dtype, self.a_desc.stride_order, (tensor_m, k, l)) + and check_contiguous_16B_alignment(self.ab_dtype, b_stride_order_for_check, b_shape_for_check) + and check_contiguous_16B_alignment(self.d_dtype, self.d_row_desc.stride_order, (tensor_m, n_out, 1)) + ), + "Invalid tensor alignment: tensors must be 16B aligned", + ) + + # ---- Expert count limit ---- + self._value_error_if( + self.expert_cnt > 1024, + f"expert_cnt must be <= 1024, got {self.expert_cnt}", + ) + + # ---- Disabled configurations ---- + self._not_implemented_error_if( + self.dbias_desc is None and self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype == torch.float32, + "Invalid configuration: fp4 ab_dtype, sf_vec_size 16, d_dtype float32 is not supported. " "Please use sf_vec_size 32 or d_dtype bf16 instead", + ) + + # ---- SM100+ check ---- + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + device = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(device) + compute_capability = major * 10 + minor + if compute_capability < 100: + raise RuntimeError(f"GroupedGemmDglu requires SM100+ compute capability, " f"but found SM{compute_capability} on device {device}") + + self._is_supported = True + self._logger.debug("check_support completed successfully") + return True + + # --------------------------------------------------------------------- # + # compile + # --------------------------------------------------------------------- # + + def compile(self) -> None: + """Compile the kernel.""" + self._logger.debug("Entering compile") + self._ensure_support_checked() + if self._compiled_kernel is not None: + self._logger.debug("Kernel already compiled; skipping recompilation") + return + if self.a_desc.shape[0] == 0: + self._logger.debug("sample valid_m is zero, skipping kernel compilation") + return + + # ---- Instantiate the unified kernel ---- + gemm_dglu = self._kernel( + sf_vec_size=self.sf_vec_size, + acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), + use_2cta_instrs=self.use_2cta_instrs, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + vectorized_f32=self.vector_f32, + discrete_col_sfd=self.discrete_col_sfd, + expert_cnt=self.expert_cnt, + weight_mode=self.weight_mode, + act_func=self.act_func, + use_dynamic_sched=self.use_dynamic_sched, + ) + + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) + max_active_clusters -= self.num_cluster_overlap_margin + self._value_error_if( + max_active_clusters <= 0, + "max_active_clusters must be > 0 after applying overlap margin; reduce CUDNNFE_CLUSTER_OVERLAP_MARGIN", + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + + # ---- Allocate workspace ---- + workspace_bytes = gemm_dglu.get_workspace_bytes() + self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device="cuda") + + if self.weight_mode == MoEWeightMode.DENSE: + self._compile_dense(gemm_dglu, max_active_clusters, fake_stream) + else: + self._compile_discrete(gemm_dglu, max_active_clusters, fake_stream) + + self._logger.debug("Kernel compiled successfully") + + # -- Dense compile path ------------------------------------------------- # + + def _compile_dense(self, gemm_dglu, max_active_clusters, fake_stream) -> None: + """Compile for dense (contiguous) weight mode.""" + use_full_dynamic = os.environ.get("CUDNN_FE_GROUPED_GEMM_DYNAMIC_MNKL", "1") != "0" + + fake_workspace_ptr = cute.runtime.nullptr( + dtype=cutlass.Uint8, + assumed_align=128, + ) + + if not use_full_dynamic: + valid_m = cute.sym_int(divisibility=256) + + a_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.a_desc.dtype, + shape=(valid_m, *self.a_desc.shape[1:]), + stride_order=self.a_desc.stride_order, + ) + b_cute_fake = self._make_fake_cute_tensor_from_desc(self.b_desc, assumed_align=16) + c_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.c_desc.dtype, + shape=(valid_m, *self.c_desc.shape[1:]), + stride_order=self.c_desc.stride_order, + ) + d_row_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.d_row_desc.dtype, + shape=(valid_m, *self.d_row_desc.shape[1:]), + stride_order=self.d_row_desc.stride_order, + ) + d_col_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.d_col_desc.dtype, + shape=(valid_m, *self.d_col_desc.shape[1:]), + stride_order=self.d_col_desc.stride_order, + ) + + tensor_m_128 = cute.sym_int() + stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfa_cute_fake = self._make_fake_cute_tensor( + dtype=self.sfa_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfa_desc.shape[4], 1), + stride=(16, 4, self.sfa_desc.stride[2], 1, 512, stride_tensor_m_128), + ) + + sfb_cute_fake = self._make_fake_cute_tensor_from_desc(self.sfb_desc, assumed_align=16) + + beta_cute_fake = self._make_fake_cute_tensor_from_desc(self.beta_desc, assumed_align=16) + prob_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.prob_desc.dtype, + shape=(valid_m, 1, 1), + stride_order=self.prob_desc.stride_order, + ) + dprob_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.dprob_desc.dtype, + shape=(valid_m, 1, 1), + stride_order=self.dprob_desc.stride_order, + ) + + sfd_row_fake = None + sfd_col_fake = None + if self.sfd_row_desc is not None: + stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) + sfd_row_fake = self._make_fake_cute_tensor( + dtype=self.sfd_row_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), + stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), + ) + if self.sfd_col_desc is not None: + rest_m = cute.sym_int(divisibility=1) + stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) + stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) + sfd_col_fake = self._make_fake_cute_tensor( + dtype=self.sfd_col_desc.dtype, + shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), + stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), + ) + else: + valid_m = cute.sym_int(divisibility=256) + n_sym = cute.sym_int() + n_out_sym = cute.sym_int() + k_sym = cute.sym_int() + l_sym = cute.sym_int() + + a_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.a_desc.dtype, + shape=(valid_m, k_sym, 1), + stride_order=self.a_desc.stride_order, + dynamic_mode=self.a_desc.stride_order[0], + divisibility=32 if self._is_fp4x2(self.ab_dtype) else 16, + ) + b_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.b_desc.dtype, + shape=(n_sym, k_sym, l_sym), + stride_order=self.b_desc.stride_order, + dynamic_mode=self.b_desc.stride_order[0], + divisibility=32 if self._is_fp4x2(self.ab_dtype) else 16, + ) + + c_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.c_desc.dtype, + shape=(valid_m, n_out_sym, 1), + stride_order=self.c_desc.stride_order, + dynamic_mode=self.c_desc.stride_order[0], + divisibility=8 if self._is_f16(self.c_desc.dtype) else 16, + ) + + d_row_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.d_row_desc.dtype, + shape=(valid_m, n_out_sym, 1), + stride_order=self.d_row_desc.stride_order, + dynamic_mode=self.d_row_desc.stride_order[0], + divisibility=8 if self._is_f16(self.d_row_desc.dtype) else 16, + ) + + d_col_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.d_col_desc.dtype, + shape=(valid_m, n_out_sym, 1), + stride_order=self.d_col_desc.stride_order, + dynamic_mode=self.d_col_desc.stride_order[0], + divisibility=8 if self._is_f16(self.d_col_desc.dtype) else 16, + ) + + tensor_m_128 = cute.sym_int() + rest_k = cute.sym_int() + stride_rest_k = cute.sym_int(divisibility=32 * 4 * 4) + stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfa_shape = list(self.sfa_desc.shape) + sfa_shape[2] = tensor_m_128 + sfa_shape[4] = rest_k + sfa_stride = list(self.sfa_desc.stride) + sfa_stride[2] = stride_rest_k + sfa_stride[5] = stride_tensor_m_128 + sfa_cute_fake = self._make_fake_cute_tensor( + dtype=self.sfa_desc.dtype, + shape=tuple(sfa_shape), + stride=tuple(sfa_stride), + ) + + tensor_n_128 = cute.sym_int() + stride_sfb_rest_k = cute.sym_int(divisibility=32 * 4 * 4) + stride_sfb_tensor_n_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfb_cute_fake = self._make_fake_cute_tensor( + dtype=self.sfb_desc.dtype, + shape=(32, 4, tensor_n_128, 4, rest_k, l_sym), + stride=(16, 4, stride_sfb_tensor_n_128, 1, 512, stride_sfb_rest_k), + ) + + beta_cute_fake = self._make_fake_cute_tensor_from_desc(self.beta_desc, assumed_align=16) + prob_cute_fake = self._make_fake_cute_tensor( + dtype=self.prob_desc.dtype, + shape=(valid_m, *self.prob_desc.shape[1:]), + stride=self.prob_desc.stride, + ) + dprob_cute_fake = self._make_fake_cute_tensor( + dtype=self.dprob_desc.dtype, + shape=(valid_m, *self.dprob_desc.shape[1:]), + stride=self.dprob_desc.stride, + ) + + sfd_row_fake = None + sfd_col_fake = None + if self.sfd_row_desc is not None: + rest_n_out = cute.sym_int() + stride_sfd_rest_n_out = cute.sym_int(divisibility=32 * 4 * 4) + stride_sfd_rest_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfd_row_fake = self._make_fake_cute_tensor( + dtype=self.sfd_row_desc.dtype, + shape=(32, 4, tensor_m_128, 4, rest_n_out, 1), + stride=(16, 4, stride_sfd_rest_n_out, 1, 512, stride_sfd_rest_tensor_m_128), + ) + if self.sfd_col_desc is not None: + tensor_n_out_128 = cute.sym_int() + rest_m_dyn = cute.sym_int() + stride_sfd_rest_m = cute.sym_int(divisibility=32 * 4 * 4) + stride_sfd_n_out = cute.sym_int(divisibility=32 * 4 * 4) + sfd_col_fake = self._make_fake_cute_tensor( + dtype=self.sfd_col_desc.dtype, + shape=(32, 4, tensor_n_out_128, 4, rest_m_dyn, 1), + stride=(16, 4, stride_sfd_rest_m, 1, 512, stride_sfd_n_out), + ) + + # Compile with keyword args (dense mode uses the unified __call__ positional order). + dbias_fake = self._make_fake_cute_tensor_from_desc(self.dbias_desc, assumed_align=16) + + _compiled_kernel = cute.compile( + gemm_dglu, + a=a_cute_fake, + b=b_cute_fake, + sfb=sfb_cute_fake, + n=cutlass.Int32(0), + k=cutlass.Int32(0), + b_stride_size=cutlass.Int64(0), + b_major_mode=OperandMajorMode.K, + workspace_ptr=fake_workspace_ptr, + c=c_cute_fake, + d=d_row_cute_fake, + d_col=d_col_cute_fake, + sfa=sfa_cute_fake, + sfd_row_tensor=sfd_row_fake, + sfd_col_tensor=sfd_col_fake, + amax_tensor=self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16), + norm_const_tensor=self._make_fake_cute_tensor_from_desc(self.norm_const_desc, assumed_align=16), + padded_offsets=self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16), + alpha=self._make_fake_cute_tensor_from_desc(self.alpha_desc, assumed_align=16), + beta=beta_cute_fake, + prob=prob_cute_fake, + dprob=dprob_cute_fake, + dbias_tensor=dbias_fake, + max_active_clusters=max_active_clusters, + stream=fake_stream, + epilogue_op=self.epilogue_op, + linear_offset=self.linear_offset, + geglu_alpha=self.geglu_alpha, + glu_clamp_max=self.glu_clamp_max, + glu_clamp_min=self.glu_clamp_min, + options="--enable-tvm-ffi", + ) + + # Cache workspace pointer for the tensor_api closure + cached_workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator + + def tensor_api( + a_tensor: torch.Tensor, + b_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_row_tensor: torch.Tensor, + d_col_tensor: Optional[torch.Tensor], + sfa_tensor: torch.Tensor, + sfb_tensor: torch.Tensor, + sfd_row_tensor: Optional[torch.Tensor], + sfd_col_tensor: Optional[torch.Tensor], + amax_tensor: Optional[torch.Tensor], + norm_const_tensor: Optional[torch.Tensor], + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + beta_tensor: torch.Tensor, + prob_tensor: torch.Tensor, + dprob_tensor: torch.Tensor, + dbias_tensor: Optional[torch.Tensor], + stream: cuda.CUstream, + ) -> None: + norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") + _compiled_kernel( + a_tensor, + b_tensor, + sfb_tensor, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + cached_workspace_ptr, + c_tensor, + d_row_tensor, + d_col_tensor, + sfa_tensor, + sfd_row_tensor, + sfd_col_tensor, + amax_tensor, + norm_const_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + dprob_tensor, + dbias_tensor, + stream, + ) + + self._compiled_kernel = tensor_api + + # -- Discrete compile path ---------------------------------------------- # + + def _compile_discrete(self, gemm_dglu, max_active_clusters, fake_stream) -> None: + """Compile for discrete (per-expert pointer) weight mode.""" + if len(self.b_shape) == 2: + n, k = self.b_shape + else: + n, k, _ = self.b_shape + + b_major_mode = OperandMajorMode.K if self.b_major == "k" else OperandMajorMode.MN + if self.b_major == "k": + b_stride_size = k + else: + b_stride_size = n + + ab_cutlass_dtype = _convert_to_cutlass_data_type(self.a_desc.dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2) + align = 32 if ab_cutlass_dtype.width == 4 else 16 + + valid_m = cute.sym_int(divisibility=256) + a_tensor = self._make_fake_cute_tensor( + dtype=self.a_desc.dtype, + shape=(valid_m, *self.a_desc.shape[1:]), + stride=(self.a_desc.stride[0], *self.a_desc.stride[1:]), + assumed_align=align, + ) + c_tensor = self._make_fake_cute_tensor( + dtype=self.c_desc.dtype, + shape=(valid_m, *self.c_desc.shape[1:]), + stride=(self.c_desc.stride[0], *self.c_desc.stride[1:]), + ) + d_row_tensor = self._make_fake_cute_compact_tensor( + dtype=self.d_row_desc.dtype, + shape=(valid_m, *self.d_row_desc.shape[1:]), + stride_order=self.d_row_desc.stride_order, + ) + d_col_tensor = self._make_fake_cute_compact_tensor( + dtype=self.d_col_desc.dtype, + shape=(valid_m, *self.d_col_desc.shape[1:]), + stride_order=self.d_col_desc.stride_order, + ) + + tensor_m_128 = cute.sym_int() + stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfa_shape = list(self.sfa_desc.shape) + sfa_shape[2] = tensor_m_128 + sfa_stride = list(self.sfa_desc.stride) + sfa_stride[5] = stride_tensor_m_128 + sfa_tensor = self._make_fake_cute_tensor( + dtype=self.sfa_desc.dtype, + shape=tuple(sfa_shape), + stride=tuple(sfa_stride), + assumed_align=16, + ) + sfd_row_tensor = None + if self.sfd_row_desc is not None: + stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) + sfd_row_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_row_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), + stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), + assumed_align=16, + ) + sfd_col_tensor = None + if self.sfd_col_desc is not None: + rest_m = cute.sym_int(divisibility=1) + stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) + stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) + sfd_col_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_col_desc.dtype, + shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), + stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), + assumed_align=16, + ) + amax_tensor = self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16) + norm_const_tensor_cute = self._make_fake_cute_tensor_from_desc(self.norm_const_desc, assumed_align=16) + padded_offsets_tensor = self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16) + alpha_tensor = self._make_fake_cute_tensor_from_desc(self.alpha_desc, assumed_align=16) + beta_tensor = self._make_fake_cute_tensor_from_desc(self.beta_desc, assumed_align=16) + prob_tensor = self._make_fake_cute_tensor( + dtype=self.prob_desc.dtype, + shape=(valid_m, *self.prob_desc.shape[1:]), + stride=self.prob_desc.stride, + assumed_align=16, + ) + dprob_tensor = self._make_fake_cute_tensor( + dtype=self.dprob_desc.dtype, + shape=(valid_m, *self.dprob_desc.shape[1:]), + stride=self.dprob_desc.stride, + assumed_align=16, + ) + dbias_tensor = self._make_fake_cute_tensor_from_desc(self.dbias_desc, assumed_align=16) + + # Compile-time pointer placeholders + b_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") + sfb_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") + b_ptrs_cute = from_dlpack(b_ptrs_placeholder, assumed_align=8).iterator + sfb_ptrs_cute = from_dlpack(sfb_ptrs_placeholder, assumed_align=8).iterator + + workspace_ptr_cute = from_dlpack(self._workspace, assumed_align=128).iterator + + self._logger.debug("Compiling discrete grouped GEMM dGLU kernel") + _compiled_kernel = cute.compile( + gemm_dglu, + a=a_tensor, + b=b_ptrs_cute, + sfb=sfb_ptrs_cute, + n=cutlass.Int32(n), + k=cutlass.Int32(k), + b_stride_size=cutlass.Int64(b_stride_size), + b_major_mode=b_major_mode, + workspace_ptr=workspace_ptr_cute, + c=c_tensor, + d=d_row_tensor, + d_col=d_col_tensor, + sfa=sfa_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + amax_tensor=amax_tensor, + norm_const_tensor=norm_const_tensor_cute, + padded_offsets=padded_offsets_tensor, + alpha=alpha_tensor, + beta=beta_tensor, + prob=prob_tensor, + dprob=dprob_tensor, + dbias_tensor=dbias_tensor, + max_active_clusters=max_active_clusters, + stream=fake_stream, + epilogue_op=self.epilogue_op, + linear_offset=self.linear_offset, + geglu_alpha=self.geglu_alpha, + glu_clamp_max=self.glu_clamp_max, + glu_clamp_min=self.glu_clamp_min, + options="--enable-tvm-ffi", + ) + + self._n = n + self._k = k + self._b_stride_size = b_stride_size + + # Cache constant values for execute() closure + cached_workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator + cached_n = cutlass.Int32(self._n) + cached_k = cutlass.Int32(self._k) + cached_b_stride = cutlass.Int64(self._b_stride_size) + + def tensor_api( + a_tensor: torch.Tensor, + b_ptrs_device: torch.Tensor, + sfb_ptrs_device: torch.Tensor, + c_tensor: torch.Tensor, + d_row_tensor: torch.Tensor, + d_col_tensor: Optional[torch.Tensor], + sfa_tensor: torch.Tensor, + sfd_row_tensor: Optional[torch.Tensor], + sfd_col_tensor: Optional[torch.Tensor], + amax_tensor: Optional[torch.Tensor], + norm_const_tensor: Optional[torch.Tensor], + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + beta_tensor: torch.Tensor, + prob_tensor: torch.Tensor, + dprob_tensor: torch.Tensor, + dbias_tensor: Optional[torch.Tensor], + stream: cuda.CUstream, + ) -> None: + norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") + b_ptrs_addr = int(b_ptrs_device.data_ptr()) + sfb_ptrs_addr = int(sfb_ptrs_device.data_ptr()) + + _compiled_kernel( + a_tensor, + b_ptrs_addr, + sfb_ptrs_addr, + cached_n, + cached_k, + cached_b_stride, + cached_workspace_ptr, + c_tensor, + d_row_tensor, + d_col_tensor, + sfa_tensor, + sfd_row_tensor, + sfd_col_tensor, + amax_tensor, + norm_const_tensor, + padded_offsets, + alpha_tensor, + beta_tensor, + prob_tensor, + dprob_tensor, + dbias_tensor, + stream, + ) + + self._compiled_kernel = tensor_api + + # --------------------------------------------------------------------- # + # execute + # --------------------------------------------------------------------- # + + def execute( + self, + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_row_tensor: torch.Tensor, + d_col_tensor: torch.Tensor, + sfa_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + beta_tensor: torch.Tensor, + prob_tensor: torch.Tensor, + dprob_tensor: torch.Tensor, + # Dense mode: + b_tensor: Optional[torch.Tensor] = None, + sfb_tensor: Optional[torch.Tensor] = None, + dbias_tensor: Optional[torch.Tensor] = None, + # Discrete mode: + b_ptrs: Optional[torch.Tensor] = None, + sfb_ptrs: Optional[torch.Tensor] = None, + # Optional: + sfd_row_tensor: Optional[torch.Tensor] = None, + sfd_col_tensor: Optional[torch.Tensor] = None, + amax_tensor: Optional[torch.Tensor] = None, + norm_const_tensor: Optional[torch.Tensor] = None, + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + """Execute the compiled kernel. + + For dense mode, supply ``b_tensor`` and ``sfb_tensor``. + For discrete mode, supply ``b_ptrs`` and ``sfb_ptrs``. + + :param a_tensor: Input A tensor (gradient input) + :param c_tensor: Forward activations input + :param d_row_tensor: Output D row tensor + :param d_col_tensor: Output D column tensor + :param sfa_tensor: Scale factor A + :param padded_offsets: End offset per expert after padding + :param alpha_tensor: Per-group alpha scaling factors + :param beta_tensor: Per-group beta scaling factors + :param prob_tensor: Per-row probability (from forward) + :param dprob_tensor: Gradient of probability (output, must be zero-initialized) + :param b_tensor: (Dense) Input B tensor (weights) + :param sfb_tensor: (Dense) Scale factor B + :param dbias_tensor: Optional dbias output tensor. + :param b_ptrs: (Discrete) 1-D int64 device tensor of per-expert B data pointers + :param sfb_ptrs: (Discrete) 1-D int64 device tensor of per-expert SFB data pointers + :param sfd_row_tensor: Optional row scale factor D + :param sfd_col_tensor: Optional column scale factor D + :param amax_tensor: Optional amax tensor + :param norm_const_tensor: Optional normalization constant + :param current_stream: CUDA stream + """ + self._logger.debug("Entering execute") + current_stream = self._get_default_stream(current_stream) + + if a_tensor.shape[0] == 0: + self._logger.debug("execute: valid_m is zero, skipping kernel execution") + return + self._runtime_error_if( + self._compiled_kernel is None, + "Kernel not compiled; call compile() first", + ) + + self._logger.debug("Executing grouped GEMM dGLU kernel") + if self._has_dbias: + self._value_error_if( + dbias_tensor is None, + "dbias_tensor is required when GroupedGemmDgluSm100 is configured with sample_dbias", + ) + + if self.weight_mode == MoEWeightMode.DENSE: + self._compiled_kernel( + a_tensor=a_tensor, + b_tensor=b_tensor, + c_tensor=c_tensor, + d_row_tensor=d_row_tensor, + d_col_tensor=d_col_tensor, + sfa_tensor=sfa_tensor, + sfb_tensor=sfb_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + amax_tensor=amax_tensor, + norm_const_tensor=norm_const_tensor, + padded_offsets=padded_offsets, + alpha_tensor=alpha_tensor, + beta_tensor=beta_tensor, + prob_tensor=prob_tensor, + dprob_tensor=dprob_tensor, + dbias_tensor=dbias_tensor, + stream=current_stream, + ) + else: + self._record_pointer_stream(b_ptrs, current_stream) + self._record_pointer_stream(sfb_ptrs, current_stream) + self._compiled_kernel( + a_tensor=a_tensor, + b_ptrs_device=b_ptrs, + sfb_ptrs_device=sfb_ptrs, + c_tensor=c_tensor, + d_row_tensor=d_row_tensor, + d_col_tensor=d_col_tensor, + sfa_tensor=sfa_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + amax_tensor=amax_tensor, + norm_const_tensor=norm_const_tensor, + padded_offsets=padded_offsets, + alpha_tensor=alpha_tensor, + beta_tensor=beta_tensor, + prob_tensor=prob_tensor, + dprob_tensor=dprob_tensor, + dbias_tensor=dbias_tensor, + stream=current_stream, + ) + + self._logger.debug("Execute completed") + + +__all__ = ["GroupedGemmDgluBlockScaledAPI"] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dglu/api.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/api.py index 1d026812d..82c32f677 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_dglu/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/api.py @@ -43,89 +43,142 @@ at execution time. """ -from .moe_blockscaled_grouped_gemm_dglu_dbias import BlockScaledMoEGroupedGemmDgluDbiasKernel +from dataclasses import dataclass, replace + +from ..grouped_gemm_utils import ( + GroupedGemmBackend, + _torch_stream_context, + backend_cache_key, + select_grouped_gemm_backend, +) from ..moe_utils import MoEWeightMode from cuda.bindings import driver as cuda import logging import os import torch -from typing import Tuple, Optional +from typing import Any, Tuple, Optional, overload -import cutlass -import cutlass.cute as cute -from cutlass.cute.nvgpu import OperandMajorMode -from cutlass.cute.runtime import from_dlpack, make_fake_stream +from cudnn.api_base import APIBase, TupleDict, ceil_div -from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +_BLOCK_SCALED_DTYPE_PAIRS = { + (dtype, dtype) + for dtype in ( + torch.float4_e2m1fn_x2, + torch.uint8, + torch.float8_e5m2, + torch.float8_e4m3fn, + ) +} + + +from ._bf16_api import GroupedGemmDgluBf16API +from ._blockscaled_api import GroupedGemmDgluBlockScaledAPI + + +@dataclass(frozen=True) +class DgluCall: + """Immutable normalized input for dGLU dispatch, allocation, and caching.""" + + a_tensor: torch.Tensor + c_tensor: torch.Tensor + sfa_tensor: Optional[torch.Tensor] + padded_offsets: torch.Tensor + alpha_tensor: torch.Tensor + beta_tensor: torch.Tensor + prob_tensor: torch.Tensor + dprob_tensor: torch.Tensor + b_tensor: Optional[torch.Tensor] = None + sfb_tensor: Optional[torch.Tensor] = None + generate_dbias: bool = False + b_ptrs: Optional[torch.Tensor] = None + sfb_ptrs: Optional[torch.Tensor] = None + n: Optional[int] = None + b_dtype: Optional[torch.dtype] = None + b_major: str = "k" + norm_const_tensor: Optional[torch.Tensor] = None + acc_dtype: torch.dtype = torch.float32 + d_dtype: torch.dtype = torch.bfloat16 + cd_major: str = "n" + mma_tiler_mn: Tuple[int, int] = (256, 256) + cluster_shape_mn: Optional[Tuple[int, int]] = None + sf_vec_size: int = 16 + vector_f32: bool = False + m_aligned: int = 256 + discrete_col_sfd: bool = False + act_func: str = "dswiglu" + linear_offset: Optional[float] = None + geglu_alpha: float = 1.702 + glu_clamp_max: float = 7.0 + glu_clamp_min: float = -7.0 + epilogue_op: Optional[str] = None + use_dynamic_sched: bool = False + current_stream: Optional[cuda.CUstream] = None + weight_mode: Optional[MoEWeightMode] = None + b_shape: Optional[Tuple[int, ...]] = None + num_experts: Optional[int] = None class GroupedGemmDgluSm100(APIBase): - """Unified API for grouped GEMM dGLU backward operation on SM100+ GPUs. - - This kernel performs block-scaled grouped GEMM with dGLU activation - gradient (dSwiGLU or dGeGLU), designed for MoE workloads. It supports - both dense (contiguous) and discrete (per-expert pointer) weight layouts - through the ``BlockScaledMoEGroupedGemmDgluDbiasKernel``. - - Weight mode is auto-detected from the constructor arguments: - - - **Dense**: provide ``sample_b`` and ``sample_sfb``. - - **Discrete**: provide ``num_experts``, ``b_shape``, and ``b_dtype``. - - Example:: - - # Dense mode - api = GroupedGemmDgluSm100( - sample_a=a, sample_c=c, - sample_d_row=d_row, sample_d_col=d_col, - sample_sfa=sfa, sample_padded_offsets=offsets, - sample_alpha=alpha, sample_beta=beta, - sample_prob=prob, sample_dprob=dprob, - sample_b=b, sample_sfb=sfb, - ) - - # Discrete mode - api = GroupedGemmDgluSm100( - sample_a=a, sample_c=c, - sample_d_row=d_row, sample_d_col=d_col, - sample_sfa=sfa, sample_padded_offsets=offsets, - sample_alpha=alpha, sample_beta=beta, - sample_prob=prob, sample_dprob=dprob, - num_experts=8, b_shape=(n, k), b_dtype=torch.uint8, - ) + """Stable public facade that selects the dGLU backend during support checking.""" - api.check_support() - api.compile() - api.execute(...) - """ + # BF16 implementation + @overload + def __init__( + self, + sample_a: torch.Tensor, + sample_c: torch.Tensor, + sample_d_row: torch.Tensor, + sample_d_col: None, + sample_sfa: None, + sample_padded_offsets: torch.Tensor, + sample_alpha: torch.Tensor, + sample_beta: torch.Tensor, + sample_prob: torch.Tensor, + sample_dprob: torch.Tensor, + *args: Any, + **kwargs: Any, + ) -> None: ... + # Block-scaled implementation + @overload def __init__( self, sample_a: torch.Tensor, sample_c: torch.Tensor, sample_d_row: torch.Tensor, - sample_d_col: torch.Tensor, + sample_d_col: Optional[torch.Tensor], sample_sfa: torch.Tensor, sample_padded_offsets: torch.Tensor, sample_alpha: torch.Tensor, sample_beta: torch.Tensor, sample_prob: torch.Tensor, sample_dprob: torch.Tensor, - # Dense mode (contiguous) -- provide these. sample_dbias is optional: + *args: Any, + **kwargs: Any, + ) -> None: ... + + def __init__( + self, + sample_a: torch.Tensor, + sample_c: torch.Tensor, + sample_d_row: torch.Tensor, + sample_d_col: Optional[torch.Tensor], + sample_sfa: Optional[torch.Tensor], + sample_padded_offsets: torch.Tensor, + sample_alpha: torch.Tensor, + sample_beta: torch.Tensor, + sample_prob: torch.Tensor, + sample_dprob: torch.Tensor, sample_b: Optional[torch.Tensor] = None, sample_sfb: Optional[torch.Tensor] = None, sample_dbias: Optional[torch.Tensor] = None, - # Discrete mode -- provide these instead: num_experts: Optional[int] = None, b_shape: Optional[Tuple[int, ...]] = None, b_dtype: Optional[torch.dtype] = None, - # Optional quantization output arguments sample_sfd_row: Optional[torch.Tensor] = None, sample_sfd_col: Optional[torch.Tensor] = None, sample_amax: Optional[torch.Tensor] = None, sample_norm_const: Optional[torch.Tensor] = None, - # Configuration acc_dtype: torch.dtype = torch.float32, mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, @@ -141,1140 +194,212 @@ def __init__( geglu_alpha: float = 1.702, glu_clamp_max: float = 7.0, glu_clamp_min: float = -7.0, - ): - """Initialize the GroupedGemmDgluSm100 API. - - :param sample_a: Sample A tensor (valid_m, k, 1) - :param sample_c: Sample C tensor -- forward activations (valid_m, n, 1) - :param sample_d_row: Sample D row output tensor (valid_m, n*2, 1) - :param sample_d_col: Sample D col output tensor (valid_m, n*2, 1) - :param sample_sfa: Sample scale factor A tensor - :param sample_padded_offsets: End offset for each expert after padding - :param sample_alpha: Per-group alpha scaling factors - :param sample_beta: Per-group beta scaling factors - :param sample_prob: Per-row probability tensor (valid_m, 1, 1) - :param sample_dprob: Gradient of probability tensor (valid_m, 1, 1), must be zero-initialized - :param sample_b: (Dense) Sample B tensor (n, k, l) - :param sample_sfb: (Dense) Sample scale factor B tensor - :param sample_dbias: Optional dbias output tensor (expert_cnt, 2*n, 1) - :param num_experts: (Discrete) Number of experts - :param b_shape: (Discrete) Shape of a single expert B tensor, e.g. (n, k) - :param b_dtype: (Discrete) Data type of B tensors - :param sample_sfd_row: Optional row scale factor for D - :param sample_sfd_col: Optional column scale factor for D - :param sample_amax: Optional amax tensor for quantization, shape (expert_cnt, 2, 1) - :param sample_norm_const: Optional normalization constant - :param acc_dtype: Accumulator data type - :param mma_tiler_mn: MMA tiler shape (M, N) - :param cluster_shape_mn: Cluster shape (M, N) - :param sf_vec_size: Scale factor vector size - :param vector_f32: Use vectorized f32 operations - :param m_aligned: Alignment for group M dimension - :param discrete_col_sfd: Generate discrete col-major scale factor tensor - :param act_func: Activation function, one of "dswiglu" or "dgeglu" - :param b_major: Major dimension for B tensor, one of "k" or "n" - :param epilogue_op: Optional epilogue operation. Valid: None, "none", "identity", "relu", "srelu" - :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing - :param linear_offset: Compile-time linear offset for dGeGLU. When None, - defaults to 1.0 for dGeGLU and 0.0 for dSwiGLU. - Ignored when ``act_func == "dswiglu"``. - :param geglu_alpha: Compile-time dGeGLU GeGLU alpha. Ignored when - ``act_func == "dswiglu"``. - :param glu_clamp_max: Compile-time dGeGLU upper clamp. Ignored when - ``act_func == "dswiglu"``. - :param glu_clamp_min: Compile-time dGeGLU lower clamp. Ignored when - ``act_func == "dswiglu"``. - """ + ) -> None: super().__init__() - - self._warn_experimental_api() - self._logger.debug("Entering __init__") - - # ---- Weight mode auto-detection ---- - if sample_b is not None and num_experts is None: - self.weight_mode = MoEWeightMode.DENSE - if sample_sfb is None: - raise ValueError("sample_sfb is required when sample_b is provided (dense mode)") - elif num_experts is not None and sample_b is None: - self.weight_mode = MoEWeightMode.DISCRETE - if b_shape is None or b_dtype is None: - raise ValueError("b_shape and b_dtype are required in discrete mode") - else: - raise ValueError("Provide either (sample_b, sample_sfb) for dense mode " "or (num_experts, b_shape, b_dtype) for discrete mode, but not both.") - - # ---- Common tensor descriptors ---- - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") - self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.d_row_desc = self._make_tensor_desc(sample_d_row, name="sample_d_row") - self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") - self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") - self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") - self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") - self.beta_desc = self._make_tensor_desc(sample_beta, name="sample_beta") - self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") - self.dprob_desc = self._make_tensor_desc(sample_dprob, name="sample_dprob") - self.dbias_desc = self._make_tensor_desc(sample_dbias, name="sample_dbias") - - self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row") - self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col") - self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") - self.norm_const_desc = self._unpad_tensor_to_ndim( - self._make_tensor_desc(sample_norm_const, name="sample_norm_const"), - 1, - "norm_const", - ) - - # ---- Mode-specific state ---- - if self.weight_mode == MoEWeightMode.DENSE: - self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") - self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb") - self.expert_cnt = self.padded_offsets_desc.shape[0] - else: - self._value_error_if(num_experts == 0, "num_experts must be > 0") - self.expert_cnt = num_experts - self.b_shape = b_shape - self.b_dtype = b_dtype - self.b_major = b_major - self._value_error_if( - self.padded_offsets_desc.shape[0] != self.expert_cnt, - f"padded_offsets length ({self.padded_offsets_desc.shape[0]}) " f"must equal num_experts ({self.expert_cnt})", - ) - - # ---- Configuration ---- - self.acc_dtype = acc_dtype - self.mma_tiler_mn = mma_tiler_mn - self.use_2cta_instrs = mma_tiler_mn[0] == 256 - if cluster_shape_mn is None: - self.cluster_shape_mn = (2, 1) if self.use_2cta_instrs else (1, 1) - else: - self.cluster_shape_mn = cluster_shape_mn - self.sf_vec_size = sf_vec_size - self.vector_f32 = vector_f32 - self.m_aligned = m_aligned - self.discrete_col_sfd = discrete_col_sfd - self.act_func = act_func - if self.weight_mode == MoEWeightMode.DENSE: - self.b_major = b_major # stored for both modes - - # Epilogue operation - if epilogue_op in [None, "none", "identity"]: - self.epilogue_op = lambda x: x - elif epilogue_op == "relu": - self.epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) - elif epilogue_op == "srelu": - self.epilogue_op = lambda x: cute.where(x > 0, x, cute.full_like(x, 0)) ** 2 - else: - raise ValueError(f"Invalid epilogue operation: {epilogue_op}. " f"Valid values: None, 'none', 'identity', 'relu', 'srelu'") - - self.use_dynamic_sched = use_dynamic_sched - if linear_offset is None: - self.linear_offset = 1.0 if self.act_func == "dgeglu" else 0.0 - else: - self.linear_offset = float(linear_offset) - self.geglu_alpha = geglu_alpha - self.glu_clamp_max = glu_clamp_max - self.glu_clamp_min = glu_clamp_min - - self._interpret_uint8_as_fp4x2 = True - self._has_dbias = self.dbias_desc is not None - self._kernel = BlockScaledMoEGroupedGemmDgluDbiasKernel - - self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) - self._logger.debug(f"setting num_cluster_overlap_margin: {self.num_cluster_overlap_margin}") - - self._workspace = None - - self._logger.debug("__init__ completed") - - # --------------------------------------------------------------------- # - # check_support - # --------------------------------------------------------------------- # + self._pending_init_kwargs = dict(locals()) + self._pending_init_kwargs.pop("self") + self._pending_init_kwargs.pop("__class__", None) + self._implementation = None def check_support(self) -> bool: - """Check if the kernel configuration is supported. - - :return: True if supported, raises exception otherwise - """ - self._logger.debug("Entering check_support") - - # ---- SFD group validation ---- - all_none = all(x is None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) - all_provided = all(x is not None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) - self._value_error_if( - not (all_none or all_provided), - "sfd_row_desc, sfd_col_desc, and norm_const_desc must be all None or all not None", - ) - self._user_requested_sfd = all_provided - - # ---- Shapes and strides ---- - self._logger.debug("Checking tensor shapes and strides") - tensor_m, k, _one = self._tensor_shape(self.a_desc, name="sample_a") - - if self.weight_mode == MoEWeightMode.DENSE: - n, _, l = self._tensor_shape(self.b_desc, name="sample_b") - else: - # Discrete: extract n, k from b_shape - if len(self.b_shape) == 2: - n, b_k = self.b_shape - else: - n, b_k, _ = self.b_shape - self._value_error_if(b_k != k, f"B K dimension ({b_k}) must match A K dimension ({k})") - l = self.expert_cnt # for shape checks that use l - - n_out = 2 * n - - self._value_error_if( - n % 32 != 0, - f"N must be divisible by 32 for dGLU (32-column input/gate interleaving), got N={n}", - ) - - self._check_tensor_shape(self.a_desc, (tensor_m, k, 1), "A") - if self.weight_mode == MoEWeightMode.DENSE: - self._check_tensor_shape(self.b_desc, (n, k, l), "B") - self._check_tensor_shape(self.c_desc, (tensor_m, n_out, 1), "C") - self._check_tensor_shape(self.d_row_desc, (tensor_m, n_out, 1), "D_row") - self._check_tensor_shape(self.d_col_desc, (tensor_m, n_out, 1), "D_col") - - rest_k = ceil_div(ceil_div(k, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfa_desc, (32, 4, ceil_div(tensor_m, 128), 4, rest_k, 1), "SFA") - if self.weight_mode == MoEWeightMode.DENSE: - self._check_tensor_shape(self.sfb_desc, (32, 4, ceil_div(n, 128), 4, rest_k, l), "SFB") - - # SFD uses n_out dimension since D has n_out columns - rest_n_out = ceil_div(ceil_div(n_out, self.sf_vec_size), 4) - self._check_tensor_shape( - self.sfd_row_desc, - (32, 4, ceil_div(tensor_m, 128), 4, rest_n_out, 1), - "SFD_row", - ) - rest_m = ceil_div(ceil_div(tensor_m, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfd_col_desc, (32, 4, ceil_div(n_out, 128), 4, rest_m, 1), "SFD_col") - - self._check_tensor_shape(self.alpha_desc, (self.expert_cnt,), "alpha") - self._check_tensor_shape(self.beta_desc, (self.expert_cnt,), "beta") - self._check_tensor_shape(self.prob_desc, (tensor_m, 1, 1), "prob") - self._check_tensor_shape(self.dprob_desc, (tensor_m, 1, 1), "dprob") - self._check_tensor_shape(self.dbias_desc, (self.expert_cnt, n_out, 1), "dbias") - self._check_tensor_shape(self.amax_desc, (self.expert_cnt, 2, 1), "amax") - self._check_tensor_shape(self.norm_const_desc, (1,), "norm_const") - self._check_tensor_shape(self.padded_offsets_desc, (self.expert_cnt,), "padded_offsets") - - # Strides - _ = self._check_tensor_stride( - self.a_desc, - stride=[(k, 1, tensor_m * k)], - extra_error_msg="A must have k-major layout", - ) - if self.weight_mode == MoEWeightMode.DENSE: - if self._is_fp8(self.a_desc): - _ = self._check_tensor_stride( - self.b_desc, - stride=[(k, 1, n * k), (1, n, n * k)], - extra_error_msg="For fp8 ab_dtype, B must have k- or n-major layout", + if self._implementation is None: + kwargs = self._pending_init_kwargs + defining_b_dtype = kwargs["sample_b"].dtype if kwargs["sample_b"] is not None else kwargs["b_dtype"] + backend = select_grouped_gemm_backend( + operation="grouped_gemm_dglu_sm100", + a_dtype=kwargs["sample_a"].dtype, + b_dtype=defining_b_dtype, + scale_controls=( + ("sample_sfa", kwargs["sample_sfa"]), + ("sample_sfb", kwargs["sample_sfb"]), + ("sample_d_col", kwargs["sample_d_col"]), + ("sample_sfd_row", kwargs["sample_sfd_row"]), + ("sample_sfd_col", kwargs["sample_sfd_col"]), + ("sample_amax", kwargs["sample_amax"]), + ("sample_norm_const", kwargs["sample_norm_const"]), + ("sf_vec_size", kwargs["sf_vec_size"] if kwargs["sf_vec_size"] != 16 else None), + ("discrete_col_sfd", kwargs["discrete_col_sfd"] if kwargs["discrete_col_sfd"] else None), + ("geglu_alpha", kwargs["geglu_alpha"] if kwargs["geglu_alpha"] != 1.702 else None), + ("glu_clamp_max", kwargs["glu_clamp_max"] if kwargs["glu_clamp_max"] != 7.0 else None), + ("glu_clamp_min", kwargs["glu_clamp_min"] if kwargs["glu_clamp_min"] != -7.0 else None), + ("epilogue_op", kwargs["epilogue_op"] if kwargs["epilogue_op"] not in (None, "none", "identity") else None), + ), + block_scaled_dtype_pairs=_BLOCK_SCALED_DTYPE_PAIRS, + ) + self.backend = backend + self.linear_offset = kwargs["linear_offset"] + if self.linear_offset is None: + self.linear_offset = 1.0 if kwargs["act_func"] == "dgeglu" else 0.0 + if backend is GroupedGemmBackend.BF16: + self._implementation = GroupedGemmDgluBf16API( + sample_a=kwargs["sample_a"], + sample_c=kwargs["sample_c"], + sample_d_row=kwargs["sample_d_row"], + sample_padded_offsets=kwargs["sample_padded_offsets"], + sample_alpha=kwargs["sample_alpha"], + sample_beta=kwargs["sample_beta"], + sample_prob=kwargs["sample_prob"], + sample_dprob=kwargs["sample_dprob"], + sample_b=kwargs["sample_b"], + sample_dbias=kwargs["sample_dbias"], + num_experts=kwargs["num_experts"], + b_shape=kwargs["b_shape"], + b_dtype=kwargs["b_dtype"], + acc_dtype=kwargs["acc_dtype"], + mma_tiler_mn=kwargs["mma_tiler_mn"], + cluster_shape_mn=kwargs["cluster_shape_mn"], + vector_f32=kwargs["vector_f32"], + m_aligned=kwargs["m_aligned"], + act_func=kwargs["act_func"], + b_major=kwargs["b_major"], + use_dynamic_sched=kwargs["use_dynamic_sched"], ) else: - _ = self._check_tensor_stride( - self.b_desc, - stride=[(k, 1, n * k)], - extra_error_msg="For fp4 ab_dtype, B must have k-major layout", - ) - _ = self._check_tensor_stride( - self.c_desc, - stride=[(n_out, 1, tensor_m * n_out)], - extra_error_msg="C must have n-major layout", - ) - _ = self._check_tensor_stride( - self.d_row_desc, - stride=[(n_out, 1, tensor_m * n_out)], - extra_error_msg="D_row must have n-major layout", - ) - _ = self._check_tensor_stride( - self.d_col_desc, - stride=[(n_out, 1, tensor_m * n_out)], - extra_error_msg="D_col must have n-major layout", - ) - - # ---- Data types ---- - self._logger.debug("Checking data types") - self.ab_dtype = self._check_dtype( - self.a_desc, - dtype=[ - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, - ], - name="A/B", - ) - if self.weight_mode == MoEWeightMode.DENSE: - self._check_dtype( - self.b_desc, - dtype=self.ab_dtype, - name="B", - extra_error_msg="B must have the same dtype as A", - ) - else: - self._value_error_if( - self.b_dtype != self.ab_dtype, - f"b_dtype ({self.b_dtype}) must match A dtype ({self.ab_dtype})", - ) - - self.sf_dtype = self._check_dtype( - self.sfa_desc, - dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], - name="SFA/SFB/SFD", - ) - if self.weight_mode == MoEWeightMode.DENSE: - self._check_dtype( - self.sfb_desc, - dtype=self.sf_dtype, - name="SFB", - extra_error_msg="SFB must have the same dtype as SFA", - ) - self._check_dtype( - self.sfd_row_desc, - dtype=self.sf_dtype, - name="SFD_row", - extra_error_msg="SFD_row must have the same dtype as SFA", - ) - self._check_dtype( - self.sfd_col_desc, - dtype=self.sf_dtype, - name="SFD_col", - extra_error_msg="SFD_col must have the same dtype as SFA", - ) - - self._value_error_if( - self.sf_vec_size not in [16, 32], - f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}", - ) - self._value_error_if( - self.sf_dtype in [torch.float8_e4m3fn] and self.sf_vec_size == 32, - f"sf_dtype {self.sf_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", - ) - self._value_error_if( - self._is_fp8(self.ab_dtype) and self.sf_vec_size == 16, - f"ab_dtype {self.ab_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", - ) - - self._check_dtype( - self.acc_dtype, - dtype=torch.float32, - name="Accumulator", - extra_error_msg="Accumulator must be float32", - ) - self._check_dtype( - self.prob_desc, - dtype=torch.float32, - name="Prob", - extra_error_msg="Prob must be float32", - ) - self._check_dtype( - self.dprob_desc, - dtype=torch.float32, - name="Dprob", - extra_error_msg="Dprob must be float32", - ) - self._check_dtype( - self.dbias_desc, - dtype=torch.bfloat16, - name="Dbias", - extra_error_msg="dbias must be bfloat16", - ) - self.c_dtype = self._check_dtype( - self.c_desc, - dtype=[torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2], - name="C", - ) - if self._is_fp8(self.c_dtype) and self.vector_f32: - raise ValueError("Invalid configuration: fp8 c_dtype and vector_f32 is not supported. " "Please use vector_f32=False or c_dtype=bfloat16 instead") - - if self._is_fp4x2(self.ab_dtype): - self.d_dtype = self._check_dtype( - self.d_row_desc, - dtype=[torch.float16, torch.bfloat16, torch.float32], - name="D_row", - extra_error_msg="D_row must be fp16, bf16, or float32 when ab_dtype is fp4", - ) - elif self._is_fp8(self.ab_dtype): - self.d_dtype = self._check_dtype( - self.d_row_desc, - dtype=[ - torch.float8_e4m3fn, - torch.float8_e5m2, - ], - name="D_row", - extra_error_msg="D_row must be fp8 dtype when ab_dtype is fp8", - ) - else: - raise NotImplementedError(f"Invalid ab_dtype: {self.ab_dtype}, expected fp4 or fp8") - self._check_dtype( - self.d_col_desc, - dtype=self.d_dtype, - name="D_col", - extra_error_msg="D_col must have the same dtype as D_row", - ) - - # ---- SFD generation logic ---- - kernel_generate_sfd = self._is_fp8(self.ab_dtype) and self.sf_dtype == torch.float8_e8m0fnu and self._is_fp8(self.d_dtype) - self._value_error_if( - kernel_generate_sfd and not self._user_requested_sfd, - "sfd_row, sfd_col, and norm_const are required for FP8 input/FP8 output with sf_dtype=torch.float8_e8m0fnu", - ) - if not kernel_generate_sfd and self._user_requested_sfd: - self._logger.warning( - "sfd_row/sfd_col/norm_const were provided, but this configuration does not generate SFD outputs; " "the tensors will be ignored by the kernel", - ) - self.generate_sfd = kernel_generate_sfd - if self.discrete_col_sfd and not self.generate_sfd: - self._logger.warning("discrete_col_sfd is True but generate_sfd is False, discrete_col_sfd will be ignored") - self.discrete_col_sfd = False - - # ---- Activation function validation ---- - self._value_error_if( - self.act_func not in ["dswiglu", "dgeglu"], - f"act_func must be 'dswiglu' or 'dgeglu', got {self.act_func}", - ) - - # ---- Discrete-mode-specific validation ---- - if self.weight_mode == MoEWeightMode.DISCRETE: - self._value_error_if( - self.b_major not in ["k", "n"], - f"b_major must be 'k' or 'n', got {self.b_major}", - ) - self._value_error_if( - self._is_fp4x2(self.ab_dtype) and self.b_major != "k", - "b_major must be 'k' when ab_dtype is fp4", - ) - - # ---- MMA tile / cluster shape ---- - self._logger.debug("Checking MMA tile shape and cluster shape") - self._value_error_if( - not self.use_2cta_instrs and self.mma_tiler_mn[0] != 128, - f"MMA tiler M must be 128 when use_2cta_instrs=False, got {self.mma_tiler_mn[0]}", - ) - self._value_error_if( - self.use_2cta_instrs and self.mma_tiler_mn[0] != 256, - f"MMA tiler M must be 256 when use_2cta_instrs=True, got {self.mma_tiler_mn[0]}", - ) - self._value_error_if( - self.mma_tiler_mn[1] != 256, - f"MMA tiler N must be 256, got {self.mma_tiler_mn[1]}", - ) - self._value_error_if( - self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0, - f"cluster_shape_mn[0] must be divisible by 2 when use_2cta_instrs=True, got {self.cluster_shape_mn[0]}", - ) - self._value_error_if( - not ( - self.cluster_shape_mn[0] * self.cluster_shape_mn[1] <= 16 - and self.cluster_shape_mn[0] > 0 - and self.cluster_shape_mn[1] > 0 - and self.cluster_shape_mn[0] <= 4 - and self.cluster_shape_mn[1] <= 4 - and is_power_of_2(self.cluster_shape_mn[0]) - and is_power_of_2(self.cluster_shape_mn[1]) - ), - f"Invalid cluster shape: expected values to be powers of 2 and product <= 16, got {self.cluster_shape_mn}", - ) - cluster_tiler_m = (self.cluster_shape_mn[0] // (2 if self.use_2cta_instrs else 1)) * self.mma_tiler_mn[0] - self._value_error_if( - cluster_tiler_m not in [128, 256], - f"Invalid cluster tiler shape: expected cluster_tiler_m in {{128, 256}}, got {cluster_tiler_m}", - ) - self._value_error_if( - self.m_aligned % self.mma_tiler_mn[0] != 0, - f"m_aligned must be divisible by mma_tiler_mn[0], got {self.m_aligned} % {self.mma_tiler_mn[0]} != 0", - ) - self._value_error_if( - self.m_aligned != BlockScaledMoEGroupedGemmDgluDbiasKernel.FIX_PAD_SIZE, - f"m_aligned must be {BlockScaledMoEGroupedGemmDgluDbiasKernel.FIX_PAD_SIZE} (FIX_PAD_SIZE), got {self.m_aligned}", - ) - - # ---- Tensor alignment ---- - self._logger.debug("Checking tensor alignment") - - def check_contiguous_16B_alignment(dtype, stride_order, tensor_shape): - is_mode0_major = stride_order == (0, 1, 2) - major_mode_idx = 0 if is_mode0_major else 1 - num_major_elements = tensor_shape[major_mode_idx] - num_contiguous_elements = 16 * 8 // (_convert_to_cutlass_data_type(dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2).width) - return num_major_elements % num_contiguous_elements == 0 - - if self.weight_mode == MoEWeightMode.DENSE: - b_stride_order_for_check = self.b_desc.stride_order - b_shape_for_check = (n, k, l) - else: - b_stride_order_for_check = (0, 1, 2) if self.b_major == "n" else (1, 0, 2) - b_shape_for_check = (n, k, 1) - - self._value_error_if( - not ( - check_contiguous_16B_alignment(self.ab_dtype, self.a_desc.stride_order, (tensor_m, k, l)) - and check_contiguous_16B_alignment(self.ab_dtype, b_stride_order_for_check, b_shape_for_check) - and check_contiguous_16B_alignment(self.d_dtype, self.d_row_desc.stride_order, (tensor_m, n_out, 1)) - ), - "Invalid tensor alignment: tensors must be 16B aligned", - ) - - # ---- Expert count limit ---- - self._value_error_if( - self.expert_cnt > 1024, - f"expert_cnt must be <= 1024, got {self.expert_cnt}", - ) - - # ---- Disabled configurations ---- - self._not_implemented_error_if( - self.dbias_desc is None and self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype == torch.float32, - "Invalid configuration: fp4 ab_dtype, sf_vec_size 16, d_dtype float32 is not supported. " "Please use sf_vec_size 32 or d_dtype bf16 instead", - ) - - # ---- SM100+ check ---- - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is not available") - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) - compute_capability = major * 10 + minor - if compute_capability < 100: - raise RuntimeError(f"GroupedGemmDglu requires SM100+ compute capability, " f"but found SM{compute_capability} on device {device}") - - self._is_supported = True - self._logger.debug("check_support completed successfully") - return True - - # --------------------------------------------------------------------- # - # compile - # --------------------------------------------------------------------- # + self._implementation = GroupedGemmDgluBlockScaledAPI(**kwargs) + self._kernel = self._implementation._kernel + self.weight_mode = self._implementation.weight_mode + supported = self._implementation.check_support() + self._is_supported = self._implementation._is_supported + if supported: + self._pending_init_kwargs = None + return supported def compile(self) -> None: - """Compile the kernel.""" - self._logger.debug("Entering compile") - self._ensure_support_checked() - if self._compiled_kernel is not None: - self._logger.debug("Kernel already compiled; skipping recompilation") - return - if self.a_desc.shape[0] == 0: - self._logger.debug("sample valid_m is zero, skipping kernel compilation") - return - - # ---- Instantiate the unified kernel ---- - gemm_dglu = self._kernel( - sf_vec_size=self.sf_vec_size, - acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), - use_2cta_instrs=self.use_2cta_instrs, - mma_tiler_mn=self.mma_tiler_mn, - cluster_shape_mn=self.cluster_shape_mn, - vectorized_f32=self.vector_f32, - discrete_col_sfd=self.discrete_col_sfd, - expert_cnt=self.expert_cnt, - weight_mode=self.weight_mode, - act_func=self.act_func, - use_dynamic_sched=self.use_dynamic_sched, - ) - - hardware_info = cutlass.utils.HardwareInfo() - max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - max_active_clusters -= self.num_cluster_overlap_margin - self._value_error_if( - max_active_clusters <= 0, - "max_active_clusters must be > 0 after applying overlap margin; reduce CUDNNFE_CLUSTER_OVERLAP_MARGIN", - ) - fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) - - # ---- Allocate workspace ---- - workspace_bytes = gemm_dglu.get_workspace_bytes() - self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device="cuda") - - if self.weight_mode == MoEWeightMode.DENSE: - self._compile_dense(gemm_dglu, max_active_clusters, fake_stream) - else: - self._compile_discrete(gemm_dglu, max_active_clusters, fake_stream) - - self._logger.debug("Kernel compiled successfully") - - # -- Dense compile path ------------------------------------------------- # - - def _compile_dense(self, gemm_dglu, max_active_clusters, fake_stream) -> None: - """Compile for dense (contiguous) weight mode.""" - use_full_dynamic = os.environ.get("CUDNN_FE_GROUPED_GEMM_DYNAMIC_MNKL", "1") != "0" - - fake_workspace_ptr = cute.runtime.nullptr( - dtype=cutlass.Uint8, - assumed_align=128, - ) - - if not use_full_dynamic: - valid_m = cute.sym_int(divisibility=256) - - a_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.a_desc.dtype, - shape=(valid_m, *self.a_desc.shape[1:]), - stride_order=self.a_desc.stride_order, - ) - b_cute_fake = self._make_fake_cute_tensor_from_desc(self.b_desc, assumed_align=16) - c_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.c_desc.dtype, - shape=(valid_m, *self.c_desc.shape[1:]), - stride_order=self.c_desc.stride_order, - ) - d_row_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.d_row_desc.dtype, - shape=(valid_m, *self.d_row_desc.shape[1:]), - stride_order=self.d_row_desc.stride_order, - ) - d_col_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.d_col_desc.dtype, - shape=(valid_m, *self.d_col_desc.shape[1:]), - stride_order=self.d_col_desc.stride_order, - ) - - tensor_m_128 = cute.sym_int() - stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfa_cute_fake = self._make_fake_cute_tensor( - dtype=self.sfa_desc.dtype, - shape=(32, 4, tensor_m_128, 4, self.sfa_desc.shape[4], 1), - stride=(16, 4, self.sfa_desc.stride[2], 1, 512, stride_tensor_m_128), - ) - - sfb_cute_fake = self._make_fake_cute_tensor_from_desc(self.sfb_desc, assumed_align=16) - - beta_cute_fake = self._make_fake_cute_tensor_from_desc(self.beta_desc, assumed_align=16) - prob_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.prob_desc.dtype, - shape=(valid_m, 1, 1), - stride_order=self.prob_desc.stride_order, - ) - dprob_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.dprob_desc.dtype, - shape=(valid_m, 1, 1), - stride_order=self.dprob_desc.stride_order, - ) - - sfd_row_fake = None - sfd_col_fake = None - if self.sfd_row_desc is not None: - stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_row_fake = self._make_fake_cute_tensor( - dtype=self.sfd_row_desc.dtype, - shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), - stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), - ) - if self.sfd_col_desc is not None: - rest_m = cute.sym_int(divisibility=1) - stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) - stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_col_fake = self._make_fake_cute_tensor( - dtype=self.sfd_col_desc.dtype, - shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), - stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), - ) - else: - valid_m = cute.sym_int(divisibility=256) - n_sym = cute.sym_int() - n_out_sym = cute.sym_int() - k_sym = cute.sym_int() - l_sym = cute.sym_int() - - a_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.a_desc.dtype, - shape=(valid_m, k_sym, 1), - stride_order=self.a_desc.stride_order, - dynamic_mode=self.a_desc.stride_order[0], - divisibility=32 if self._is_fp4x2(self.ab_dtype) else 16, - ) - b_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.b_desc.dtype, - shape=(n_sym, k_sym, l_sym), - stride_order=self.b_desc.stride_order, - dynamic_mode=self.b_desc.stride_order[0], - divisibility=32 if self._is_fp4x2(self.ab_dtype) else 16, - ) - - c_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.c_desc.dtype, - shape=(valid_m, n_out_sym, 1), - stride_order=self.c_desc.stride_order, - dynamic_mode=self.c_desc.stride_order[0], - divisibility=8 if self._is_f16(self.c_desc.dtype) else 16, - ) - - d_row_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.d_row_desc.dtype, - shape=(valid_m, n_out_sym, 1), - stride_order=self.d_row_desc.stride_order, - dynamic_mode=self.d_row_desc.stride_order[0], - divisibility=8 if self._is_f16(self.d_row_desc.dtype) else 16, - ) - - d_col_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.d_col_desc.dtype, - shape=(valid_m, n_out_sym, 1), - stride_order=self.d_col_desc.stride_order, - dynamic_mode=self.d_col_desc.stride_order[0], - divisibility=8 if self._is_f16(self.d_col_desc.dtype) else 16, - ) - - tensor_m_128 = cute.sym_int() - rest_k = cute.sym_int() - stride_rest_k = cute.sym_int(divisibility=32 * 4 * 4) - stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfa_shape = list(self.sfa_desc.shape) - sfa_shape[2] = tensor_m_128 - sfa_shape[4] = rest_k - sfa_stride = list(self.sfa_desc.stride) - sfa_stride[2] = stride_rest_k - sfa_stride[5] = stride_tensor_m_128 - sfa_cute_fake = self._make_fake_cute_tensor( - dtype=self.sfa_desc.dtype, - shape=tuple(sfa_shape), - stride=tuple(sfa_stride), - ) - - tensor_n_128 = cute.sym_int() - stride_sfb_rest_k = cute.sym_int(divisibility=32 * 4 * 4) - stride_sfb_tensor_n_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfb_cute_fake = self._make_fake_cute_tensor( - dtype=self.sfb_desc.dtype, - shape=(32, 4, tensor_n_128, 4, rest_k, l_sym), - stride=(16, 4, stride_sfb_tensor_n_128, 1, 512, stride_sfb_rest_k), - ) - - beta_cute_fake = self._make_fake_cute_tensor_from_desc(self.beta_desc, assumed_align=16) - prob_cute_fake = self._make_fake_cute_tensor( - dtype=self.prob_desc.dtype, - shape=(valid_m, *self.prob_desc.shape[1:]), - stride=self.prob_desc.stride, - ) - dprob_cute_fake = self._make_fake_cute_tensor( - dtype=self.dprob_desc.dtype, - shape=(valid_m, *self.dprob_desc.shape[1:]), - stride=self.dprob_desc.stride, - ) - - sfd_row_fake = None - sfd_col_fake = None - if self.sfd_row_desc is not None: - rest_n_out = cute.sym_int() - stride_sfd_rest_n_out = cute.sym_int(divisibility=32 * 4 * 4) - stride_sfd_rest_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfd_row_fake = self._make_fake_cute_tensor( - dtype=self.sfd_row_desc.dtype, - shape=(32, 4, tensor_m_128, 4, rest_n_out, 1), - stride=(16, 4, stride_sfd_rest_n_out, 1, 512, stride_sfd_rest_tensor_m_128), - ) - if self.sfd_col_desc is not None: - tensor_n_out_128 = cute.sym_int() - rest_m_dyn = cute.sym_int() - stride_sfd_rest_m = cute.sym_int(divisibility=32 * 4 * 4) - stride_sfd_n_out = cute.sym_int(divisibility=32 * 4 * 4) - sfd_col_fake = self._make_fake_cute_tensor( - dtype=self.sfd_col_desc.dtype, - shape=(32, 4, tensor_n_out_128, 4, rest_m_dyn, 1), - stride=(16, 4, stride_sfd_rest_m, 1, 512, stride_sfd_n_out), - ) - - # Compile with keyword args (dense mode uses the unified __call__ positional order). - dbias_fake = self._make_fake_cute_tensor_from_desc(self.dbias_desc, assumed_align=16) - - _compiled_kernel = cute.compile( - gemm_dglu, - a=a_cute_fake, - b=b_cute_fake, - sfb=sfb_cute_fake, - n=cutlass.Int32(0), - k=cutlass.Int32(0), - b_stride_size=cutlass.Int64(0), - b_major_mode=OperandMajorMode.K, - workspace_ptr=fake_workspace_ptr, - c=c_cute_fake, - d=d_row_cute_fake, - d_col=d_col_cute_fake, - sfa=sfa_cute_fake, - sfd_row_tensor=sfd_row_fake, - sfd_col_tensor=sfd_col_fake, - amax_tensor=self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16), - norm_const_tensor=self._make_fake_cute_tensor_from_desc(self.norm_const_desc, assumed_align=16), - padded_offsets=self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16), - alpha=self._make_fake_cute_tensor_from_desc(self.alpha_desc, assumed_align=16), - beta=beta_cute_fake, - prob=prob_cute_fake, - dprob=dprob_cute_fake, - dbias_tensor=dbias_fake, - max_active_clusters=max_active_clusters, - stream=fake_stream, - epilogue_op=self.epilogue_op, - linear_offset=self.linear_offset, - geglu_alpha=self.geglu_alpha, - glu_clamp_max=self.glu_clamp_max, - glu_clamp_min=self.glu_clamp_min, - options="--enable-tvm-ffi", - ) - - # Cache workspace pointer for the tensor_api closure - cached_workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator - - def tensor_api( - a_tensor: torch.Tensor, - b_tensor: torch.Tensor, - c_tensor: torch.Tensor, - d_row_tensor: torch.Tensor, - d_col_tensor: Optional[torch.Tensor], - sfa_tensor: torch.Tensor, - sfb_tensor: torch.Tensor, - sfd_row_tensor: Optional[torch.Tensor], - sfd_col_tensor: Optional[torch.Tensor], - amax_tensor: Optional[torch.Tensor], - norm_const_tensor: Optional[torch.Tensor], - padded_offsets: torch.Tensor, - alpha_tensor: torch.Tensor, - beta_tensor: torch.Tensor, - prob_tensor: torch.Tensor, - dprob_tensor: torch.Tensor, - dbias_tensor: Optional[torch.Tensor], - stream: cuda.CUstream, - ) -> None: - norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") - _compiled_kernel( - a_tensor, - b_tensor, - sfb_tensor, - cutlass.Int32(0), - cutlass.Int32(0), - cutlass.Int64(0), - cached_workspace_ptr, - c_tensor, - d_row_tensor, - d_col_tensor, - sfa_tensor, - sfd_row_tensor, - sfd_col_tensor, - amax_tensor, - norm_const_tensor, - padded_offsets, - alpha_tensor, - beta_tensor, - prob_tensor, - dprob_tensor, - dbias_tensor, - stream, - ) - - self._compiled_kernel = tensor_api - - # -- Discrete compile path ---------------------------------------------- # - - def _compile_discrete(self, gemm_dglu, max_active_clusters, fake_stream) -> None: - """Compile for discrete (per-expert pointer) weight mode.""" - if len(self.b_shape) == 2: - n, k = self.b_shape - else: - n, k, _ = self.b_shape - - b_major_mode = OperandMajorMode.K if self.b_major == "k" else OperandMajorMode.MN - if self.b_major == "k": - b_stride_size = k - else: - b_stride_size = n - - ab_cutlass_dtype = _convert_to_cutlass_data_type(self.a_desc.dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2) - align = 32 if ab_cutlass_dtype.width == 4 else 16 - - valid_m = cute.sym_int(divisibility=256) - a_tensor = self._make_fake_cute_tensor( - dtype=self.a_desc.dtype, - shape=(valid_m, *self.a_desc.shape[1:]), - stride=(self.a_desc.stride[0], *self.a_desc.stride[1:]), - assumed_align=align, - ) - c_tensor = self._make_fake_cute_tensor( - dtype=self.c_desc.dtype, - shape=(valid_m, *self.c_desc.shape[1:]), - stride=(self.c_desc.stride[0], *self.c_desc.stride[1:]), - ) - d_row_tensor = self._make_fake_cute_compact_tensor( - dtype=self.d_row_desc.dtype, - shape=(valid_m, *self.d_row_desc.shape[1:]), - stride_order=self.d_row_desc.stride_order, - ) - d_col_tensor = self._make_fake_cute_compact_tensor( - dtype=self.d_col_desc.dtype, - shape=(valid_m, *self.d_col_desc.shape[1:]), - stride_order=self.d_col_desc.stride_order, - ) - - tensor_m_128 = cute.sym_int() - stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfa_shape = list(self.sfa_desc.shape) - sfa_shape[2] = tensor_m_128 - sfa_stride = list(self.sfa_desc.stride) - sfa_stride[5] = stride_tensor_m_128 - sfa_tensor = self._make_fake_cute_tensor( - dtype=self.sfa_desc.dtype, - shape=tuple(sfa_shape), - stride=tuple(sfa_stride), - assumed_align=16, - ) - sfd_row_tensor = None - if self.sfd_row_desc is not None: - stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_row_tensor = self._make_fake_cute_tensor( - dtype=self.sfd_row_desc.dtype, - shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), - stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), - assumed_align=16, - ) - sfd_col_tensor = None - if self.sfd_col_desc is not None: - rest_m = cute.sym_int(divisibility=1) - stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) - stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_col_tensor = self._make_fake_cute_tensor( - dtype=self.sfd_col_desc.dtype, - shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), - stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), - assumed_align=16, - ) - amax_tensor = self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16) - norm_const_tensor_cute = self._make_fake_cute_tensor_from_desc(self.norm_const_desc, assumed_align=16) - padded_offsets_tensor = self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16) - alpha_tensor = self._make_fake_cute_tensor_from_desc(self.alpha_desc, assumed_align=16) - beta_tensor = self._make_fake_cute_tensor_from_desc(self.beta_desc, assumed_align=16) - prob_tensor = self._make_fake_cute_tensor( - dtype=self.prob_desc.dtype, - shape=(valid_m, *self.prob_desc.shape[1:]), - stride=self.prob_desc.stride, - assumed_align=16, - ) - dprob_tensor = self._make_fake_cute_tensor( - dtype=self.dprob_desc.dtype, - shape=(valid_m, *self.dprob_desc.shape[1:]), - stride=self.dprob_desc.stride, - assumed_align=16, - ) - dbias_tensor = self._make_fake_cute_tensor_from_desc(self.dbias_desc, assumed_align=16) - - # Compile-time pointer placeholders - b_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - sfb_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - b_ptrs_cute = from_dlpack(b_ptrs_placeholder, assumed_align=8).iterator - sfb_ptrs_cute = from_dlpack(sfb_ptrs_placeholder, assumed_align=8).iterator - - workspace_ptr_cute = from_dlpack(self._workspace, assumed_align=128).iterator - - self._logger.debug("Compiling discrete grouped GEMM dGLU kernel") - _compiled_kernel = cute.compile( - gemm_dglu, - a=a_tensor, - b=b_ptrs_cute, - sfb=sfb_ptrs_cute, - n=cutlass.Int32(n), - k=cutlass.Int32(k), - b_stride_size=cutlass.Int64(b_stride_size), - b_major_mode=b_major_mode, - workspace_ptr=workspace_ptr_cute, - c=c_tensor, - d=d_row_tensor, - d_col=d_col_tensor, - sfa=sfa_tensor, - sfd_row_tensor=sfd_row_tensor, - sfd_col_tensor=sfd_col_tensor, - amax_tensor=amax_tensor, - norm_const_tensor=norm_const_tensor_cute, - padded_offsets=padded_offsets_tensor, - alpha=alpha_tensor, - beta=beta_tensor, - prob=prob_tensor, - dprob=dprob_tensor, - dbias_tensor=dbias_tensor, - max_active_clusters=max_active_clusters, - stream=fake_stream, - epilogue_op=self.epilogue_op, - linear_offset=self.linear_offset, - geglu_alpha=self.geglu_alpha, - glu_clamp_max=self.glu_clamp_max, - glu_clamp_min=self.glu_clamp_min, - options="--enable-tvm-ffi", - ) - - self._n = n - self._k = k - self._b_stride_size = b_stride_size - - # Cache constant values for execute() closure - cached_workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator - cached_n = cutlass.Int32(self._n) - cached_k = cutlass.Int32(self._k) - cached_b_stride = cutlass.Int64(self._b_stride_size) - - def tensor_api( - a_tensor: torch.Tensor, - b_ptrs_device: torch.Tensor, - sfb_ptrs_device: torch.Tensor, - c_tensor: torch.Tensor, - d_row_tensor: torch.Tensor, - d_col_tensor: Optional[torch.Tensor], - sfa_tensor: torch.Tensor, - sfd_row_tensor: Optional[torch.Tensor], - sfd_col_tensor: Optional[torch.Tensor], - amax_tensor: Optional[torch.Tensor], - norm_const_tensor: Optional[torch.Tensor], - padded_offsets: torch.Tensor, - alpha_tensor: torch.Tensor, - beta_tensor: torch.Tensor, - prob_tensor: torch.Tensor, - dprob_tensor: torch.Tensor, - dbias_tensor: Optional[torch.Tensor], - stream: cuda.CUstream, - ) -> None: - norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") - b_ptrs_addr = int(b_ptrs_device.data_ptr()) - sfb_ptrs_addr = int(sfb_ptrs_device.data_ptr()) - - _compiled_kernel( - a_tensor, - b_ptrs_addr, - sfb_ptrs_addr, - cached_n, - cached_k, - cached_b_stride, - cached_workspace_ptr, - c_tensor, - d_row_tensor, - d_col_tensor, - sfa_tensor, - sfd_row_tensor, - sfd_col_tensor, - amax_tensor, - norm_const_tensor, - padded_offsets, - alpha_tensor, - beta_tensor, - prob_tensor, - dprob_tensor, - dbias_tensor, - stream, - ) - - self._compiled_kernel = tensor_api - - # --------------------------------------------------------------------- # - # execute - # --------------------------------------------------------------------- # - + if self._implementation is None: + self.check_support() + if self._is_supported: + self._implementation._is_supported = True + self._implementation.compile() + self._is_supported = self._implementation._is_supported + self._compiled_kernel = self._implementation._compiled_kernel + + # BF16 implementation + @overload + def execute( + self, + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_row_tensor: torch.Tensor, + d_col_tensor: None, + sfa_tensor: None, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + beta_tensor: torch.Tensor, + prob_tensor: torch.Tensor, + dprob_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor] = None, + *, + sfb_tensor: None = None, + sfb_ptrs: None = None, + sfd_row_tensor: None = None, + sfd_col_tensor: None = None, + amax_tensor: None = None, + norm_const_tensor: None = None, + ) -> None: ... + + # Block-scaled implementation + @overload def execute( self, a_tensor: torch.Tensor, c_tensor: torch.Tensor, d_row_tensor: torch.Tensor, - d_col_tensor: torch.Tensor, + d_col_tensor: Optional[torch.Tensor], sfa_tensor: torch.Tensor, padded_offsets: torch.Tensor, alpha_tensor: torch.Tensor, beta_tensor: torch.Tensor, prob_tensor: torch.Tensor, dprob_tensor: torch.Tensor, - # Dense mode: + b_tensor: Optional[torch.Tensor] = None, + *, + sfb_tensor: Optional[torch.Tensor] = None, + sfb_ptrs: Optional[torch.Tensor] = None, + sfd_row_tensor: Optional[torch.Tensor] = None, + sfd_col_tensor: Optional[torch.Tensor] = None, + amax_tensor: Optional[torch.Tensor] = None, + norm_const_tensor: Optional[torch.Tensor] = None, + ) -> None: ... + + def execute( + self, + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_row_tensor: torch.Tensor, + d_col_tensor: Optional[torch.Tensor], + sfa_tensor: Optional[torch.Tensor], + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + beta_tensor: torch.Tensor, + prob_tensor: torch.Tensor, + dprob_tensor: torch.Tensor, b_tensor: Optional[torch.Tensor] = None, sfb_tensor: Optional[torch.Tensor] = None, dbias_tensor: Optional[torch.Tensor] = None, - # Discrete mode: b_ptrs: Optional[torch.Tensor] = None, sfb_ptrs: Optional[torch.Tensor] = None, - # Optional: sfd_row_tensor: Optional[torch.Tensor] = None, sfd_col_tensor: Optional[torch.Tensor] = None, amax_tensor: Optional[torch.Tensor] = None, norm_const_tensor: Optional[torch.Tensor] = None, current_stream: Optional[cuda.CUstream] = None, ) -> None: - """Execute the compiled kernel. - - For dense mode, supply ``b_tensor`` and ``sfb_tensor``. - For discrete mode, supply ``b_ptrs`` and ``sfb_ptrs``. - - :param a_tensor: Input A tensor (gradient input) - :param c_tensor: Forward activations input - :param d_row_tensor: Output D row tensor - :param d_col_tensor: Output D column tensor - :param sfa_tensor: Scale factor A - :param padded_offsets: End offset per expert after padding - :param alpha_tensor: Per-group alpha scaling factors - :param beta_tensor: Per-group beta scaling factors - :param prob_tensor: Per-row probability (from forward) - :param dprob_tensor: Gradient of probability (output, must be zero-initialized) - :param b_tensor: (Dense) Input B tensor (weights) - :param sfb_tensor: (Dense) Scale factor B - :param dbias_tensor: Optional dbias output tensor. - :param b_ptrs: (Discrete) 1-D int64 device tensor of per-expert B data pointers - :param sfb_ptrs: (Discrete) 1-D int64 device tensor of per-expert SFB data pointers - :param sfd_row_tensor: Optional row scale factor D - :param sfd_col_tensor: Optional column scale factor D - :param amax_tensor: Optional amax tensor - :param norm_const_tensor: Optional normalization constant - :param current_stream: CUDA stream - """ - self._logger.debug("Entering execute") - current_stream = self._get_default_stream(current_stream) - - if a_tensor.shape[0] == 0: - self._logger.debug("execute: valid_m is zero, skipping kernel execution") - return - self._runtime_error_if( - self._compiled_kernel is None, - "Kernel not compiled; call compile() first", - ) - - self._logger.debug("Executing grouped GEMM dGLU kernel") - if self._has_dbias: - self._value_error_if( - dbias_tensor is None, - "dbias_tensor is required when GroupedGemmDgluSm100 is configured with sample_dbias", + if self._implementation is None: + raise RuntimeError("Kernel not compiled; call compile() first") + if self.backend is GroupedGemmBackend.BF16: + controls = ( + ("sfa_tensor", sfa_tensor), + ("sfb_tensor", sfb_tensor), + ("sfb_ptrs", sfb_ptrs), + ("d_col_tensor", d_col_tensor), + ("sfd_row_tensor", sfd_row_tensor), + ("sfd_col_tensor", sfd_col_tensor), + ("amax_tensor", amax_tensor), + ("norm_const_tensor", norm_const_tensor), ) - - if self.weight_mode == MoEWeightMode.DENSE: - self._compiled_kernel( + forbidden = [name for name, value in controls if value is not None] + if forbidden: + raise ValueError(f"grouped_gemm_dglu_sm100: BF16 forbids scale control {forbidden[0]}") + self._implementation.execute( a_tensor=a_tensor, - b_tensor=b_tensor, c_tensor=c_tensor, d_row_tensor=d_row_tensor, - d_col_tensor=d_col_tensor, - sfa_tensor=sfa_tensor, - sfb_tensor=sfb_tensor, - sfd_row_tensor=sfd_row_tensor, - sfd_col_tensor=sfd_col_tensor, - amax_tensor=amax_tensor, - norm_const_tensor=norm_const_tensor, padded_offsets=padded_offsets, alpha_tensor=alpha_tensor, beta_tensor=beta_tensor, prob_tensor=prob_tensor, dprob_tensor=dprob_tensor, + b_tensor=b_tensor, + b_ptrs=b_ptrs, dbias_tensor=dbias_tensor, - stream=current_stream, + linear_offset=self.linear_offset, + current_stream=current_stream, ) else: - self._compiled_kernel( + self._implementation.execute( a_tensor=a_tensor, - b_ptrs_device=b_ptrs, - sfb_ptrs_device=sfb_ptrs, c_tensor=c_tensor, d_row_tensor=d_row_tensor, d_col_tensor=d_col_tensor, sfa_tensor=sfa_tensor, - sfd_row_tensor=sfd_row_tensor, - sfd_col_tensor=sfd_col_tensor, - amax_tensor=amax_tensor, - norm_const_tensor=norm_const_tensor, padded_offsets=padded_offsets, alpha_tensor=alpha_tensor, beta_tensor=beta_tensor, prob_tensor=prob_tensor, dprob_tensor=dprob_tensor, + b_tensor=b_tensor, + sfb_tensor=sfb_tensor, dbias_tensor=dbias_tensor, - stream=current_stream, + b_ptrs=b_ptrs, + sfb_ptrs=sfb_ptrs, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + amax_tensor=amax_tensor, + norm_const_tensor=norm_const_tensor, + current_stream=current_stream, ) - - self._logger.debug("Execute completed") + self._is_supported = self._implementation._is_supported + self._compiled_kernel = self._implementation._compiled_kernel # --------------------------------------------------------------------------- # @@ -1285,45 +410,7 @@ def execute( _cache_of_GroupedGemmDgluSm100Objects = {} -def grouped_gemm_dglu_wrapper_sm100( - a_tensor: torch.Tensor, - c_tensor: torch.Tensor, - sfa_tensor: torch.Tensor, - padded_offsets: torch.Tensor, - alpha_tensor: torch.Tensor, - beta_tensor: torch.Tensor, - prob_tensor: torch.Tensor, - dprob_tensor: torch.Tensor, - # generate_dbias is optional in both modes: - b_tensor: Optional[torch.Tensor] = None, - sfb_tensor: Optional[torch.Tensor] = None, - generate_dbias: bool = False, - # Discrete mode: - b_ptrs: Optional[torch.Tensor] = None, - sfb_ptrs: Optional[torch.Tensor] = None, - n: Optional[int] = None, - b_dtype: Optional[torch.dtype] = None, - b_major: str = "k", - # Common: - norm_const_tensor: Optional[torch.Tensor] = None, - acc_dtype: torch.dtype = torch.float32, - d_dtype: torch.dtype = torch.bfloat16, - cd_major: str = "n", - mma_tiler_mn: Tuple[int, int] = (256, 256), - cluster_shape_mn: Optional[Tuple[int, int]] = None, - sf_vec_size: int = 16, - vector_f32: bool = False, - m_aligned: int = 256, - discrete_col_sfd: bool = False, - act_func: str = "dswiglu", - linear_offset: Optional[float] = None, - geglu_alpha: float = 1.702, - glu_clamp_max: float = 7.0, - glu_clamp_min: float = -7.0, - epilogue_op: Optional[str] = None, - use_dynamic_sched: bool = False, - current_stream: Optional[cuda.CUstream] = None, -) -> TupleDict: +def _grouped_gemm_dglu_block_scaled_call(call: DgluCall) -> TupleDict: """Convenience wrapper for grouped GEMM dGLU backward operation. Auto-detects dense vs. discrete mode based on which weight arguments @@ -1388,6 +475,41 @@ def grouped_gemm_dglu_wrapper_sm100( """ from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor + a_tensor = call.a_tensor + c_tensor = call.c_tensor + sfa_tensor = call.sfa_tensor + padded_offsets = call.padded_offsets + alpha_tensor = call.alpha_tensor + beta_tensor = call.beta_tensor + prob_tensor = call.prob_tensor + dprob_tensor = call.dprob_tensor + b_tensor = call.b_tensor + sfb_tensor = call.sfb_tensor + generate_dbias = call.generate_dbias + b_ptrs = call.b_ptrs + sfb_ptrs = call.sfb_ptrs + n = call.n + b_dtype = call.b_dtype + b_major = call.b_major + norm_const_tensor = call.norm_const_tensor + acc_dtype = call.acc_dtype + d_dtype = call.d_dtype + cd_major = call.cd_major + mma_tiler_mn = call.mma_tiler_mn + cluster_shape_mn = call.cluster_shape_mn + sf_vec_size = call.sf_vec_size + vector_f32 = call.vector_f32 + m_aligned = call.m_aligned + discrete_col_sfd = call.discrete_col_sfd + act_func = call.act_func + linear_offset = call.linear_offset + geglu_alpha = call.geglu_alpha + glu_clamp_max = call.glu_clamp_max + glu_clamp_min = call.glu_clamp_min + epilogue_op = call.epilogue_op + use_dynamic_sched = call.use_dynamic_sched + current_stream = call.current_stream + # Resolve linear_offset default: None means "use the activation-derived # default" (1.0 for dgeglu, 0.0 for dswiglu). if linear_offset is None: @@ -1720,3 +842,359 @@ def dynamic_m_tensor_signature( sfd_row_tensor=sfd_row_tensor, sfd_col_tensor=sfd_col_tensor, ) + + +def _normalize_dglu_call( + call: DgluCall, +) -> tuple[DgluCall, GroupedGemmBackend]: + from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor + + is_dense = call.b_tensor is not None + is_discrete = call.b_ptrs is not None + if is_dense and is_discrete: + raise ValueError("Provide either (b_tensor, sfb_tensor) or (b_ptrs, sfb_ptrs), not both") + if not is_dense and not is_discrete: + raise ValueError("Must provide either (b_tensor, sfb_tensor) or (b_ptrs, sfb_ptrs)") + if call.a_tensor.ndim != 3 or call.a_tensor.shape[2] != 1: + raise ValueError(f"a_tensor must have shape (m, k, 1), got {tuple(call.a_tensor.shape)}") + + valid_m, k, _ = call.a_tensor.shape + if is_dense: + if call.b_tensor.ndim != 3: + raise ValueError(f"b_tensor must have shape (n, k, experts), got {tuple(call.b_tensor.shape)}") + n_weight, b_k, num_experts = call.b_tensor.shape + if b_k != k: + raise ValueError(f"b_tensor K dimension ({b_k}) must match a_tensor ({k})") + defining_b_dtype = call.b_tensor.dtype + b_shape = None + weight_mode = MoEWeightMode.DENSE + if call.n is not None or call.b_dtype is not None: + raise ValueError("Dense mode forbids n and b_dtype") + else: + _require_pointer_tensor(call.b_ptrs, "b_ptrs") + num_experts = call.b_ptrs.numel() + if call.n is None or call.b_dtype is None: + raise ValueError("n and b_dtype are required for discrete mode") + n_weight = call.n + defining_b_dtype = call.b_dtype + b_shape = (n_weight, k) + weight_mode = MoEWeightMode.DISCRETE + + backend = select_grouped_gemm_backend( + operation="grouped_gemm_dglu_sm100", + a_dtype=call.a_tensor.dtype, + b_dtype=defining_b_dtype, + scale_controls=( + ("sfa_tensor", call.sfa_tensor), + ("sfb_tensor", call.sfb_tensor), + ("sfb_ptrs", call.sfb_ptrs), + ("norm_const_tensor", call.norm_const_tensor), + ("sf_vec_size", call.sf_vec_size if call.sf_vec_size != 16 else None), + ( + "discrete_col_sfd", + call.discrete_col_sfd if call.discrete_col_sfd else None, + ), + ("geglu_alpha", call.geglu_alpha if call.geglu_alpha != 1.702 else None), + ( + "glu_clamp_max", + call.glu_clamp_max if call.glu_clamp_max != 7.0 else None, + ), + ( + "glu_clamp_min", + call.glu_clamp_min if call.glu_clamp_min != -7.0 else None, + ), + ( + "epilogue_op", + call.epilogue_op if call.epilogue_op not in (None, "none", "identity") else None, + ), + ), + block_scaled_dtype_pairs=_BLOCK_SCALED_DTYPE_PAIRS, + ) + linear_offset = call.linear_offset + if linear_offset is None: + linear_offset = 1.0 if call.act_func == "dgeglu" else 0.0 + normalized = replace( + call, + linear_offset=linear_offset, + weight_mode=weight_mode, + b_shape=b_shape, + num_experts=num_experts, + ) + if backend is GroupedGemmBackend.BLOCK_SCALED: + return normalized, backend + + if call.cd_major != "n": + raise ValueError(f"cd_major must be 'n', got {call.cd_major}") + if call.act_func not in ("dswiglu", "dgeglu"): + raise ValueError(f"act_func must be 'dswiglu' or 'dgeglu', got {call.act_func}") + if call.d_dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise ValueError(f"d_dtype must be BF16, FP16, or FP32, got {call.d_dtype}") + if call.m_aligned != 256: + raise ValueError(f"m_aligned must be 256, got {call.m_aligned}") + if valid_m % 256 != 0: + raise ValueError(f"a_tensor M dimension must be 256-aligned, got {valid_m}") + if n_weight <= 0 or n_weight % 32 != 0: + raise ValueError(f"N must be positive and divisible by 32, got {n_weight}") + two_n = 2 * n_weight + if tuple(call.c_tensor.shape) != (valid_m, two_n, 1): + raise ValueError(f"c_tensor must have shape {(valid_m, two_n, 1)}, got {tuple(call.c_tensor.shape)}") + if tuple(call.prob_tensor.shape) != (valid_m, 1, 1): + raise ValueError(f"prob_tensor must have shape {(valid_m, 1, 1)}, got {tuple(call.prob_tensor.shape)}") + if tuple(call.dprob_tensor.shape) != (valid_m, 1, 1): + raise ValueError(f"dprob_tensor must have shape {(valid_m, 1, 1)}, got {tuple(call.dprob_tensor.shape)}") + if call.dprob_tensor.dtype != torch.float32: + raise ValueError(f"dprob_tensor must have dtype torch.float32, got {call.dprob_tensor.dtype}") + if is_discrete and call.b_ptrs.numel() != call.padded_offsets.numel(): + raise ValueError(f"b_ptrs length mismatch: expected {call.padded_offsets.numel()}, " f"got {call.b_ptrs.numel()}") + if tuple(call.padded_offsets.shape) != (num_experts,): + raise ValueError(f"padded_offsets length mismatch: expected {num_experts}, got {call.padded_offsets.numel()}") + if tuple(call.alpha_tensor.shape) != (num_experts,): + raise ValueError(f"alpha_tensor must have shape {(num_experts,)}, got {tuple(call.alpha_tensor.shape)}") + if tuple(call.beta_tensor.shape) != (num_experts,): + raise ValueError(f"beta_tensor must have shape {(num_experts,)}, got {tuple(call.beta_tensor.shape)}") + if is_discrete: + if call.b_ptrs.numel() != num_experts: + raise ValueError(f"b_ptrs length mismatch: expected {num_experts}, got {call.b_ptrs.numel()}") + if call.b_ptrs.device != call.a_tensor.device: + raise ValueError(f"b_ptrs must be on the same device as a_tensor ({call.a_tensor.device}), " f"got {call.b_ptrs.device}") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + major, minor = torch.cuda.get_device_capability(call.a_tensor.device) + capability = major * 10 + minor + if capability < 100: + raise RuntimeError(f"GroupedGemmDgluSm100 requires SM100+, found SM{capability}") + return normalized, backend + + +def _dglu_stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: + return tuple( + index + for index, _ in sorted( + enumerate(tensor.stride()), + key=lambda item: (item[1], tensor.shape[item[0]]), + ) + ) + + +def _dglu_tensor_signature(tensor: Optional[torch.Tensor], *, dynamic_m: bool = False) -> tuple: + if tensor is None: + return (None, None, None, None) + shape = (None, *tuple(tensor.shape[1:])) if dynamic_m else tuple(tensor.shape) + return ( + shape, + _dglu_stride_order(tensor), + tensor.dtype, + (tensor.device.type, tensor.device.index), + ) + + +def _grouped_gemm_dglu_bf16_call(call: DgluCall) -> TupleDict: + valid_m = call.a_tensor.shape[0] + n_weight = call.b_tensor.shape[0] if call.b_tensor is not None else call.n + two_n = 2 * n_weight + with _torch_stream_context(call.current_stream, call.a_tensor.device): + d_row_tensor = torch.empty_strided( + (valid_m, two_n, 1), + (two_n, 1, valid_m * two_n), + dtype=call.d_dtype, + device=call.a_tensor.device, + ) + dbias_tensor = ( + torch.zeros( + (call.num_experts, two_n, 1), + dtype=torch.bfloat16, + device=call.a_tensor.device, + ) + if call.generate_dbias + else None + ) + + overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + workspace_bytes = (128 * call.num_experts if call.weight_mode == MoEWeightMode.DISCRETE else 0) + (4 if call.use_dynamic_sched else 0) + cache_key = backend_cache_key( + GroupedGemmBackend.BF16, + call.weight_mode, + call.act_func, + _dglu_tensor_signature(call.a_tensor, dynamic_m=True), + _dglu_tensor_signature(call.b_tensor), + call.b_shape, + call.b_dtype, + _dglu_tensor_signature(call.c_tensor, dynamic_m=True), + _dglu_tensor_signature(d_row_tensor, dynamic_m=True), + _dglu_tensor_signature(call.padded_offsets), + _dglu_tensor_signature(call.alpha_tensor), + _dglu_tensor_signature(call.beta_tensor), + _dglu_tensor_signature(call.prob_tensor, dynamic_m=True), + _dglu_tensor_signature(call.dprob_tensor, dynamic_m=True), + _dglu_tensor_signature(dbias_tensor), + ( + ( + tuple(call.b_ptrs.shape), + tuple(call.b_ptrs.stride()), + call.b_ptrs.dtype, + (call.b_ptrs.device.type, call.b_ptrs.device.index), + ) + if call.b_ptrs is not None + else None + ), + call.acc_dtype, + call.d_dtype, + call.mma_tiler_mn, + call.cluster_shape_mn, + call.vector_f32, + call.m_aligned, + call.b_major, + call.use_dynamic_sched, + workspace_bytes, + (call.a_tensor.device.type, call.a_tensor.device.index), + overlap_margin, + ) + + if cache_key in _cache_of_GroupedGemmDgluSm100Objects: + api = _cache_of_GroupedGemmDgluSm100Objects[cache_key] + else: + api = GroupedGemmDgluSm100( + sample_a=call.a_tensor, + sample_c=call.c_tensor, + sample_d_row=d_row_tensor, + sample_d_col=None, + sample_sfa=None, + sample_padded_offsets=call.padded_offsets, + sample_alpha=call.alpha_tensor, + sample_beta=call.beta_tensor, + sample_prob=call.prob_tensor, + sample_dprob=call.dprob_tensor, + sample_b=call.b_tensor, + sample_sfb=None, + sample_dbias=dbias_tensor, + num_experts=(call.num_experts if call.weight_mode == MoEWeightMode.DISCRETE else None), + b_shape=call.b_shape, + b_dtype=call.b_dtype, + sample_sfd_row=None, + sample_sfd_col=None, + sample_amax=None, + sample_norm_const=None, + acc_dtype=call.acc_dtype, + mma_tiler_mn=call.mma_tiler_mn, + cluster_shape_mn=call.cluster_shape_mn, + sf_vec_size=16, + vector_f32=call.vector_f32, + m_aligned=call.m_aligned, + discrete_col_sfd=False, + act_func=call.act_func, + b_major=call.b_major, + epilogue_op=None, + use_dynamic_sched=call.use_dynamic_sched, + linear_offset=call.linear_offset, + geglu_alpha=1.702, + glu_clamp_max=7.0, + glu_clamp_min=-7.0, + ) + if not api.check_support(): + raise RuntimeError("Unsupported BF16 configuration") + api.compile() + _cache_of_GroupedGemmDgluSm100Objects[cache_key] = api + + api._implementation.execute( + a_tensor=call.a_tensor, + c_tensor=call.c_tensor, + d_row_tensor=d_row_tensor, + padded_offsets=call.padded_offsets, + alpha_tensor=call.alpha_tensor, + beta_tensor=call.beta_tensor, + prob_tensor=call.prob_tensor, + dprob_tensor=call.dprob_tensor, + b_tensor=call.b_tensor, + b_ptrs=call.b_ptrs, + dbias_tensor=dbias_tensor, + linear_offset=call.linear_offset, + current_stream=call.current_stream, + ) + return TupleDict( + d_row_tensor=d_row_tensor, + d_col_tensor=None, + dprob_tensor=call.dprob_tensor, + dbias_tensor=dbias_tensor, + amax_tensor=None, + sfd_row_tensor=None, + sfd_col_tensor=None, + ) + + +def grouped_gemm_dglu_wrapper_sm100( + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + sfa_tensor: Optional[torch.Tensor], + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + beta_tensor: torch.Tensor, + prob_tensor: torch.Tensor, + dprob_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor] = None, + sfb_tensor: Optional[torch.Tensor] = None, + generate_dbias: bool = False, + b_ptrs: Optional[torch.Tensor] = None, + sfb_ptrs: Optional[torch.Tensor] = None, + n: Optional[int] = None, + b_dtype: Optional[torch.dtype] = None, + b_major: str = "k", + norm_const_tensor: Optional[torch.Tensor] = None, + acc_dtype: torch.dtype = torch.float32, + d_dtype: torch.dtype = torch.bfloat16, + cd_major: str = "n", + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "dswiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + epilogue_op: Optional[str] = None, + use_dynamic_sched: bool = False, + current_stream: Optional[cuda.CUstream] = None, +) -> TupleDict: + """Dispatch grouped GEMM dGLU once from an immutable normalized call.""" + call = DgluCall( + a_tensor=a_tensor, + c_tensor=c_tensor, + sfa_tensor=sfa_tensor, + padded_offsets=padded_offsets, + alpha_tensor=alpha_tensor, + beta_tensor=beta_tensor, + prob_tensor=prob_tensor, + dprob_tensor=dprob_tensor, + b_tensor=b_tensor, + sfb_tensor=sfb_tensor, + generate_dbias=generate_dbias, + b_ptrs=b_ptrs, + sfb_ptrs=sfb_ptrs, + n=n, + b_dtype=b_dtype, + b_major=b_major, + norm_const_tensor=norm_const_tensor, + acc_dtype=acc_dtype, + d_dtype=d_dtype, + cd_major=cd_major, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + act_func=act_func, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + epilogue_op=epilogue_op, + use_dynamic_sched=use_dynamic_sched, + current_stream=current_stream, + ) + normalized, backend = _normalize_dglu_call(call) + if backend is GroupedGemmBackend.BF16: + return _grouped_gemm_dglu_bf16_call(normalized) + return _grouped_gemm_dglu_block_scaled_call(normalized) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_dglu/moe_grouped_gemm_dglu_dbias.py b/python/cudnn/grouped_gemm/grouped_gemm_dglu/moe_grouped_gemm_dglu_dbias.py new file mode 100644 index 000000000..36a11dbdc --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_dglu/moe_grouped_gemm_dglu_dbias.py @@ -0,0 +1,2311 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +MoE BF16 Grouped GEMM Kernel with dGLU (dSwiGLU/dGeGLU) Backward Fusion. + +Supports: + - Static / Dynamic persistent tile scheduling (MoEPersistentTileScheduler) + - Dense (contiguous 3-D B) / Discrete (per-expert pointer array B) weight layout + - dGLU backward activation fusion (dSwiGLU / dGeGLU) + - Optional dBias reduction + +This module contains only the kernel class. +MoE scheduler components live in moe_persistent_scheduler.py / moe_sched_extension.py / moe_utils.py. +""" + +from typing import Type, Tuple, Union, Optional +from functools import partial + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 +from cutlass.cutlass_dsl import T +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils + +from cutlass.cute.typing import Float32, Int32, AddressSpace +from cutlass._mlir.dialects import llvm +from cutlass._mlir.dialects import vector + +from ..moe_persistent_scheduler import ( + MoEPersistentTileScheduler, + MoESchedulerParams, + MoEWorkTileInfo, +) +from ..moe_utils import ( + MoEWeightMode, + TensormapWorkspace, + store_tma_desc, +) +from ..moe_sched_extension import ( + DiscreteWeightGroupedGemmSchedExtension, + ContiguousGroupedGemmSchedExtension, +) +from ..moe_kernel_helpers import ( + fmin, + fmax, + atomic_add_float32, + sigmoid_f32, + can_implement_bf16_grouped_gemm, + compute_grid, +) + + +def atomic_add_bf16x2(ptr, val_fp32_lo, val_fp32_hi, *, loc=None, ip=None): + """Packed BF16x2 atomic reduction to global memory.""" + lo_ir = val_fp32_lo.ir_value(loc=loc, ip=ip) + hi_ir = val_fp32_hi.ir_value(loc=loc, ip=ip) + llvm.inline_asm( + None, + [ptr, hi_ir, lo_ir], + "{ .reg .b32 packed; cvt.rn.bf16x2.f32 packed, $1, $2; red.global.add.noftz.bf16x2 [$0], packed; }", + "l,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +class MoEGroupedGemmDgluDbiasBf16Kernel: + """Plain BF16 grouped GEMM kernel with MoE scheduling and dGLU backward fusion. + + The kernel is organized as persistent scheduler, A/B TMA load, MMA, C-load, + and epilogue warps. The epilogue computes dSwiGLU or dGeGLU gradients, + reduces dprob, optionally reduces dbias, and writes D. + """ + + # Fixed pad size for user-side padding (decoupled from kernel tile size) + FIX_PAD_SIZE = 256 + + @staticmethod + def can_implement( + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m: int, + n: int, + k: int, + l: int, + a_major: str, + b_major: str, + cd_major: str, + m_aligned: int, + act_func: str, + ) -> bool: + return can_implement_bf16_grouped_gemm( + ab_dtype, + c_dtype, + d_dtype, + acc_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + m, + n, + k, + l, + a_major, + b_major, + cd_major, + m_aligned, + fix_pad_size=MoEGroupedGemmDgluDbiasBf16Kernel.FIX_PAD_SIZE, + n_align=32, + tile_n_align=32, + ) and act_func in ["dswiglu", "dgeglu"] + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + vectorized_f32: bool, + expert_cnt: int, + weight_mode: MoEWeightMode = MoEWeightMode.DISCRETE, + use_dynamic_sched: bool = False, + act_func: str = "dswiglu", + ): + # Validate FIX_PAD_SIZE compatibility with tile size + mma_tile_m = mma_tiler_mn[0] + if self.FIX_PAD_SIZE % mma_tile_m != 0: + raise ValueError( + f"FIX_PAD_SIZE ({self.FIX_PAD_SIZE}) must be divisible by " f"mma_tiler_mn[0] ({mma_tile_m}). " f"Supported mma_tiler_mn[0] values: 128, 256." + ) + if expert_cnt > 1024: + raise ValueError("Expert count > 1024 is not supported.") + if not isinstance(weight_mode, MoEWeightMode): + raise TypeError(f"weight_mode must be a MoEWeightMode, got {type(weight_mode)}") + if act_func not in ["dswiglu", "dgeglu"]: + raise ValueError(f"Invalid activation function: {act_func}") + + self.expert_cnt = expert_cnt + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler = (*mma_tiler_mn, 1) + + self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + + self.occupancy = 1 + self.epilog_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.epilog_load_tma_id = 6 + self.sched_warp_id = 7 + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + ( + *self.epilog_warp_id, + self.mma_warp_id, + self.tma_warp_id, + self.epilog_load_tma_id, + self.sched_warp_id, + ) + ) + self.threads_wo_sched = self.threads_per_warp * len( + ( + *self.epilog_warp_id, + self.mma_warp_id, + self.tma_warp_id, + self.epilog_load_tma_id, + ) + ) + + # Set barrier for cta sync, epilogue sync and tmem ptr sync + self.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=32 * len(self.epilog_warp_id), + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), + ) + self.sched_sync_barrier = pipeline.NamedBarrier( + barrier_id=4, + num_threads=self.threads_per_warp, + ) + self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + + self.vectorized_f32 = vectorized_f32 + self.use_dynamic_sched = use_dynamic_sched + self.store_d_directly = False + + self.num_epilog_warps = len(self.epilog_warp_id) + + self.weight_mode = weight_mode + + self.act_func = act_func + + def _setup_attributes(self): + """Set up input-dependent BF16 GEMM attributes.""" + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + self.mma_tiler_d = ( + self.mma_tiler[0], + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_d = ( + self.mma_tiler_d[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_d[1], + self.mma_tiler_d[2], + ) + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + # Set epilogue subtile + self.epi_tile = (128, 32) + self.epi_tile_cnt = ( + self.cta_tile_shape_mnk_d[0] // self.epi_tile[0], + self.cta_tile_shape_mnk_d[1] // self.epi_tile[1], + ) + + ( + self.num_acc_stage, + self.num_ab_stage, + self.num_c_stage, + self.num_d_stage, + self.num_tile_stage, + ) = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.epi_tile, + self.c_dtype, + self.c_layout, + self.d_dtype, + self.d_layout, + self.num_smem_capacity, + self.occupancy, + self.store_d_directly, + self.generate_dbias, + ) + + # TMEM accumulator columns: derive from the actual accumulator fragment and + # round up to a valid power-of-two allocation (handles non-power-of-two tile N), + # matching the plain BF16 grouped GEMM kernel instead of always reserving 512. + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + + # Compute A/B/C/D shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + + self.d_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.d_dtype, + self.d_layout, + self.epi_tile, + 1 if self.store_d_directly else self.num_d_stage, + ) + + self.epilogue_prefetch_more = False + self.generate_dprob = True + + def get_desc_workspace_bytes(self) -> int: + """Return descriptor workspace size in bytes.""" + if self.weight_mode == MoEWeightMode.DISCRETE: + return TensormapWorkspace.size_bytes(1, self.expert_cnt) + return 0 + + def get_workspace_bytes(self) -> int: + """Return descriptor workspace plus optional dynamic scheduler state.""" + desc_workspace_bytes = self.get_desc_workspace_bytes() + dynamic_sched_bytes = 4 if self.use_dynamic_sched else 0 + return desc_workspace_bytes + dynamic_sched_bytes + + @cute.jit + def _get_sched_counter_ptr(self, workspace_ptr): + counter_addr = workspace_ptr.toint() + self.get_desc_workspace_bytes() + return cute.make_ptr( + cutlass.Int32, + counter_addr, + AddressSpace.gmem, + assumed_align=4, + ) + + @cute.kernel + def helper_kernel( + self, + ptrs_b: cute.Pointer, + n: Int32, + k: Int32, + b_stride_size: cutlass.Int64, + b_major_mode: cutlass.Constexpr, + workspace_ptr, + tiled_mma_arg: cute.TiledMma, + b_smem_layout_arg, + cluster_layout_vmnk_shape_arg: cutlass.Constexpr, + ): + """Pre-main-kernel initialization. + + Launched with grid=(expert_cnt, 1, 1) for discrete mode, or + grid=(1, 1, 1) for dense+dynamic mode. + + Discrete weight: each block builds a B TMA descriptor for one expert. + Dynamic sched: block 0 resets the atomic tile counter to 0. + """ + expert_idx = cute.arch.block_idx()[0] + + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + b_tma_op_arg = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma_arg.thr_id) + + b_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), cute.make_layout((self.expert_cnt,)) + ) + + c1 = cutlass.Int32(1) + c0 = cutlass.Int64(0) + c1_64 = 1 + if cutlass.const_expr(b_major_mode == OperandMajorMode.K): + stride_n = b_stride_size + stride_k = c1_64 + else: + stride_n = c1_64 + stride_k = b_stride_size + + b_ptr_val = b_ptr_tensor[expert_idx] + b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) + b_tensor_i = cute.make_tensor( + b_ptr, + cute.make_layout((n, k, c1), stride=(stride_n, stride_k, c0)), + ) + tma_atom_b, _ = cute.nvgpu.make_tiled_tma_atom_B( + b_tma_op_arg, + b_tensor_i, + b_smem_layout_arg, + self.mma_tiler, + tiled_mma_arg, + cluster_layout_vmnk_shape_arg, + ) + + workspace = TensormapWorkspace(workspace_ptr, ["b"]) + store_tma_desc(tma_atom_b, workspace.get_ptr("b", expert_idx)) + + if cutlass.const_expr(self.use_dynamic_sched): + if expert_idx == cutlass.Int32(0): + sched_counter = cute.make_tensor( + self._get_sched_counter_ptr(workspace_ptr), + cute.make_layout(1), + ) + sched_counter[0] = cutlass.Int32(0) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b, # Dense: cute.Tensor (N,K,L) | Discrete: cute.Pointer to int64[] + n: Int32, # Ignored for dense mode + k: Int32, # Ignored for dense mode + b_stride_size: cutlass.Int64, # Ignored for dense mode + b_major_mode: cutlass.Constexpr, # Ignored for dense mode + workspace_ptr, # Descriptor workspace, plus dynamic scheduler counter when enabled + c: cute.Tensor, + d: cute.Tensor, + padded_offsets: cute.Tensor, + alpha: cute.Tensor, + beta: cute.Tensor, + prob: cute.Tensor, + dprob: cute.Tensor, + linear_offset: Float32, + dbias_tensor: Optional[cute.Tensor], + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + ): + """Execute the GEMM. + + Dense mode: ``b`` is a 3-D cute.Tensor (N, K, L). + Discrete mode: ``b`` is a cute.Pointer to a device int64[] array of + per-expert base addresses; ``n``, ``k``, ``b_stride_size``, and + ``b_major_mode`` describe the uniform per-expert layout. + """ + # Setup static attributes before smem/grid/tma computation + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = a.element_type # B must match A dtype + self.c_dtype: Type[cutlass.Numeric] = c.element_type + self.d_dtype: Type[cutlass.Numeric] = d.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + else: + self.b_major_mode = b_major_mode + self.c_layout = utils.LayoutEnum.from_tensor(c) + self.d_layout = utils.LayoutEnum.from_tensor(d) + + # dBias configuration + self.generate_dbias = dbias_tensor is not None + self.dbias_cross_warp_reduce = self.generate_dbias # always cross-warp reduce + + # Check if input data types are compatible with MMA instruction + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + # ---- B setup (mode-dependent) ---- + b_from_call_arg = b + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + c1 = cutlass.Int32(1) + c0 = cutlass.Int64(0) + c1_64 = 1 + if cutlass.const_expr(b_major_mode == OperandMajorMode.K): + b_template_stride = (b_stride_size, c1_64, c0) + else: + b_template_stride = (c1_64, b_stride_size, c0) + b_template_layout = cute.make_layout((n, k, c1), stride=b_template_stride) + b_ptr_typed = cute.make_ptr(self.b_dtype, b.toint(), AddressSpace.gmem, assumed_align=16) + b = cute.make_tensor(b_ptr_typed, b_template_layout) + + # Compute grid size + m, n_d, l = cute.shape(d) + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # Setup TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + # Setup TMA load for C + c_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + self.tma_c_load_bytes = cute.size_in_bytes(self.c_dtype, c_smem_layout) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + c, + c_smem_layout, + self.epi_tile, + ) + + # Setup TMA store for D unless the tuning knob asks epilogue warps to + # write D directly from registers to GMEM. + tma_atom_d = None + tma_tensor_d = d + if cutlass.const_expr(not self.store_d_directly): + d_smem_layout = cute.slice_(self.d_smem_layout_staged, (None, None, 0)) + tma_atom_d, tma_tensor_d = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + d, + d_smem_layout, + self.epi_tile, + ) + + # Compute grid size using MoE scheduler + # dGLU output has shape (m, 2*N_half, 1), but scheduling is over (m, N_half) + # expert_shape = (expert_cnt, N_half, K) + n_half = n_d // 2 + sched_params = MoESchedulerParams( + scenario="2Dx3D", + expert_shape=(self.expert_cnt, n_half, cute.size(a.shape, mode=[1])), + cta_tile_shape_mnk=self.cta_tile_shape_mnk_d, + cluster_shape_mn=self.cluster_shape_mn, + use_dynamic_sched=self.use_dynamic_sched, + ) + self.sched_params, grid = compute_grid(sched_params, max_active_clusters, self.use_2cta_instrs) + + self.buffer_align_bytes = 1024 + + # Define shared storage for kernel + SchedulerStorage = MoEPersistentTileScheduler.make_storage_struct(self.num_tile_stage, self.use_dynamic_sched) + + sD_size = 0 if self.store_d_directly else cute.cosize(self.d_smem_layout_staged.outer) + + @cute.struct + class SharedStorage: + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + scheduler: SchedulerStorage + c_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage] + c_empty_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_c_stage] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.cosize(self.c_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + sD: cute.struct.Align[ + cute.struct.MemRange[ + self.d_dtype, + sD_size, + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sB: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + # dBias SMEM transpose buffer: (128, epi_tile_n*2) col-major FP32 + sDbias: cute.struct.Align[ + cute.struct.MemRange[ + cutlass.Float32, + 128 * self.epi_tile[1] * 2 if self.generate_dbias else 1, + ], + 128 if self.generate_dbias else 4, + ] + + self.shared_storage = SharedStorage + + # Initialize per-expert B TMA descriptors in workspace + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + _need_helper = cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE or self.use_dynamic_sched) + if cutlass.const_expr(_need_helper): + _helper_grid_x = self.expert_cnt if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else 1 + _helper_args = ( + b_from_call_arg if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cute.make_ptr(cutlass.Int64, 0, AddressSpace.gmem), + n if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int32(0), + k if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int32(0), + b_stride_size if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int64(0), + b_major_mode if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else self.b_major_mode, + workspace_ptr, + tiled_mma, + b_smem_layout, + self.cluster_layout_vmnk.shape, + ) + self.helper_kernel(*_helper_args).launch( + grid=(_helper_grid_x, 1, 1), + block=(1, 1, 1), + stream=stream, + min_blocks_per_mp=1, + ) + + # Launch the main kernel + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + tma_atom_d, + tma_tensor_d, + padded_offsets, + alpha, + beta, + prob, + dprob, + linear_offset, + dbias_tensor, + workspace_ptr, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.d_smem_layout_staged, + self.epi_tile, + self.sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + max_number_threads=[self.threads_per_cta, 1, 1], + smem=self.shared_storage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + return + + # ------------------------------------------------------------------ + # Internal: create extension based on weight_mode + # ------------------------------------------------------------------ + + @cute.jit + def _make_extension(self, workspace_ptr): + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + desc_workspace = TensormapWorkspace(workspace_ptr, ["b"]) + return DiscreteWeightGroupedGemmSchedExtension(tensormap_ctor=desc_workspace) + else: + return ContiguousGroupedGemmSchedExtension() + + @cute.jit + def dbias_reduction( + self, + d1_vec, + d2_vec, + warp_idx, + sDbias, + dbias_gmem_2d, + expert_idx, + n_base_d1, + n_base_d2, + dbias_n_total, + ) -> None: + """Merged dy1+dy2 dbias reduction via SMEM transpose.""" + epi_n = self.epi_tile[1] + lane_idx = cute.arch.lane_idx() + warp_local = warp_idx - self.epilog_warp_id[0] + + for n in cutlass.range(epi_n, unroll_full=True): + sDbias[(n, lane_idx, warp_local)] = d1_vec[n] + sDbias[(epi_n + n, lane_idx, warp_local)] = d2_vec[n] + + self.epilog_sync_barrier.arrive_and_wait() + + col_a = 2 * lane_idx if lane_idx < 16 else epi_n + 2 * (lane_idx - 16) + col_b = col_a + 1 + + copy_128bit_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=128) + warp_base_ptr = sDbias.iterator + warp_local * epi_n * 2 * 32 + swizzle_a = ((col_a >> 1) & 0x7) << 2 + swizzle_b = ((col_b >> 1) & 0x7) << 2 + + sum_a = cutlass.Float32(0.0) + sum_b = cutlass.Float32(0.0) + rDst_a = cute.make_rmem_tensor(cute.make_layout((4,)), cutlass.Float32) + rDst_b = cute.make_rmem_tensor(cute.make_layout((4,)), cutlass.Float32) + for g in cutlass.range(8, unroll_full=True): + m_base = g * 4 + sw_offset_a = col_a * 32 + (m_base ^ swizzle_a) + sSrc_a = cute.make_tensor(warp_base_ptr + sw_offset_a, cute.make_layout((4,))) + cute.copy_atom_call(copy_128bit_atom, sSrc_a, rDst_a) + + sw_offset_b = col_b * 32 + (m_base ^ swizzle_b) + sSrc_b = cute.make_tensor(warp_base_ptr + sw_offset_b, cute.make_layout((4,))) + cute.copy_atom_call(copy_128bit_atom, sSrc_b, rDst_b) + + for i in cutlass.range(4, unroll_full=True): + sum_a = sum_a + rDst_a[i] + sum_b = sum_b + rDst_b[i] + + n_offset = (n_base_d1 + 2 * lane_idx) if lane_idx < 16 else (n_base_d2 + 2 * (lane_idx - 16)) + + if cutlass.const_expr(self.dbias_cross_warp_reduce): + reduce_base = sDbias.iterator + copy_64bit_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=64) + + self.epilog_sync_barrier.arrive_and_wait() + rSrc_partial = cute.make_rmem_tensor(cute.make_layout((2,)), cutlass.Float32) + rSrc_partial[0] = sum_a + rSrc_partial[1] = sum_b + sDst_partial = cute.make_tensor(reduce_base + warp_local * 64 + lane_idx * 2, cute.make_layout((2,))) + cute.copy_atom_call(copy_64bit_atom, rSrc_partial, sDst_partial) + self.epilog_sync_barrier.arrive_and_wait() + + if warp_idx == self.epilog_warp_id[0]: + cta_sum_a = cutlass.Float32(0.0) + cta_sum_b = cutlass.Float32(0.0) + rDst_w = cute.make_rmem_tensor(cute.make_layout((2,)), cutlass.Float32) + for w in cutlass.range(self.num_epilog_warps): + sSrc_w = cute.make_tensor(reduce_base + w * 64 + lane_idx * 2, cute.make_layout((2,))) + cute.copy_atom_call(copy_64bit_atom, sSrc_w, rDst_w) + cta_sum_a = cta_sum_a + rDst_w[0] + cta_sum_b = cta_sum_b + rDst_w[1] + if n_offset < dbias_n_total: + gmem_ptr = dbias_gmem_2d[(expert_idx, n_offset, None)].iterator.llvm_ptr + atomic_add_bf16x2(gmem_ptr, cta_sum_a, cta_sum_b) + else: + if n_offset < dbias_n_total: + gmem_ptr = dbias_gmem_2d[(expert_idx, n_offset, None)].iterator.llvm_ptr + atomic_add_bf16x2(gmem_ptr, sum_a, sum_b) + + @cute.jit + def dswiglu( + self, + acc_vec: cute.Tensor, + ab1_vec_load: cute.Tensor, + ab2_vec_load: cute.Tensor, + mProb: cute.Tensor, + beta_val: Float32, + square_alpha: Float32, + dprob_swiglu: Optional[cute.Tensor] = None, + ): + LOG2_E = cutlass.Float32(1.4426950408889634) + if cutlass.const_expr(self.vectorized_f32): + d1_vec = cute.make_rmem_tensor(acc_vec.shape, cutlass.Float32) + d2_vec = cute.make_rmem_tensor(acc_vec.shape, cutlass.Float32) + for i in cutlass.range(0, cute.size(acc_vec), 2, unroll_full=True): + # Apply dGLU alpha/beta/prob scaling. + ( + acc_vec[i + 0], + acc_vec[i + 1], + ) = cute.arch.mul_packed_f32x2( + (acc_vec[i + 0], acc_vec[i + 1]), + (square_alpha, square_alpha), + rnd="rn", + ftz=False, + ) + ab1_vec_acc_type = cute.arch.mul_packed_f32x2( + ( + ab1_vec_load[i + 0].to(self.acc_dtype), + ab1_vec_load[i + 1].to(self.acc_dtype), + ), + (beta_val, beta_val), + rnd="rn", + ftz=False, + ) + ab2_vec_acc_type = cute.arch.mul_packed_f32x2( + ( + ab2_vec_load[i + 0].to(self.acc_dtype), + ab2_vec_load[i + 1].to(self.acc_dtype), + ), + (beta_val, beta_val), + rnd="rn", + ftz=False, + ) + sig_rcp_0, sig_rcp_1 = cute.arch.mul_packed_f32x2( + (ab1_vec_acc_type), + (-LOG2_E, -LOG2_E), + rnd="rn", + ftz=False, + ) + sig_rcp_0, sig_rcp_1 = cute.arch.add_packed_f32x2( + ( + cute.math.exp2(sig_rcp_0, fastmath=True), + cute.math.exp2(sig_rcp_1, fastmath=True), + ), + (1.0, 1.0), + rnd="rn", + ftz=False, + ) + sig = ( + cute.arch.rcp_approx(sig_rcp_0), + cute.arch.rcp_approx(sig_rcp_1), + ) + swish = cute.arch.mul_packed_f32x2( + ab1_vec_acc_type, + sig, + rnd="rn", + ftz=False, + ) + # calculate dprob + if cutlass.const_expr(self.generate_dprob): + ( + dprob_swiglu[i + 0], + dprob_swiglu[i + 1], + ) = cute.arch.mul_packed_f32x2( + (ab2_vec_acc_type[0], ab2_vec_acc_type[1]), + swish, + ) + ( + dprob_swiglu[i + 0], + dprob_swiglu[i + 1], + ) = cute.arch.mul_packed_f32x2( + (dprob_swiglu[i + 0], dprob_swiglu[i + 1]), + (acc_vec[i + 0], acc_vec[i + 1]), + ) + # calculate dswiglu + acc_vec_prob = cute.arch.mul_packed_f32x2( + (acc_vec[i + 0], acc_vec[i + 1]), + (mProb, mProb), + ) + # calculate d2_vec + ( + d2_vec[i + 0], + d2_vec[i + 1], + ) = cute.arch.mul_packed_f32x2( + (acc_vec_prob[0], acc_vec_prob[1]), + swish, + rnd="rn", + ftz=False, + ) + # calculate d1_vec + ( + d1_vec[i + 0], + d1_vec[i + 1], + ) = cute.arch.mul_packed_f32x2( + (acc_vec_prob[0], acc_vec_prob[1]), + (ab2_vec_acc_type[0], ab2_vec_acc_type[1]), + rnd="rn", + ftz=False, + ) + ( + d1_vec[i + 0], + d1_vec[i + 1], + ) = cute.arch.mul_packed_f32x2( + (d1_vec[i + 0], d1_vec[i + 1]), + sig, + rnd="rn", + ftz=False, + ) + one_minus_sig = cute.arch.add_packed_f32x2( + (1.0, 1.0), + (-sig[0], -sig[1]), + rnd="rn", + ftz=False, + ) + dsig = cute.arch.mul_packed_f32x2( + ab1_vec_acc_type, + one_minus_sig, + rnd="rn", + ftz=False, + ) + dsig_add_1 = cute.arch.add_packed_f32x2( + (dsig[0], dsig[1]), + (1.0, 1.0), + rnd="rn", + ftz=False, + ) + ( + d1_vec[i + 0], + d1_vec[i + 1], + ) = cute.arch.mul_packed_f32x2( + (d1_vec[i + 0], d1_vec[i + 1]), + dsig_add_1, + rnd="rn", + ftz=False, + ) + d1_vec = d1_vec.load() + d2_vec = d2_vec.load() + if cutlass.const_expr(self.generate_dprob): + dprob_swiglu = dprob_swiglu.load() + return d1_vec, d2_vec, dprob_swiglu + else: + acc_vec = acc_vec.load() + ab1_vec_load = ab1_vec_load.load() + ab2_vec_load = ab2_vec_load.load() + + acc_vec = acc_vec * square_alpha # apply scale for A*B + ab1_vec_load = ab1_vec_load * beta_val # apply scale for C + ab2_vec_load = ab2_vec_load * beta_val # apply scale for C + + sig_rcp = (1 + cute.math.exp(-1 * ab1_vec_load, True)).to(self.acc_dtype) + res = cute.make_rmem_tensor(sig_rcp.shape, cutlass.Float32) + res.store(sig_rcp) + # let every res[?] be cute.arch.rcp_approx(res[?]) + [res.__setitem__(i, cute.arch.rcp_approx(res[i])) for i in range(cute.size(res.shape))] + sig = res.load() + swish = ab1_vec_load * sig + + # calculate dprob + if cutlass.const_expr(self.generate_dprob): + dprob_swiglu = ab2_vec_load * swish + dprob_swiglu = acc_vec * dprob_swiglu + + # calculate dswiglu + d1_vec = acc_vec * mProb * ab2_vec_load * sig * (1 + ab1_vec_load * (1 - sig)) + d2_vec = acc_vec * mProb * swish + return d1_vec, d2_vec, dprob_swiglu + + @cute.jit + def dgeglu( + self, + acc_vec: cute.Tensor, + x1_vec_load: cute.Tensor, + x2_vec_load: cute.Tensor, + mProb: cute.Tensor, + beta_val: Float32, + square_alpha: Float32, + linear_offset: Float32, + dprob_swiglu: Optional[cute.Tensor] = None, + ): + geglu_max_value = cutlass.Float32(7.0) + geglu_min_value = cutlass.Float32(-7.0) + fmul2 = partial(cute.arch.mul_packed_f32x2, rnd="rn", ftz=False) + fadd2 = partial(cute.arch.add_packed_f32x2, rnd="rn", ftz=False) + scale_1702 = (1.702, 1.702) + ones2 = (1.0, 1.0) + mprob2 = (mProb, mProb) + beta2 = (beta_val, beta_val) + square_alpha2 = (square_alpha, square_alpha) + linear_offset2 = (linear_offset, linear_offset) + + if cutlass.const_expr(self.vectorized_f32): + dx1_vec = cute.make_rmem_tensor(acc_vec.shape, cutlass.Float32) + dx2_vec = cute.make_rmem_tensor(acc_vec.shape, cutlass.Float32) + for i in cutlass.range(0, cute.size(acc_vec), 2, unroll_full=True): + acc = fmul2((acc_vec[i], acc_vec[i + 1]), square_alpha2) + x1_0, x1_1 = fmul2( + ( + x1_vec_load[i].to(self.acc_dtype), + x1_vec_load[i + 1].to(self.acc_dtype), + ), + beta2, + ) + x2_0, x2_1 = fmul2( + ( + x2_vec_load[i].to(self.acc_dtype), + x2_vec_load[i + 1].to(self.acc_dtype), + ), + beta2, + ) + + y1_0 = fmin(x1_0, geglu_max_value) + y1_1 = fmin(x1_1, geglu_max_value) + y2_0 = fmin(x2_0, geglu_max_value) + y2_1 = fmin(x2_1, geglu_max_value) + y2_0 = fmax(y2_0, geglu_min_value) + y2_1 = fmax(y2_1, geglu_min_value) + + y1 = (y1_0, y1_1) + y2 = (y2_0, y2_1) + + # y1 = 1.702 * x1 + y1_scaled = fmul2(y1, scale_1702) + + sigmoid_out_0 = sigmoid_f32(y1_scaled[0], fastmath=True) + sigmoid_out_1 = sigmoid_f32(y1_scaled[1], fastmath=True) + + # g * sigmoid_out + acc_mul_sigmoid_out = fmul2(acc, (sigmoid_out_0, sigmoid_out_1)) + acc_mul_sigmoid_prob = fmul2(acc_mul_sigmoid_out, mprob2) + + # y1 = 1 + 1.702 * y1 * (1 - sigmoid_out) + one_minus_sigmoid_0, one_minus_sigmoid_1 = fadd2(ones2, (-sigmoid_out_0, -sigmoid_out_1)) + y1_scaled = fadd2( + fmul2(y1_scaled, (one_minus_sigmoid_0, one_minus_sigmoid_1)), + ones2, + ) + + # y2 + linear_offset + y2_with_linear_offset_0, y2_with_linear_offset_1 = fadd2(y2, linear_offset2) + + # dy1 = g * sigmoid_out * (y2 + linear_offset) + dy1_pre_0, dy1_pre_1 = fmul2( + (y2_with_linear_offset_0, y2_with_linear_offset_1), + acc_mul_sigmoid_out, + ) + # dy1 = g * sigmoid_out * (y2 + linear_offset) * (1 + 1.702 * y1 * (1 - sigmoid_out)) * mProb + dy1_0, dy1_1 = fmul2((dy1_pre_0, dy1_pre_1), y1_scaled) + dy1_0, dy1_1 = fmul2((dy1_0, dy1_1), mprob2) + + x1_filter_0 = y1_0 if x1_0 <= geglu_max_value else cutlass.Float32(0.0) + x1_filter_1 = y1_1 if x1_1 <= geglu_max_value else cutlass.Float32(0.0) + + dx1_vec[i], dx1_vec[i + 1] = fmul2((dy1_0, dy1_1), (cutlass.Float32(x1_filter_0), cutlass.Float32(x1_filter_1))) + + # dy2 = g * y1 * sigmoid_out * mProb + dy2_0, dy2_1 = fmul2(y1, acc_mul_sigmoid_prob) + x2_filter_0 = x2_0 if x2_0 <= geglu_max_value else cutlass.Float32(0.0) + x2_filter_1 = x2_1 if x2_1 <= geglu_max_value else cutlass.Float32(0.0) + x2_filter_0 = y2_0 if x2_filter_0 >= geglu_min_value else cutlass.Float32(0.0) + x2_filter_1 = y2_1 if x2_filter_1 >= geglu_min_value else cutlass.Float32(0.0) + dx2_vec[i], dx2_vec[i + 1] = fmul2((dy2_0, dy2_1), (cutlass.Float32(x2_filter_0), cutlass.Float32(x2_filter_1))) + + if cutlass.const_expr(self.generate_dprob): + prob_grad, prob_grad_1 = fmul2( + (dy1_pre_0, dy1_pre_1), + y1, + ) + dprob_swiglu[i] = prob_grad + dprob_swiglu[i + 1] = prob_grad_1 + dx1_vec = dx1_vec.load() + dx2_vec = dx2_vec.load() + if cutlass.const_expr(self.generate_dprob): + dprob_swiglu = dprob_swiglu.load() + return dx1_vec, dx2_vec, dprob_swiglu + else: + element_count = cute.size(x1_vec_load) + acc_vec = acc_vec.load() * square_alpha + x1_vec_load = x1_vec_load.load().to(cutlass.Float32) * beta_val + x2_vec_load = x2_vec_load.load().to(cutlass.Float32) * beta_val + dx1_vec = cute.make_rmem_tensor(acc_vec.shape, cutlass.Float32) + dx2_vec = cute.make_rmem_tensor(acc_vec.shape, cutlass.Float32) + + # y1 = clamp(x1, max=7.0); y2 = clamp(x2, min=-7.0, max=7.0) + for i in cutlass.range_constexpr(element_count): + fc2_dgrad = acc_vec[i] + g = fc2_dgrad * mProb + y1 = min(x1_vec_load[i], 7.0) + y2 = min(x2_vec_load[i], 7.0) + y2 = max(y2, -7.0) + + sigmoid_out = sigmoid_f32(y1 * 1.702, fastmath=True) + + dy1 = g * sigmoid_out * (1 + 1.702 * y1 * (1 - sigmoid_out)) * (y2 + linear_offset) + dy2 = g * y1 * sigmoid_out + + x1_filter = x1_vec_load[i] if x1_vec_load[i] <= 7.0 else 0.0 + x2_filter = x2_vec_load[i] if x2_vec_load[i] <= 7.0 else 0.0 + x2_filter = x2_filter if x2_filter >= -7.0 else 0.0 + + dx1_vec[i] = x1_filter * dy1 + dx2_vec[i] = x2_filter * dy2 + + if cutlass.const_expr(self.generate_dprob): + prob_grad = y1 * sigmoid_out * (y2 + linear_offset) * fc2_dgrad + dprob_swiglu[i] = prob_grad + + return dx1_vec.load(), dx2_vec.load(), dprob_swiglu.load() + + @cute.jit + def stg_256(self, ptr, vec8_f32, *, loc=None, ip=None): + """Store 256 bits to global memory with L1 no-allocate.""" + dst = ptr.ir_value(loc=loc, ip=ip) if hasattr(ptr, "ir_value") else ptr + src = vec8_f32.ir_value(loc=loc, ip=ip) if hasattr(vec8_f32, "ir_value") else vec8_f32 + llvm.inline_asm( + T.i32(), + [ + dst, + vector.extract(src, [], [0], loc=loc, ip=ip), + vector.extract(src, [], [1], loc=loc, ip=ip), + vector.extract(src, [], [2], loc=loc, ip=ip), + vector.extract(src, [], [3], loc=loc, ip=ip), + vector.extract(src, [], [4], loc=loc, ip=ip), + vector.extract(src, [], [5], loc=loc, ip=ip), + vector.extract(src, [], [6], loc=loc, ip=ip), + vector.extract(src, [], [7], loc=loc, ip=ip), + ], + "st.global.L1::no_allocate.v8.f32 [$1], {$2, $3, $4, $5, $6, $7, $8, $9}; mov.u32 $0, 0;", + "=r,l,f,f,f,f,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + @cute.jit + def store_global_memory_256b(self, dst: cute.Tensor, src: cute.Tensor): + vec_shape = cute.make_layout(8) + dst_f32 = cute.flatten(cute.recast_tensor(dst, cutlass.Float32)) + src_f32 = cute.flatten(cute.recast_tensor(src, cutlass.Float32)) + dst_vf32x8 = cute.logical_divide(dst_f32, vec_shape) + src_vf32x8 = cute.logical_divide(src_f32, vec_shape) + for ei in cutlass.range_constexpr(dst_vf32x8.shape[1]): + self.stg_256( + dst_vf32x8[None, ei].iterator.llvm_ptr, + src_vf32x8[None, ei].load(), + ) + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + tma_atom_d: Optional[cute.CopyAtom], + mD_mnl: cute.Tensor, + padded_offsets: cute.Tensor, + alpha: cute.Tensor, + beta: cute.Tensor, + prob: cute.Tensor, + dprob: cute.Tensor, + linear_offset: Float32, + mDbias_tensor: Optional[cute.Tensor], + workspace_ptr, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + d_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + epi_tile: cute.Tile, + sched_params: MoESchedulerParams, + ): + """ + GPU device kernel performing the Persistent batched GEMM computation. + """ + tidx, _, _ = cute.arch.thread_idx() + warp_idx = tidx // 32 + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + total_tokens = padded_offsets[self.expert_cnt - 1] + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_c) + if cutlass.const_expr(not self.store_d_directly): + cpasync.prefetch_descriptor(tma_atom_d) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + sched_storage = storage.scheduler + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_warp_id) * (2 if use_2cta_instrs else 1) + acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Load C pipeline + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + c_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.c_full_mbar_ptr.data_ptr(), + num_stages=self.num_c_stage, + producer_group=c_producer_group, + consumer_group=c_consumer_group, + tx_count=self.tma_c_load_bytes, + defer_sync=True, + ) + + # Initialize tile info pipeline (barrier) and states + tile_info_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * 1, + ) + tile_info_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_wo_sched, + ) + tile_info_pipeline = pipeline.PipelineAsync.create( + barrier_storage=sched_storage.tile_info_mbar.data_ptr(), + num_stages=self.num_tile_stage, + producer_group=tile_info_pipeline_producer_group, + consumer_group=tile_info_pipeline_consumer_group, + ) + + scheduler = MoEPersistentTileScheduler.create( + sched_params, + padded_offsets, + cute.arch.block_idx(), + cute.arch.grid_dim(), + counter_ptr=self._get_sched_counter_ptr(workspace_ptr), + sched_storage=sched_storage, + ) + scheduler.internal_init() + + # dBias SMEM setup + if cutlass.const_expr(self.generate_dbias): + sDbias = storage.sDbias.get_tensor( + cute.make_layout( + (self.epi_tile[1] * 2, 32, len(self.epilog_warp_id)), + stride=(32, 1, self.epi_tile[1] * 2 * 32), + ) + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, + ) + + # Cluster arrive after barrier init + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_arrive_relaxed() + + # + # Setup smem tensor A/B/C/D + # + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner) + sD = None + if cutlass.const_expr(not self.store_d_directly): + sD = storage.sD.get_tensor(d_smem_layout_staged.outer, swizzle=d_smem_layout_staged.inner) + # (MMA, MMA_M, MMA_K, STAGE) + sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner) + # (MMA, MMA_N, MMA_K, STAGE) + sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner) + # (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + info_layout = cute.make_layout((4, self.num_tile_stage), stride=(1, 4)) + sInfo = sched_storage.sInfo.get_tensor(info_layout) + + # + # Compute multicast mask for A/B buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2) + b_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/D + # (SMEM/TMEM partitions stay global - they don't depend on per-expert tensors) + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + # + # Cluster wait before tensor memory alloc + # + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_wait() + else: + self.cta_sync_barrier.arrive_and_wait() + + if total_tokens <= 0: + cute.arch.nvvm.exit() + k_tile_cnt = cute.ceil_div(cute.size(mB_nkl, mode=[1]), self.mma_tiler[2]) + + # + # Specialized Schedule warp (MoE Persistent Tile Scheduler) + # + if warp_idx == self.sched_warp_id: + work_tile_info = scheduler.initial_work_tile_info() + + tile_info_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_tile_stage) + + while work_tile_info.is_valid_tile: + # sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = work_tile_info.expert_idx + sInfo[(1, tile_info_producer_state.index)] = work_tile_info.tile_m_idx + sInfo[(2, tile_info_producer_state.index)] = work_tile_info.tile_n_idx + sInfo[(3, tile_info_producer_state.index)] = work_tile_info.k_tile_cnt + cute.arch.fence_proxy("async.shared", space="cta") + + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + + work_tile_info = scheduler.advance_to_next_work() + + # Send invalid tile signal: expert_idx = -1 + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = cutlass.Int32(-1) + sInfo[(1, tile_info_producer_state.index)] = cutlass.Int32(0) + sInfo[(2, tile_info_producer_state.index)] = cutlass.Int32(0) + sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + tile_info_pipeline.producer_tail(tile_info_producer_state) + + # + # Specialized TMA load warp + # + if warp_idx == self.tma_warp_id: + ext = self._make_extension(workspace_ptr) + + ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage) + + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + # Get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + # assert(k_tile_cnt == work_tile_info.k_tile_cnt) + ext.update_expert_info(padded_offsets, work_tile_info.expert_idx) + + # Get per-expert real tensors + TMA desc ptrs via extension + real_a, _ = ext.get_gmem_tensor("a", mA_mkl, padded_offsets, work_tile_info) + real_b, desc_ptr_b = ext.get_gmem_tensor("b", mB_nkl, padded_offsets, work_tile_info) + + # local_tile on per-expert tensors + gA_mkl = cute.local_tile(real_a, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) + gB_nkl = cute.local_tile(real_b, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) + + # MMA partition + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + + # TMA partition A + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA partition B + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + # Slice to per mma tile index (L=0 since domain already offset'd) + mma_tile_coord_m = work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape) + mma_tile_coord_n = work_tile_info.tile_n_idx + tAgA_slice = tAgA[(None, mma_tile_coord_m, None, 0)] + tBgB_slice = tBgB[(None, mma_tile_coord_n, None, 0)] + + # Peek (try_wait) AB buffer empty + peek_ab_empty_status = cutlass.Boolean(1) + if k_tile_cnt > 0: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + tAgA_k = tAgA_slice[(None, k_tile)] + tBgB_k = tBgB_slice[(None, k_tile)] + tAsA_pipe = tAsA[(None, ab_producer_state.index)] + tBsB_pipe = tBsB[(None, ab_producer_state.index)] + + tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state) + + # Conditionally wait for AB buffer empty + ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) + ab_producer_state_next = ab_producer_state.clone() + ab_producer_state_next.advance() + if k_tile < k_tile_cnt - 1: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state_next) + + # TMA load A (contiguous, global desc) + cute.copy( + tma_atom_a, + tAgA_k, + tAsA_pipe, + tma_bar_ptr=tma_bar, + mcast_mask=a_full_mcast_mask, + ) + # TMA load B (discrete, per-expert desc from workspace) + cute.copy( + tma_atom_b, + tBgB_k, + tBsB_pipe, + tma_bar_ptr=tma_bar, + mcast_mask=b_full_mcast_mask, + tma_desc_ptr=desc_ptr_b, + ) + + # Peek (try_wait) AB buffer empty for next k_tile + ab_producer_state = ab_producer_state_next + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # + # Wait A/B buffer empty + # + ab_pipeline.producer_tail(ab_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Bar sync for retrieve tensor memory ptr from shared mem + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + ab_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_ab_stage) + acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) + + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + # Get the first tile info (sInfo format: expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + # assert(k_tile_cnt == tile_info[3]) + + # Peek (try_wait) AB buffer full for k_tile = 0 + peek_ab_full_status = cutlass.Boolean(1) + if k_tile_cnt > 0 and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + # sInfo: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + mma_tile_coord_mnl = ( + tile_info[1] // cute.size(tiled_mma.thr_id.shape), + tile_info[2], + cutlass.Int32(0), + ) + + tCtAcc = tCtAcc_base[(None, None, None, acc_producer_state.index)] + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for AB buffer full + ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + ab_consumer_state_next = ab_consumer_state.clone() + ab_consumer_state_next.advance() + if k_tile < k_tile_cnt - 1: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state_next) + + # tCtAcc += tCrA * tCrB + num_kblocks = cute.size(tCrA, mode=[2]) + + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord = ( + None, + None, + kblock_idx, + ab_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord], + tCrB[kblock_coord], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + ab_pipeline.consumer_release(ab_consumer_state) + ab_consumer_state = ab_consumer_state_next + acc_pipeline.producer_commit(acc_producer_state) + + # Peek (try_wait) Acc buffer empty for k_tile = k_tile + 1 + acc_producer_state.advance() + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Partition for epilogue (SMEM/TMEM/register - invariant across experts) + # + epi_tidx = tidx + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = self.epilog_tmem_copy_and_partition(epi_tidx, tCtAcc_base, epi_tile, use_2cta_instrs) + + tTR_rC1 = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tTR_rC2 = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_s2r, tRS_rC1, tRS_rC2, tRS_sC = self.epilog_smem_copy_and_partition_load(tiled_copy_t2r, tTR_rC1, tTR_rC2, epi_tidx, sC) + + tTR_rD1 = cute.make_rmem_tensor(tTR_rAcc.shape, self.d_dtype) + tTR_rD2 = cute.make_rmem_tensor(tTR_rAcc.shape, self.d_dtype) + tiled_copy_r2s, tRS_rD1, tRS_rD2, tRS_sD = self.epilog_smem_copy_and_partition_store(tiled_copy_t2r, tTR_rD1, tTR_rD2, epi_tidx, sD) + + # Extension for per-expert domain conversion in epilogue + epi_ext = self._make_extension(workspace_ptr) + + acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) + + # Load C pipeline + c_pipeline_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_c_stage) + + # Threads/warps participating in tma store pipeline + d_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + d_pipeline = None + if cutlass.const_expr(not self.store_d_directly): + num_d_stages = self.num_d_stage // 2 + d_pipeline = pipeline.PipelineTmaStore.create( + num_stages=num_d_stages, + producer_group=d_producer_group, + ) + + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + # Get the first tile info (sInfo format: expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + num_prev_subtiles = cutlass.Int32(0) + while is_valid_tile: + # sInfo: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + epi_work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + expert_idx = epi_work_tile_info.expert_idx + # N is doubled for dGLU dual output + mma_tile_coord_mnl = ( + epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape), + epi_work_tile_info.tile_n_idx * 2, + cutlass.Int32(0), + ) + + # + # Get alpha/beta for current expert + # + alpha_val = alpha[expert_idx] + beta_val = beta[expert_idx] + epi_ext.update_expert_info(padded_offsets, expert_idx) + + # + # Per-expert gmem tensor setup via extension + # + real_d, _ = epi_ext.get_gmem_tensor("d", mD_mnl, padded_offsets, epi_work_tile_info) + thr_mma_epi = tiled_mma.get_slice(mma_tile_coord_v) + + if cutlass.const_expr(not self.store_d_directly): + gD_mnl_loop = cute.local_tile(real_d, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None)) + tCgD_loop = thr_mma_epi.partition_C(gD_mnl_loop) + + bSG_sD, bSG_gD_partitioned = self.epilog_gmem_copy_and_partition(epi_tidx, tma_atom_d, tCgD_loop, epi_tile, sD) + bSG_gD = bSG_gD_partitioned[(None, None, None, mma_tile_coord_mnl[0], mma_tile_coord_mnl[1], 0)] + bSG_gD = cute.group_modes(bSG_gD, 1, cute.rank(bSG_gD)) + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] + + # + # Get PROB (per-expert local M position) + # + real_prob, _ = epi_ext.get_gmem_tensor("prob", prob, padded_offsets, epi_work_tile_info) + mPosition = ( + (epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape)) * self.mma_tiler[0] + + mma_tile_coord_v * (self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape)) + + tidx + ) + mProb = real_prob[mPosition, 0, 0] + if cutlass.const_expr(self.generate_dprob): + dProbVal = cutlass.Float32(0.0) + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + if cutlass.const_expr(self.epilogue_prefetch_more): + tTR_rAcc_0 = cute.make_rmem_tensor(tTR_rAcc.shape, cutlass.Float32) + tTR_rAcc_1 = cute.make_rmem_tensor(tTR_rAcc.shape, cutlass.Float32) + for subtile_idx in cutlass.range(0, subtile_cnt, 1, unroll=1): + real_subtile_idx = subtile_idx + real_subtile_idx_next = subtile_idx + 1 + # + # Load accumulator from tensor memory buffer to register + # + # Don't ask why, AST is shit tracking the constexpr values to loop args. + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.d_layout, + self.d_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + if cutlass.const_expr(self.epilogue_prefetch_more): + tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] + tTR_tAcc_mn_next = tTR_tAcc[(None, None, None, real_subtile_idx_next)] + if subtile_idx % 2 == 0: + cute.copy(copy_atom_t2r, tTR_tAcc_mn, tTR_rAcc_0) + cute.copy(copy_atom_t2r, tTR_tAcc_mn_next, tTR_rAcc_1) + tTR_rAcc = tTR_rAcc_0 + else: + tTR_rAcc = tTR_rAcc_1 + else: + tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] + cute.copy(copy_atom_t2r, tTR_tAcc_mn, tTR_rAcc) + + # Wait for C1/C2 load to complete + c_pipeline.consumer_wait(c_pipeline_consumer_state) + cute.copy( + tiled_copy_s2r, + tRS_sC[(None, None, None, c_pipeline_consumer_state.index)], + tRS_rC1, + ) + cute.arch.fence_proxy("async.shared", space="cta") + c_pipeline.consumer_release(c_pipeline_consumer_state) + c_pipeline_consumer_state.advance() + c_pipeline.consumer_wait(c_pipeline_consumer_state) + cute.copy( + tiled_copy_s2r, + tRS_sC[(None, None, None, c_pipeline_consumer_state.index)], + tRS_rC2, + ) + cute.arch.fence_proxy("async.shared", space="cta") + c_pipeline.consumer_release(c_pipeline_consumer_state) + c_pipeline_consumer_state.advance() + + acc_vec = tiled_copy_r2s.retile(tTR_rAcc) + ab1_vec_load = tiled_copy_r2s.retile(tRS_rC1) + ab2_vec_load = tiled_copy_r2s.retile(tRS_rC2) + if cutlass.const_expr(self.generate_dprob): + dprob_swiglu = cute.make_rmem_tensor(acc_vec.shape, cutlass.Float32) + else: + dprob_swiglu = None + + # + # Apply alpha, act, and prob + # + square_alpha = alpha_val * alpha_val + if cutlass.const_expr(self.act_func == "dswiglu"): + d1_vec, d2_vec, dprob_swiglu = self.dswiglu(acc_vec, ab1_vec_load, ab2_vec_load, mProb, beta_val, square_alpha, dprob_swiglu) + elif cutlass.const_expr(self.act_func == "dgeglu"): + d1_vec, d2_vec, dprob_swiglu = self.dgeglu( + acc_vec, ab1_vec_load, ab2_vec_load, mProb, beta_val, square_alpha, linear_offset, dprob_swiglu + ) + + if cutlass.const_expr(self.generate_dprob): + # dprob sum reduction + if cutlass.const_expr(self.vectorized_f32): + dprob_pair_0 = cutlass.Float32(0.0) + dprob_pair_1 = cutlass.Float32(0.0) + for j in cutlass.range(0, cute.size(dprob_swiglu.shape), 2, unroll_full=True): + ( + dprob_pair_0, + dprob_pair_1, + ) = cute.arch.add_packed_f32x2( + (dprob_pair_0, dprob_pair_1), + (dprob_swiglu[j], dprob_swiglu[j + 1]), + rnd="rn", + ftz=False, + ) + dProbVal += dprob_pair_0 + dprob_pair_1 + else: + dProbVal += dprob_swiglu.reduce( + cute.ReductionOp.ADD, + cutlass.Float32(0.0), + 0, + ) + + # + # Generate dBias + # + if cutlass.const_expr(self.generate_dbias): + n_base_d1 = epi_work_tile_info.tile_n_idx * (self.mma_tiler[1] * 2) + (2 * real_subtile_idx + 0) * self.epi_tile[1] + n_base_d2 = epi_work_tile_info.tile_n_idx * (self.mma_tiler[1] * 2) + (2 * real_subtile_idx + 1) * self.epi_tile[1] + dbias_n_total = cute.size(mDbias_tensor, mode=[1]) + self.dbias_reduction( + d1_vec, + d2_vec, + warp_idx, + sDbias, + mDbias_tensor, + expert_idx, + n_base_d1, + n_base_d2, + dbias_n_total, + ) + + # + # Convert to D type + # + tRS_rD1.store(d1_vec.to(self.d_dtype)) + tRS_rD2.store(d2_vec.to(self.d_dtype)) + + # + # Store D + # + if cutlass.const_expr(self.store_d_directly): + self.epilog_sync_barrier.arrive_and_wait() + d_idx_mn = ( + epi_work_tile_info.tile_m_idx, + epi_work_tile_info.tile_n_idx, + ) + d_epilogue_subtile = ( + cute.make_layout(128), + cute.make_layout(self.mma_tiler[1] * 2), + ) + gD_sub_loop = cute.local_tile(real_d, d_epilogue_subtile, (None, None, None)) + thr_copy_t2r = tiled_copy_t2r.get_slice(epi_tidx) + tCgD_mnl_loop = thr_copy_t2r.partition_D(gD_sub_loop) + tCgD_mnl_loop = cute.filter_zeros(tCgD_mnl_loop) + tCgD1 = tCgD_mnl_loop[ + ( + None, + 0, # T2R_M + 2 * real_subtile_idx + 0, # T2R_N + *d_idx_mn, # RestM/N + 0, # RestL + ) + ] + tCgD2 = tCgD_mnl_loop[ + ( + None, + 0, # T2R_M + 2 * real_subtile_idx + 1, # T2R_N + *d_idx_mn, # RestM/N + 0, # RestL + ) + ] + d_n_total = cute.size(real_d, mode=[1]) + n_base_d1 = epi_work_tile_info.tile_n_idx * (self.mma_tiler[1] * 2) + (2 * real_subtile_idx + 0) * self.epi_tile[1] + n_base_d2 = epi_work_tile_info.tile_n_idx * (self.mma_tiler[1] * 2) + (2 * real_subtile_idx + 1) * self.epi_tile[1] + if n_base_d1 < d_n_total: + self.store_global_memory_256b(tCgD1, tRS_rD1) + if n_base_d2 < d_n_total: + self.store_global_memory_256b(tCgD2, tRS_rD2) + else: + if warp_idx == self.epilog_warp_id[0]: + d_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + d1_buffer = num_prev_subtiles % self.num_d_stage + num_prev_subtiles = num_prev_subtiles + 1 + cute.copy( + tiled_copy_r2s, + tRS_rD1, + tRS_sD[(None, None, None, d1_buffer)], + ) + d2_buffer = num_prev_subtiles % self.num_d_stage + num_prev_subtiles = num_prev_subtiles + 1 + cute.copy( + tiled_copy_r2s, + tRS_rD2, + tRS_sD[(None, None, None, d2_buffer)], + ) + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier.arrive_and_wait() + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_d, + bSG_sD[(None, d1_buffer)], + bSG_gD[(None, 2 * real_subtile_idx + 0)], + ) + cute.copy( + tma_atom_d, + bSG_sD[(None, d2_buffer)], + bSG_gD[(None, 2 * real_subtile_idx + 1)], + ) + d_pipeline.producer_commit() + self.epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + if cutlass.const_expr(self.generate_dprob): + real_dprob, _ = epi_ext.get_gmem_tensor("dprob", dprob, padded_offsets, epi_work_tile_info) + _ = atomic_add_float32( + ptr=real_dprob[(mPosition, None, None)].iterator.llvm_ptr, + value=dProbVal, + ) + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + # + # Wait for D store complete + # + if cutlass.const_expr(not self.store_d_directly): + d_pipeline.producer_tail() + # + # Specialized epilog load warp (loads C from GMEM to SMEM via TMA) + # + if warp_idx == self.epilog_load_tma_id: + c_load_ext = self._make_extension(workspace_ptr) + + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + c_pipeline_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_c_stage) + while is_valid_tile: + c_work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + mma_tile_coord_mnl = ( + c_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape), + c_work_tile_info.tile_n_idx * 2, + cutlass.Int32(0), + ) + + # Per-expert C tensor via extension + real_c, _ = c_load_ext.get_gmem_tensor("c", mC_mnl, padded_offsets, c_work_tile_info) + gC_mnl_loop = cute.local_tile(real_c, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) + thr_mma_c_load = tiled_mma.get_slice(mma_tile_coord_v) + tCgC_loop = thr_mma_c_load.partition_C(gC_mnl_loop) + + bGS_sC, bGS_gC_partitioned = self.epilog_gmem_copy_and_partition(tidx, tma_atom_c, tCgC_loop, epi_tile, sC) + bGS_gC = bGS_gC_partitioned[(None, None, None, mma_tile_coord_mnl[0], mma_tile_coord_mnl[1], 0)] + bGS_gC = cute.group_modes(bGS_gC, 1, cute.rank(bGS_gC)) + subtile_cnt = cute.size(bGS_gC.shape, mode=[1]) + for subtile_idx in cutlass.range(subtile_cnt, unroll=1): + real_subtile_idx = subtile_idx + c_pipeline.producer_acquire(c_pipeline_producer_state) + cute.copy( + tma_atom_c, + bGS_gC[(None, 2 * real_subtile_idx + 0)], + bGS_sC[(None, c_pipeline_producer_state.index)], + tma_bar_ptr=c_pipeline.producer_get_barrier(c_pipeline_producer_state), + ) + c_pipeline_producer_state.advance() + c_pipeline.producer_acquire(c_pipeline_producer_state) + cute.copy( + tma_atom_c, + bGS_gC[(None, 2 * real_subtile_idx + 1)], + bGS_sC[(None, c_pipeline_producer_state.index)], + tma_bar_ptr=c_pipeline.producer_get_barrier(c_pipeline_producer_state), + ) + c_pipeline_producer_state.advance() + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Wait C buffer tail complete + # + c_pipeline.producer_tail(c_pipeline_producer_state) + + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[cutlass.Boolean, bool], + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) + and derive register array shape from the TMEM partition (no gmem dependency). + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param tAcc: The accumulator tensor to be copied and partitioned + :type tAcc: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param use_2cta_instrs: Whether use_2cta_instrs is enabled + :type use_2cta_instrs: bool + + :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) where: + - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + - tTR_tAcc: The partitioned accumulator tensor in TMEM + - tTR_rAcc: The register tensor for accumulator (shape derived from TMEM partition) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.d_layout, + self.d_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc[((None, None), 0, 0, None)], + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + tTR_rAcc = thr_copy_t2r.partition_D(tAcc_epi) + + # Derive register shape from TMEM partition (no gmem D needed) + per_subtile_shape = cute.coalesce(tTR_rAcc[(None, None, None, 0, 0, 0)].layout, target_profile=((1, 1), 1, 1)).shape + tTR_rAcc = cute.make_rmem_tensor(per_subtile_shape, self.acc_dtype) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilog_smem_copy_and_partition_load( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tTR_rC1: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for shared memory load, then use it to partition register array (destination) and shared memory (source). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing (tiled_copy_s2r, tSR_rC, tSR_sC) where: + - tiled_copy_s2r: The tiled copy operation for smem to register copy(s2r) + - tSR_rC: The partitioned tensor C (register destination) + - tSR_sC: The partitioned tensor C (smem source) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_s2r = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.c_dtype) + tiled_copy_s2r = cute.make_tiled_copy_D(copy_atom_s2r, tiled_copy_t2r) + # (S2R, S2R_M, S2R_N, PIPE_C) + thr_copy_s2r = tiled_copy_s2r.get_slice(tidx) + tSR_sC = thr_copy_s2r.partition_D(sC) + # (S2R, S2R_M, S2R_N) + tSR_rC = tiled_copy_s2r.retile(tTR_rC) + tSR_rC1 = tiled_copy_s2r.retile(tTR_rC1) + return tiled_copy_s2r, tSR_rC, tSR_rC1, tSR_sC + + def epilog_smem_copy_and_partition_store( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rD1: cute.Tensor, + tTR_rD2: cute.Tensor, + tidx: cutlass.Int32, + sD: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rD1: The partitioned accumulator tensor + :type tTR_rD1: cute.Tensor + :param tTR_rD2: The partitioned accumulator tensor + :type tTR_rD2: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sD: The shared memory tensor to be copied and partitioned + :type sD: cute.Tensor + + :return: A tuple containing (tiled_copy_r2s, tRS_rD, tRS_sD) where: + - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) + - tRS_rD: The partitioned tensor D (register source) + - tRS_sD: The partitioned tensor D (smem destination) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_r2s = sm100_utils.get_smem_store_op(self.d_layout, self.d_dtype, self.acc_dtype, tiled_copy_t2r) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sD = None + if cutlass.const_expr(sD is not None): + tRS_sD = thr_copy_r2s.partition_D(sD) + # (R2S, R2S_M, R2S_N) + tRS_rD1 = tiled_copy_r2s.retile(tTR_rD1) + tRS_rD2 = tiled_copy_r2s.retile(tTR_rD2) + return tiled_copy_r2s, tRS_rD1, tRS_rD2, tRS_sD + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: Union[cute.CopyAtom, cute.TiledCopy], + gD_mnl: cute.Tensor, + epi_tile: cute.Tile, + sD: cute.Tensor, + ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """Make tiledCopy for global memory store, then use it to: + - partition register array (source) and global memory (destination) for none TMA store version; + - partition shared memory (source) and global memory (destination) for TMA store version. + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version + :type atom: cute.CopyAtom or cute.TiledCopy + :param gD_mnl: The global tensor D + :type gD_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param sD: The shared memory tensor to be copied and partitioned + :type sD: cute.Tensor + + :return: A tuple containing : + - For TMA store: (tma_atom_d, bSG_sD, bSG_gD) where: + - tma_atom_d: The TMA copy atom + - bSG_sD: The partitioned shared memory tensor D + - bSG_gD: The partitioned global tensor D + :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor] + """ + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL) + gD_epi = cute.flat_divide(gD_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + tma_atom_d = atom + sD_for_tma_partition = cute.group_modes(sD, 0, 2) + gD_for_tma_partition = cute.group_modes(gD_epi, 0, 2) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, loopM, loopN, loopL) + bSG_sD, bSG_gD = cpasync.tma_partition( + tma_atom_d, + 0, + cute.make_layout(1), + sD_for_tma_partition, + gD_for_tma_partition, + ) + return bSG_sD, bSG_gD + + @staticmethod + def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + epi_tile: cute.Tile, + c_dtype: Type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + d_dtype: Type[cutlass.Numeric], + d_layout: utils.LayoutEnum, + num_smem_capacity: int, + occupancy: int, + store_d_directly: bool, + generate_dbias: bool = False, + ) -> Tuple[int, int, int]: + """Compute BF16-only pipeline stages. + + C stage count follows BF16 dGLU tuning; D staging is disabled when the + direct-store path is selected. + """ + num_acc_stage = 2 + num_c_stage = 2 + num_d_stage = 0 if store_d_directly else 2 + num_tile_stage = 2 + + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + 1, # a tmp 1 stage is provided + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + 1, # a tmp 1 stage is provided + ) + + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + + d_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + d_dtype, + d_layout, + epi_tile, + 1, + ) + + ab_bytes_per_stage = cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + mbar_helpers_bytes = 1024 + sinfo_bytes = 4 * 4 * num_tile_stage + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * num_c_stage + d_bytes_per_stage = cute.size_in_bytes(d_dtype, d_smem_layout_staged_one) + d_bytes = d_bytes_per_stage * num_d_stage + # dBias transpose buffer: (128, 64) column-major FP32 = 32 KB + dbias_bytes = 128 * 64 * cute.size_in_bytes(cutlass.Float32, cute.make_layout((1,))) if generate_dbias else 0 + epi_bytes = c_bytes + d_bytes + dbias_bytes + + num_ab_stage = (num_smem_capacity // occupancy - (mbar_helpers_bytes + epi_bytes + sinfo_bytes)) // ab_bytes_per_stage + + return num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, num_tile_stage diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu/_bf16_api.py b/python/cudnn/grouped_gemm/grouped_gemm_glu/_bf16_api.py new file mode 100644 index 000000000..5809fe328 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/_bf16_api.py @@ -0,0 +1,569 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Private BF16 API for the SM100 grouped GEMM GLU kernel.""" + +import os +import weakref +from typing import Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cuda.bindings import driver as cuda +from cutlass.cute.nvgpu import OperandMajorMode +from cutlass.cute.runtime import from_dlpack, make_fake_stream +import torch + +from cudnn.api_base import APIBase, TensorDesc +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor + +from ..moe_utils import MoEWeightMode +from .moe_grouped_gemm_glu_bias import MoEGroupedGemmGluBiasBf16Kernel + +_OUTPUT_DTYPES = [torch.bfloat16, torch.float16, torch.float32] + + +class GroupedGemmGluBf16API(APIBase): + """Descriptor-first lifecycle API for BF16 GLU forward.""" + + def __init__( + self, + sample_a: torch.Tensor, + sample_c: torch.Tensor, + sample_d: torch.Tensor, + sample_padded_offsets: torch.Tensor, + sample_alpha: torch.Tensor, + sample_b: Optional[torch.Tensor] = None, + sample_bias: Optional[torch.Tensor] = None, + sample_prob: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + b_shape: Optional[Tuple[int, ...]] = None, + b_dtype: Optional[torch.dtype] = None, + acc_dtype: torch.dtype = torch.float32, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + vector_f32: bool = False, + m_aligned: int = 256, + generate_c: bool = False, + act_func: str = "swiglu", + b_major: str = "k", + use_dynamic_sched: bool = False, + ) -> None: + super().__init__() + self._warn_experimental_api() + + if sample_b is not None and num_experts is None: + self.weight_mode = MoEWeightMode.DENSE + elif sample_b is None and num_experts is not None: + self.weight_mode = MoEWeightMode.DISCRETE + if b_shape is None or b_dtype is None: + raise ValueError("b_shape and b_dtype are required in discrete mode") + else: + raise ValueError("Provide sample_b for dense mode or (num_experts, b_shape, b_dtype) " "for discrete mode, but not both") + + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") + self.d_desc = self._make_tensor_desc(sample_d, name="sample_d") + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") + self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias") + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") + + self._sample_offset_values = self._copy_values_to_host(sample_padded_offsets) + self._sample_offsets_ref = weakref.ref(sample_padded_offsets) + self._sample_offsets_version = int(sample_padded_offsets._version) + self._sample_data_ptrs = { + name: tensor.data_ptr() + for name, tensor in ( + ("sample_a", sample_a), + ("sample_b", sample_b), + ("sample_c", sample_c), + ("sample_d", sample_d), + ("sample_padded_offsets", sample_padded_offsets), + ("sample_alpha", sample_alpha), + ("sample_bias", sample_bias), + ("sample_prob", sample_prob), + ) + if tensor is not None + } + + self.expert_cnt = self.b_desc.shape[2] if self.weight_mode == MoEWeightMode.DENSE and self.b_desc.ndim == 3 else int(num_experts or 0) + self.b_shape = tuple(b_shape) if b_shape is not None else None + self.b_dtype = b_dtype if b_dtype is not None else self.b_desc.dtype + self.acc_dtype = acc_dtype + self.mma_tiler_mn = tuple(mma_tiler_mn) + self.use_2cta_instrs = self.mma_tiler_mn[0] == 256 + self.cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if self.use_2cta_instrs else (1, 1))) + self.vector_f32 = vector_f32 + self.m_aligned = m_aligned + self.generate_c = generate_c + self.act_func = act_func + self.b_major = b_major + self.use_dynamic_sched = use_dynamic_sched + self._has_bias = self.bias_desc is not None + self._kernel = MoEGroupedGemmGluBiasBf16Kernel + self._workspace: Optional[torch.Tensor] = None + self._compile_b_ptrs: Optional[torch.Tensor] = None + self._validated_offsets: dict[int, tuple] = {} + self._validated_pointer_values: dict[int, tuple] = {} + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + @staticmethod + def _expect_shape(desc: TensorDesc, expected: Tuple[int, ...], name: str) -> None: + if desc.shape != expected: + raise ValueError(f"{name} shape mismatch: expected {expected}, got {desc.shape}") + + @staticmethod + def _expect_stride(desc: TensorDesc, expected: Tuple[int, ...], name: str) -> None: + if desc.stride != expected: + raise ValueError(f"{name} must use the source-compatible layout with stride " f"{expected}, got {desc.stride}") + + @staticmethod + def _expect_device(desc: TensorDesc, device: torch.device, name: str) -> None: + if desc.device != device: + raise ValueError(f"{name} must be on {device}, got {desc.device}") + + @staticmethod + def _copy_values_to_host(tensor: torch.Tensor) -> Tuple[int, ...]: + return tuple(int(value) for value in tensor.detach().cpu().tolist()) + + @staticmethod + def _is_validation_cached(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> bool: + cached = cache.get(id(tensor)) + return bool(cached and cached[0]() is tensor and cached[1] == int(tensor._version) and cached[2] == extra) + + @staticmethod + def _remember_validation(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> None: + key = id(tensor) + + def discard(_reference, *, cache=cache, key=key): + cache.pop(key, None) + + cache[key] = (weakref.ref(tensor, discard), int(tensor._version), extra) + + @staticmethod + def _validate_offset_sequence(values: Tuple[int, ...], *, expert_cnt: int, tensor_m: int) -> None: + if len(values) != expert_cnt: + raise ValueError(f"padded_offsets length mismatch: expected {expert_cnt}, got {len(values)}") + previous = 0 + for index, value in enumerate(values): + if value < previous: + raise ValueError("padded_offsets must be a non-decreasing cumulative sum; " f"index {index} is {value} after {previous}") + if value % MoEGroupedGemmGluBiasBf16Kernel.FIX_PAD_SIZE != 0: + raise ValueError(f"padded_offsets[{index}] must be 256-aligned, got {value}") + previous = value + if not values or values[-1] <= 0 or values[-1] > tensor_m: + raise ValueError(f"padded_offsets last value must be in [1, {tensor_m}], got " f"{values[-1] if values else None}") + + def _validate_offsets_once(self, offsets: torch.Tensor, *, tensor_m: int) -> None: + extra = (self.expert_cnt, tensor_m) + if self._is_validation_cached(self._validated_offsets, offsets, extra): + return + values = self._copy_values_to_host(offsets) + self._validate_offset_sequence(values, expert_cnt=self.expert_cnt, tensor_m=tensor_m) + self._remember_validation(self._validated_offsets, offsets, extra) + + def _validate_pointer_values_once(self, b_ptrs: torch.Tensor) -> None: + if self._is_validation_cached(self._validated_pointer_values, b_ptrs, self.expert_cnt): + return + pointer_values = self._copy_values_to_host(b_ptrs) + if any(value == 0 or value % 16 != 0 for value in pointer_values): + raise ValueError("b_ptrs entries must be non-null and 16-byte aligned") + self._remember_validation(self._validated_pointer_values, b_ptrs, self.expert_cnt) + + @staticmethod + def _validate_data_alignment(tensor: torch.Tensor, name: str) -> None: + if tensor.data_ptr() % 16 != 0: + raise ValueError(f"{name} data pointer must be 16-byte aligned") + + @staticmethod + def _validate_pointer_array_alignment(tensor: torch.Tensor) -> None: + if tensor.data_ptr() % 8 != 0: + raise ValueError("b_ptrs data pointer must be 8-byte aligned") + + @staticmethod + def _record_pointer_stream(b_ptrs: torch.Tensor, current_stream: cuda.CUstream) -> None: + handle = int(current_stream) + torch_current = torch.cuda.current_stream(b_ptrs.device) + torch_default = torch.cuda.default_stream(b_ptrs.device) + if handle == torch_current.cuda_stream: + launch_stream = torch_current + elif handle == torch_default.cuda_stream: + launch_stream = torch_default + else: + launch_stream = torch.cuda.ExternalStream(handle, device=b_ptrs.device) + b_ptrs.record_stream(launch_stream) + + def check_support(self) -> bool: + if self.a_desc.ndim != 3: + raise ValueError(f"sample_a must be rank-3, got {self.a_desc.shape}") + tensor_m, k, one = self.a_desc.shape + if one != 1: + raise ValueError(f"sample_a trailing dimension must be 1, got {one}") + + if self.weight_mode == MoEWeightMode.DENSE: + if self.b_desc.ndim != 3: + raise ValueError(f"sample_b must be rank-3, got {self.b_desc.shape}") + n, b_k, experts = self.b_desc.shape + if b_k != k: + raise ValueError(f"sample_b K dimension ({b_k}) must match sample_a ({k})") + if experts != self.expert_cnt: + raise ValueError("sample_b expert dimension is inconsistent") + self._expect_stride(self.b_desc, (k, 1, n * k), "sample_b") + else: + if len(self.b_shape) not in (2, 3): + raise ValueError(f"b_shape must be rank-2 or rank-3, got {self.b_shape}") + n, b_k = self.b_shape[:2] + if len(self.b_shape) == 3 and self.b_shape[2] != 1: + raise ValueError(f"b_shape trailing dimension must be 1, got {self.b_shape}") + if b_k != k: + raise ValueError(f"b_shape K dimension ({b_k}) must match sample_a ({k})") + if n % 64 != 0: + raise ValueError(f"N must be divisible by 64 for paired GLU blocks, got {n}") + + n_out = n // 2 + self._expect_shape(self.c_desc, (tensor_m, n, 1), "sample_c") + self._expect_shape(self.d_desc, (tensor_m, n_out, 1), "sample_d") + self._expect_shape(self.padded_offsets_desc, (self.expert_cnt,), "sample_padded_offsets") + self._expect_shape(self.alpha_desc, (self.expert_cnt,), "sample_alpha") + if self.prob_desc is None: + raise ValueError("sample_prob is required") + self._expect_shape(self.prob_desc, (tensor_m, 1, 1), "sample_prob") + + self._expect_stride(self.a_desc, (k, 1, tensor_m * k), "A tensor") + self._expect_stride(self.c_desc, (n, 1, tensor_m * n), "sample_c") + self._expect_stride(self.d_desc, (n_out, 1, tensor_m * n_out), "sample_d") + self._expect_stride(self.padded_offsets_desc, (1,), "sample_padded_offsets") + self._expect_stride(self.alpha_desc, (1,), "sample_alpha") + self._expect_stride(self.prob_desc, (1, 1, 1), "sample_prob") + + self._check_dtype(self.a_desc, torch.bfloat16, "sample_a") + if self.weight_mode == MoEWeightMode.DENSE: + self._check_dtype(self.b_desc, torch.bfloat16, "sample_b") + self._check_dtype(self.b_dtype, torch.bfloat16, "b_dtype") + self._check_dtype(self.c_desc, _OUTPUT_DTYPES, "sample_c") + self._check_dtype(self.d_desc, _OUTPUT_DTYPES, "sample_d") + self._check_dtype(self.padded_offsets_desc, torch.int32, "sample_padded_offsets") + self._check_dtype(self.alpha_desc, torch.float32, "sample_alpha") + self._check_dtype(self.prob_desc, torch.float32, "sample_prob") + + device = self.a_desc.device + for desc, name in ( + (self.c_desc, "sample_c"), + (self.d_desc, "sample_d"), + (self.padded_offsets_desc, "sample_padded_offsets"), + (self.alpha_desc, "sample_alpha"), + (self.prob_desc, "sample_prob"), + ): + self._expect_device(desc, device, name) + if self.b_desc is not None: + self._expect_device(self.b_desc, device, "sample_b") + + if self.bias_desc is not None: + self._expect_shape(self.bias_desc, (n, self.expert_cnt), "sample_bias") + self._expect_stride(self.bias_desc, (1, n), "sample_bias") + self._check_dtype(self.bias_desc, _OUTPUT_DTYPES, "sample_bias") + self._expect_device(self.bias_desc, device, "sample_bias") + + for name, data_ptr in self._sample_data_ptrs.items(): + if data_ptr % 16 != 0: + raise ValueError(f"{name} data pointer must be 16-byte aligned") + + if self.acc_dtype != torch.float32: + raise ValueError(f"acc_dtype must be torch.float32, got {self.acc_dtype}") + if self.m_aligned != MoEGroupedGemmGluBiasBf16Kernel.FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be 256, got {self.m_aligned}") + if self.act_func not in ("swiglu", "geglu"): + raise ValueError(f"act_func must be 'swiglu' or 'geglu', got {self.act_func}") + if self.b_major not in ("k", "n"): + raise ValueError(f"b_major must be 'k' or 'n', got {self.b_major}") + if self.expert_cnt <= 0 or self.expert_cnt > 1024: + raise ValueError(f"expert count must be in [1, 1024], got {self.expert_cnt}") + if tensor_m % 256 != 0: + raise ValueError(f"sample_a M dimension must be 256-aligned, got {tensor_m}") + + self._validate_offset_sequence( + self._sample_offset_values, + expert_cnt=self.expert_cnt, + tensor_m=tensor_m, + ) + sample_offsets = self._sample_offsets_ref() + if sample_offsets is not None and int(sample_offsets._version) == self._sample_offsets_version: + self._remember_validation( + self._validated_offsets, + sample_offsets, + (self.expert_cnt, tensor_m), + ) + elif sample_offsets is not None: + self._validate_offsets_once(sample_offsets, tensor_m=tensor_m) + + if not self._kernel.can_implement( + _convert_to_cutlass_data_type(torch.bfloat16), + _convert_to_cutlass_data_type(self.c_desc.dtype), + _convert_to_cutlass_data_type(self.d_desc.dtype), + _convert_to_cutlass_data_type(self.acc_dtype), + self.use_2cta_instrs, + self.mma_tiler_mn, + self.cluster_shape_mn, + tensor_m, + n, + k, + self.expert_cnt, + "k", + self.b_major, + "n", + self.m_aligned, + ): + raise ValueError("Unsupported BF16 grouped GEMM GLU tile, cluster, alignment, " "or layout configuration") + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + major, minor = torch.cuda.get_device_capability(self.a_desc.device) + compute_capability = major * 10 + minor + if compute_capability < 100: + raise RuntimeError(f"GroupedGemmGluSm100 requires SM100+, found SM{compute_capability} " f"on {self.a_desc.device}") + + self._is_supported = True + return True + + def compile(self) -> None: + self._ensure_support_checked() + if self._compiled_kernel is not None: + return + + kernel = self._kernel( + acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), + use_2cta_instrs=self.use_2cta_instrs, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + vectorized_f32=self.vector_f32, + expert_cnt=self.expert_cnt, + weight_mode=self.weight_mode, + use_dynamic_sched=self.use_dynamic_sched, + act_func=self.act_func, + enable_bias=self._has_bias, + generate_c=self.generate_c, + ) + + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - self.num_cluster_overlap_margin + if max_active_clusters <= 0: + raise ValueError("max_active_clusters must be > 0 after applying " "CUDNNFE_CLUSTER_OVERLAP_MARGIN") + + workspace_bytes = kernel.get_workspace_bytes() + self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device=self.a_desc.device) + if self._workspace.data_ptr() % 128 != 0: + raise RuntimeError("workspace allocation must be 128-byte aligned") + workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + + valid_m = cute.sym_int(divisibility=256) + a_fake = self._make_fake_cute_compact_tensor( + self.a_desc.dtype, + self.a_desc.shape, + self.a_desc.stride_order, + dynamic_mode=0, + divisibility=256, + ) + c_fake = self._make_fake_cute_compact_tensor( + self.c_desc.dtype, + self.c_desc.shape, + self.c_desc.stride_order, + dynamic_mode=0, + divisibility=256, + ) + d_fake = self._make_fake_cute_compact_tensor( + self.d_desc.dtype, + self.d_desc.shape, + self.d_desc.stride_order, + dynamic_mode=0, + divisibility=256, + ) + prob_fake = self._make_fake_cute_tensor(self.prob_desc.dtype, (valid_m, 1, 1), self.prob_desc.stride) + + if self.weight_mode == MoEWeightMode.DENSE: + b_fake = self._make_fake_cute_tensor_from_desc(self.b_desc) + n_value = cutlass.Int32(0) + k_value = cutlass.Int32(0) + b_stride = cutlass.Int64(0) + b_major_mode = OperandMajorMode.K + else: + self._compile_b_ptrs = torch.empty((self.expert_cnt,), dtype=torch.int64, device=self.a_desc.device) + self._validate_pointer_array_alignment(self._compile_b_ptrs) + b_fake = from_dlpack(self._compile_b_ptrs, assumed_align=8).iterator + n, k = self.b_shape[:2] + n_value = cutlass.Int32(n) + k_value = cutlass.Int32(k) + b_stride = cutlass.Int64(k if self.b_major == "k" else n) + b_major_mode = OperandMajorMode.K if self.b_major == "k" else OperandMajorMode.MN + + raw_compiled = cute.compile( + kernel, + a=a_fake, + b=b_fake, + n=n_value, + k=k_value, + b_stride_size=b_stride, + b_major_mode=b_major_mode, + workspace_ptr=workspace_ptr, + c=c_fake, + d=d_fake, + padded_offsets=self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc), + alpha=self._make_fake_cute_tensor_from_desc(self.alpha_desc), + prob=prob_fake, + bias=self._make_fake_cute_tensor_from_desc(self.bias_desc), + max_active_clusters=max_active_clusters, + stream=fake_stream, + linear_offset=cutlass.Float32(0.0), + options="--enable-tvm-ffi", + ) + + cached_n = n_value + cached_k = k_value + cached_b_stride = b_stride + + def tensor_api( + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor], + b_ptrs: Optional[torch.Tensor], + bias_tensor: Optional[torch.Tensor], + prob_tensor: torch.Tensor, + stream: cuda.CUstream, + linear_offset: float, + ) -> None: + b_arg = b_tensor if self.weight_mode == MoEWeightMode.DENSE else int(b_ptrs.data_ptr()) + raw_compiled( + a_tensor, + b_arg, + cached_n, + cached_k, + cached_b_stride, + workspace_ptr, + c_tensor, + d_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + bias_tensor, + stream, + cutlass.Float32(linear_offset), + ) + + self._compiled_kernel = tensor_api + + def _validate_live_tensor( + self, + tensor: torch.Tensor, + sample: TensorDesc, + name: str, + *, + dynamic_m: bool = False, + ) -> TensorDesc: + desc = self._make_tensor_desc(tensor, name=name) + if desc.dtype != sample.dtype: + raise ValueError(f"{name} dtype mismatch: expected {sample.dtype}, got {desc.dtype}") + if desc.device != sample.device: + raise ValueError(f"{name} device mismatch: expected {sample.device}, got {desc.device}") + if dynamic_m: + if desc.shape[1:] != sample.shape[1:]: + raise ValueError(f"{name} shape suffix mismatch: expected {sample.shape[1:]}, " f"got {desc.shape[1:]}") + if desc.stride_order != sample.stride_order: + raise ValueError(f"{name} layout mismatch: expected stride order " f"{sample.stride_order}, got {desc.stride_order}") + elif desc.shape != sample.shape or desc.stride != sample.stride: + raise ValueError(f"{name} descriptor mismatch: expected shape/stride " f"{sample.shape}/{sample.stride}, got {desc.shape}/{desc.stride}") + return desc + + def execute( + self, + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor] = None, + b_ptrs: Optional[torch.Tensor] = None, + bias_tensor: Optional[torch.Tensor] = None, + prob_tensor: Optional[torch.Tensor] = None, + linear_offset: float = 0.0, + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + current_stream = self._get_default_stream(current_stream) + if self._compiled_kernel is None: + raise RuntimeError("Kernel not compiled; call compile() first") + if prob_tensor is None: + raise ValueError("prob_tensor is required") + + a_desc = self._validate_live_tensor(a_tensor, self.a_desc, "a_tensor", dynamic_m=True) + c_desc = self._validate_live_tensor(c_tensor, self.c_desc, "c_tensor", dynamic_m=True) + d_desc = self._validate_live_tensor(d_tensor, self.d_desc, "d_tensor", dynamic_m=True) + self._validate_live_tensor(padded_offsets, self.padded_offsets_desc, "padded_offsets") + self._validate_live_tensor(alpha_tensor, self.alpha_desc, "alpha_tensor") + prob_desc = self._validate_live_tensor(prob_tensor, self.prob_desc, "prob_tensor", dynamic_m=True) + + tensor_m, k, _ = a_desc.shape + n = c_desc.shape[1] + n_out = n // 2 + if tensor_m % 256 != 0: + raise ValueError(f"a_tensor M dimension must be 256-aligned, got {tensor_m}") + self._expect_shape(c_desc, (tensor_m, n, 1), "c_tensor") + self._expect_shape(d_desc, (tensor_m, n_out, 1), "d_tensor") + self._expect_shape(prob_desc, (tensor_m, 1, 1), "prob_tensor") + self._expect_stride(a_desc, (k, 1, tensor_m * k), "a_tensor") + self._expect_stride(c_desc, (n, 1, tensor_m * n), "c_tensor") + self._expect_stride(d_desc, (n_out, 1, tensor_m * n_out), "d_tensor") + self._expect_stride(prob_desc, (1, 1, 1), "prob_tensor") + self._validate_offsets_once(padded_offsets, tensor_m=tensor_m) + + for tensor, name in ( + (a_tensor, "a_tensor"), + (c_tensor, "c_tensor"), + (d_tensor, "d_tensor"), + (padded_offsets, "padded_offsets"), + (alpha_tensor, "alpha_tensor"), + (prob_tensor, "prob_tensor"), + ): + self._validate_data_alignment(tensor, name) + + if self._has_bias: + if bias_tensor is None: + raise ValueError("bias_tensor is required because the API was compiled " "with sample_bias") + self._validate_live_tensor(bias_tensor, self.bias_desc, "bias_tensor") + self._validate_data_alignment(bias_tensor, "bias_tensor") + elif bias_tensor is not None: + raise ValueError("bias_tensor must be omitted because the API was compiled " "without sample_bias") + + if self.weight_mode == MoEWeightMode.DENSE: + if b_tensor is None or b_ptrs is not None: + raise ValueError("Dense execution requires b_tensor and forbids b_ptrs") + self._validate_live_tensor(b_tensor, self.b_desc, "b_tensor") + self._validate_data_alignment(b_tensor, "b_tensor") + else: + if b_tensor is not None or b_ptrs is None: + raise ValueError("Discrete execution requires b_ptrs and forbids b_tensor") + _require_pointer_tensor(b_ptrs, "b_ptrs", self.expert_cnt) + if b_ptrs.device != self.a_desc.device: + raise ValueError(f"b_ptrs must be on the same device as a_tensor " f"({self.a_desc.device}), got {b_ptrs.device}") + if b_ptrs.data_ptr() % 8 != 0: + raise ValueError("b_ptrs data pointer must be 8-byte aligned") + self._validate_pointer_values_once(b_ptrs) + self._record_pointer_stream(b_ptrs, current_stream) + + self._compiled_kernel( + a_tensor, + c_tensor, + d_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + b_ptrs, + bias_tensor, + prob_tensor, + current_stream, + linear_offset, + ) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu/_blockscaled_api.py b/python/cudnn/grouped_gemm/grouped_gemm_glu/_blockscaled_api.py new file mode 100644 index 000000000..772c299a7 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/_blockscaled_api.py @@ -0,0 +1,1263 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Unified API for Grouped GEMM GLU Forward Kernel (SM100+) + +This module provides a single API class that supports both contiguous (dense) +and discrete weight modes for block-scaled grouped GEMM with GLU activation +(SwiGLU / GeGLU) in MoE (Mixture of Experts) workloads. + +Dense mode + All expert weights are packed contiguously in a 3-D tensor (N, K, L). + Callers supply ``sample_b`` and ``sample_sfb``. + +Discrete mode + Each expert has its own memory allocation. Callers supply + ``num_experts``, ``b_shape``, ``b_dtype``, and per-expert pointer arrays + at execution time. +""" + +from .moe_blockscaled_grouped_gemm_glu_bias import BlockScaledMoEGroupedGemmGluBiasKernel +from ..grouped_gemm_utils import _torch_stream_context +from ..moe_utils import MoEWeightMode +from cuda.bindings import driver as cuda +import os +import torch +from typing import Tuple, Optional + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import OperandMajorMode +from cutlass.cute.runtime import from_dlpack, make_fake_stream + +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.api_base import APIBase, ceil_div, is_power_of_2 + + +class GroupedGemmGluBlockScaledAPI(APIBase): + """Unified API for grouped GEMM GLU forward operation on SM100+ GPUs. + + This kernel performs block-scaled grouped GEMM with GLU activation + (SwiGLU or GeGLU), designed for MoE workloads. It supports both dense + (contiguous) and discrete (per-expert pointer) weight layouts through + the ``BlockScaledMoEGroupedGemmGluBiasKernel``. + + Weight mode is auto-detected from the constructor arguments: + + - **Dense**: provide ``sample_b`` and ``sample_sfb``. + - **Discrete**: provide ``num_experts``, ``b_shape``, and ``b_dtype``. + + Example:: + + # Dense mode + api = GroupedGemmGluSm100( + sample_a=a, sample_b=b, sample_c=c, sample_d=d, + sample_sfa=sfa, sample_sfb=sfb, + sample_padded_offsets=offsets, sample_alpha=alpha, + sample_d_col=d_col, + ) + + # Discrete mode + api = GroupedGemmGluSm100( + sample_a=a, num_experts=8, b_shape=(n, k), b_dtype=torch.uint8, + sample_c=c, sample_d=d, sample_sfa=sfa, + sample_padded_offsets=offsets, sample_alpha=alpha, + sample_d_col=d_col, + ) + + api.check_support() + api.compile() + api.execute(...) + """ + + def __init__( + self, + sample_a: torch.Tensor, + sample_c: torch.Tensor, + sample_d: torch.Tensor, + sample_sfa: torch.Tensor, + sample_padded_offsets: torch.Tensor, + sample_alpha: torch.Tensor, + sample_d_col: torch.Tensor, + # Dense mode (contiguous) -- provide these. sample_bias is optional: + sample_b: Optional[torch.Tensor] = None, + sample_sfb: Optional[torch.Tensor] = None, + sample_bias: Optional[torch.Tensor] = None, + # Discrete mode -- provide these instead: + num_experts: Optional[int] = None, + b_shape: Optional[Tuple[int, ...]] = None, + b_dtype: Optional[torch.dtype] = None, + # Optional quantization output arguments + sample_sfd_row: Optional[torch.Tensor] = None, + sample_sfd_col: Optional[torch.Tensor] = None, + sample_amax: Optional[torch.Tensor] = None, + sample_norm_const: Optional[torch.Tensor] = None, + sample_prob: Optional[torch.Tensor] = None, + # Configuration + acc_dtype: torch.dtype = torch.float32, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "swiglu", + b_major: str = "k", + use_dynamic_sched: bool = False, + ): + """Initialize the GroupedGemmGluSm100 API. + + :param sample_a: Sample A tensor (valid_m, k, 1) + :param sample_c: Sample C tensor for intermediate storage + :param sample_d: Sample D output tensor (valid_m, n/2, 1) after GLU + :param sample_sfa: Sample scale factor A tensor + :param sample_padded_offsets: End offset for each expert after padding + :param sample_alpha: Per-group alpha scaling factors + :param sample_d_col: Column-quantized D tensor + :param sample_b: (Dense) Sample B tensor (n, k, l) + :param sample_sfb: (Dense) Sample scale factor B tensor + :param sample_bias: Optional bias tensor with shape (n, l) or (n, expert_cnt), stride (1, n). + Dense mode supports fp16/bfloat16/float32 bias; discrete mode supports fp16/bfloat16 bias. + :param num_experts: (Discrete) Number of experts + :param b_shape: (Discrete) Shape of a single expert B tensor, e.g. (n, k) + :param b_dtype: (Discrete) Data type of B tensors + :param sample_sfd_row: Optional row scale factor for D + :param sample_sfd_col: Optional column scale factor for D + :param sample_amax: Optional amax tensor for quantization + :param sample_norm_const: Optional normalization constant + :param sample_prob: Optional probability tensor for gating + :param acc_dtype: Accumulator data type + :param mma_tiler_mn: MMA tiler shape (M, N) + :param cluster_shape_mn: Cluster shape (M, N) + :param sf_vec_size: Scale factor vector size + :param vector_f32: Use vectorized f32 operations + :param m_aligned: Alignment for group M dimension + :param discrete_col_sfd: Generate discrete col-major scale factor tensor + :param act_func: Activation function, one of "swiglu" or "geglu" + :param b_major: Major dimension for B tensor, one of "k" or "n" + :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing + """ + super().__init__() + + self._warn_experimental_api() + self._logger.debug("Entering __init__") + + # ---- Weight mode auto-detection ---- + if sample_b is not None and num_experts is None: + self.weight_mode = MoEWeightMode.DENSE + if sample_sfb is None: + raise ValueError("sample_sfb is required when sample_b is provided (dense mode)") + elif num_experts is not None and sample_b is None: + self.weight_mode = MoEWeightMode.DISCRETE + if b_shape is None or b_dtype is None: + raise ValueError("b_shape and b_dtype are required in discrete mode") + else: + raise ValueError("Provide either (sample_b, sample_sfb) for dense mode " "or (num_experts, b_shape, b_dtype) for discrete mode, but not both.") + + # ---- Common tensor descriptors ---- + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") + self.d_desc = self._make_tensor_desc(sample_d, name="sample_d") + self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") + + self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") + self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias") + self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row") + self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col") + self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") + self.norm_const_desc = self._unpad_tensor_to_ndim( + self._make_tensor_desc(sample_norm_const, name="sample_norm_const"), + 1, + "norm_const", + ) + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") + + # ---- Mode-specific state ---- + if self.weight_mode == MoEWeightMode.DENSE: + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") + self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb") + self.expert_cnt = self.padded_offsets_desc.shape[0] + else: + self._value_error_if(num_experts == 0, "num_experts must be > 0") + self.expert_cnt = num_experts + self.b_shape = b_shape + self.b_dtype = b_dtype + self.b_major = b_major + self._value_error_if( + self.padded_offsets_desc.shape[0] != self.expert_cnt, + f"padded_offsets length ({self.padded_offsets_desc.shape[0]}) " f"must equal num_experts ({self.expert_cnt})", + ) + + # ---- Configuration ---- + self.acc_dtype = acc_dtype + self.mma_tiler_mn = mma_tiler_mn + self.use_2cta_instrs = mma_tiler_mn[0] == 256 + if cluster_shape_mn is None: + self.cluster_shape_mn = (2, 1) if self.use_2cta_instrs else (1, 1) + else: + self.cluster_shape_mn = cluster_shape_mn + self.sf_vec_size = sf_vec_size + self.vector_f32 = vector_f32 + self.m_aligned = m_aligned + self.discrete_col_sfd = discrete_col_sfd + self.act_func = act_func + if self.weight_mode == MoEWeightMode.DENSE: + self.b_major = b_major # stored for both modes + + self.use_dynamic_sched = use_dynamic_sched + + self._interpret_uint8_as_fp4x2 = True + self._has_bias = self.bias_desc is not None + self._kernel = BlockScaledMoEGroupedGemmGluBiasKernel + + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + self._logger.debug(f"setting num_cluster_overlap_margin: {self.num_cluster_overlap_margin}") + + self._workspace = None + + self._logger.debug("__init__ completed") + + # --------------------------------------------------------------------- # + # check_support + # --------------------------------------------------------------------- # + + def check_support(self) -> bool: + """Check if the kernel configuration is supported. + + :return: True if supported, raises exception otherwise + """ + self._logger.debug("Entering check_support") + + # ---- SFD group validation ---- + all_none = all(x is None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) + all_provided = all(x is not None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) + self._value_error_if( + not (all_none or all_provided), + "sfd_row_desc, sfd_col_desc, and norm_const_desc must be all None or all not None", + ) + self.generate_sfd = all_provided + if self.discrete_col_sfd and not self.generate_sfd: + self._logger.warning("discrete_col_sfd is True but generate_sfd is False, discrete_col_sfd will be ignored") + self.discrete_col_sfd = False + + # ---- Shapes and strides ---- + self._logger.debug("Checking tensor shapes and strides") + tensor_m, k, _one = self._tensor_shape(self.a_desc, name="sample_a") + + if self.weight_mode == MoEWeightMode.DENSE: + n, _, l = self._tensor_shape(self.b_desc, name="sample_b") + else: + # Discrete: extract n, k from b_shape + if len(self.b_shape) == 2: + n, b_k = self.b_shape + else: + n, b_k, _ = self.b_shape + self._value_error_if(b_k != k, f"B K dimension ({b_k}) must match A K dimension ({k})") + l = self.expert_cnt # for shape checks that use l + + _, n_2, _one = self._tensor_shape(self.d_desc, name="sample_d") + + self._value_error_if( + n % 64 != 0, + f"N must be divisible by 64 for GLU (two consecutive 32-column blocks), got N={n}", + ) + + self._check_tensor_shape(self.a_desc, (tensor_m, k, 1), "A") + if self.weight_mode == MoEWeightMode.DENSE: + self._check_tensor_shape(self.b_desc, (n, k, l), "B") + self._check_tensor_shape(self.c_desc, (tensor_m, n, 1), "C") + self._check_tensor_shape(self.d_desc, (tensor_m, n // 2, 1), "D") + self._check_tensor_shape(self.d_col_desc, (tensor_m, n // 2, 1), "D_col") + self._check_tensor_shape(self.bias_desc, (n, l), "bias") + + rest_k = ceil_div(ceil_div(k, self.sf_vec_size), 4) + self._check_tensor_shape(self.sfa_desc, (32, 4, ceil_div(tensor_m, 128), 4, rest_k, 1), "SFA") + if self.weight_mode == MoEWeightMode.DENSE: + self._check_tensor_shape(self.sfb_desc, (32, 4, ceil_div(n, 128), 4, rest_k, l), "SFB") + + rest_n2 = ceil_div(ceil_div(n // 2, self.sf_vec_size), 4) + self._check_tensor_shape( + self.sfd_row_desc, + (32, 4, ceil_div(tensor_m, 128), 4, rest_n2, 1), + "SFD_row", + ) + rest_m = ceil_div(ceil_div(tensor_m, self.sf_vec_size), 4) + self._check_tensor_shape(self.sfd_col_desc, (32, 4, ceil_div(n // 2, 128), 4, rest_m, 1), "SFD_col") + + self._check_tensor_shape(self.alpha_desc, (self.expert_cnt,), "alpha") + self._check_tensor_shape(self.prob_desc, (tensor_m, 1, 1), "prob") + self._check_tensor_shape(self.amax_desc, (self.expert_cnt, 1), "amax") + self._check_tensor_shape(self.norm_const_desc, (1,), "norm_const") + self._check_tensor_shape(self.padded_offsets_desc, (self.expert_cnt,), "padded_offsets") + + # Strides + _ = self._check_tensor_stride( + self.a_desc, + stride=[(k, 1, tensor_m * k)], + extra_error_msg="A must have k-major layout", + ) + if self.weight_mode == MoEWeightMode.DENSE: + _ = self._check_tensor_stride( + self.b_desc, + stride=[(k, 1, n * k)], + extra_error_msg="B must have k-major layout", + ) + _ = self._check_tensor_stride( + self.c_desc, + stride=[(n, 1, tensor_m * n)], + extra_error_msg="C must have n-major layout", + ) + _ = self._check_tensor_stride( + self.d_desc, + stride=[(n_2, 1, tensor_m * n_2)], + extra_error_msg="D must have n-major layout", + ) + _ = self._check_tensor_stride( + self.d_col_desc, + stride=[(n_2, 1, tensor_m * n_2)], + extra_error_msg="D_col must have n-major layout", + ) + _ = self._check_tensor_stride( + self.bias_desc, + stride=[(1, n)], + ) + + # ---- Data types ---- + self._logger.debug("Checking data types") + self.ab_dtype = self._check_dtype( + self.a_desc, + dtype=[ + torch.float4_e2m1fn_x2, + torch.uint8, + torch.float8_e5m2, + torch.float8_e4m3fn, + ], + name="A/B", + ) + if self.weight_mode == MoEWeightMode.DENSE: + self._check_dtype( + self.b_desc, + dtype=self.ab_dtype, + name="B", + extra_error_msg="B must have the same dtype as A", + ) + self._check_dtype( + self.bias_desc, + dtype=[torch.bfloat16, torch.float16, torch.float32], + name="bias", + extra_error_msg="bias must be fp16, bfloat16, or float32", + ) + else: + self._value_error_if( + self.b_dtype != self.ab_dtype, + f"b_dtype ({self.b_dtype}) must match A dtype ({self.ab_dtype})", + ) + self._check_dtype( + self.bias_desc, + dtype=[torch.bfloat16, torch.float16], + name="bias", + extra_error_msg="bias must be fp16 or bfloat16 in discrete mode", + ) + + self.sf_dtype = self._check_dtype( + self.sfa_desc, + dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], + name="SFA/SFB/SFD", + ) + if self.weight_mode == MoEWeightMode.DENSE: + self._check_dtype( + self.sfb_desc, + dtype=self.sf_dtype, + name="SFB", + extra_error_msg="SFB must have the same dtype as SFA", + ) + self._check_dtype( + self.sfd_row_desc, + dtype=self.sf_dtype, + name="SFD_row", + extra_error_msg="SFD_row must have the same dtype as SFA", + ) + self._check_dtype( + self.sfd_col_desc, + dtype=self.sf_dtype, + name="SFD_col", + extra_error_msg="SFD_col must have the same dtype as SFA", + ) + + self._value_error_if( + self.sf_vec_size not in [16, 32], + f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}", + ) + self._value_error_if( + self.sf_dtype in [torch.float8_e4m3fn] and self.sf_vec_size == 32, + f"sf_dtype {self.sf_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", + ) + self._value_error_if( + self._is_fp8(self.ab_dtype) and self.sf_vec_size == 16, + f"ab_dtype {self.ab_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", + ) + + self._check_dtype( + self.acc_dtype, + dtype=torch.float32, + name="Accumulator", + extra_error_msg="Accumulator must be float32", + ) + self.c_dtype = self._check_dtype( + self.c_desc, + dtype=[ + torch.float32, + torch.float16, + torch.bfloat16, + torch.float8_e4m3fn, + torch.float8_e5m2, + torch.float4_e2m1fn_x2, + ], + name="C", + ) + + if self._is_fp4x2(self.ab_dtype): + self.d_dtype = self._check_dtype( + self.d_desc, + dtype=[torch.float16, torch.bfloat16, torch.float32], + name="D", + extra_error_msg="D must be fp16, bf16, or float32 when ab_dtype is fp4", + ) + else: + self.d_dtype = self._check_dtype( + self.d_desc, + dtype=[ + torch.float16, + torch.bfloat16, + torch.float8_e4m3fn, + torch.float8_e5m2, + torch.float4_e2m1fn_x2, + ], + name="D", + ) + self._check_dtype( + self.d_col_desc, + dtype=self.d_dtype, + name="D_col", + extra_error_msg="D_col must have the same dtype as D", + ) + + self._not_implemented_error_if( + self.bias_desc is None and self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype == torch.float32, + "Invalid configuration: fp4 ab_dtype, sf_vec_size 16, d_dtype float32 is not supported. " "Please use sf_vec_size 32 or d_dtype bf16 instead", + ) + + # ---- Activation function validation (both modes) ---- + self._value_error_if( + self.act_func not in ["swiglu", "geglu"], + f"act_func must be 'swiglu' or 'geglu', got {self.act_func}", + ) + + # ---- Discrete-mode-specific validation ---- + if self.weight_mode == MoEWeightMode.DISCRETE: + self._value_error_if( + self.b_major not in ["k", "n"], + f"b_major must be 'k' or 'n', got {self.b_major}", + ) + self._value_error_if( + self._is_fp4x2(self.ab_dtype) and self.b_major != "k", + "b_major must be 'k' when ab_dtype is fp4", + ) + + # ---- MMA tile / cluster shape ---- + self._logger.debug("Checking MMA tile shape and cluster shape") + self._value_error_if( + not self.use_2cta_instrs and self.mma_tiler_mn[0] != 128, + f"MMA tiler M must be 128 when use_2cta_instrs=False, got {self.mma_tiler_mn[0]}", + ) + self._value_error_if( + self.use_2cta_instrs and self.mma_tiler_mn[0] != 256, + f"MMA tiler M must be 256 when use_2cta_instrs=True, got {self.mma_tiler_mn[0]}", + ) + self._value_error_if( + self.mma_tiler_mn[1] != 256, + f"MMA tiler N must be 256, got {self.mma_tiler_mn[1]}", + ) + self._value_error_if( + self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0, + f"cluster_shape_mn[0] must be divisible by 2 when use_2cta_instrs=True, got {self.cluster_shape_mn[0]}", + ) + self._value_error_if( + not ( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1] <= 16 + and self.cluster_shape_mn[0] > 0 + and self.cluster_shape_mn[1] > 0 + and self.cluster_shape_mn[0] <= 4 + and self.cluster_shape_mn[1] <= 4 + and is_power_of_2(self.cluster_shape_mn[0]) + and is_power_of_2(self.cluster_shape_mn[1]) + ), + f"Invalid cluster shape: expected values to be powers of 2 and product <= 16, got {self.cluster_shape_mn}", + ) + cluster_tiler_m = (self.cluster_shape_mn[0] // (2 if self.use_2cta_instrs else 1)) * self.mma_tiler_mn[0] + self._value_error_if( + cluster_tiler_m not in [128, 256], + f"Invalid cluster tiler shape: expected cluster_tiler_m in {{128, 256}}, got {cluster_tiler_m}", + ) + self._value_error_if( + self.m_aligned % self.mma_tiler_mn[0] != 0, + f"m_aligned must be divisible by mma_tiler_mn[0], got {self.m_aligned} % {self.mma_tiler_mn[0]} != 0", + ) + self._value_error_if( + self.m_aligned != BlockScaledMoEGroupedGemmGluBiasKernel.FIX_PAD_SIZE, + f"m_aligned must be {BlockScaledMoEGroupedGemmGluBiasKernel.FIX_PAD_SIZE} (FIX_PAD_SIZE), got {self.m_aligned}", + ) + + # ---- Tensor alignment ---- + self._logger.debug("Checking tensor alignment") + + def check_contiguous_16B_alignment(dtype, stride_order, tensor_shape): + is_mode0_major = stride_order == (0, 1, 2) + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // (_convert_to_cutlass_data_type(dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2).width) + return num_major_elements % num_contiguous_elements == 0 + + if self.weight_mode == MoEWeightMode.DENSE: + b_stride_order_for_check = self.b_desc.stride_order + b_shape_for_check = (n, k, l) + else: + b_stride_order_for_check = (0, 1, 2) if self.b_major == "n" else (1, 0, 2) + b_shape_for_check = (n, k, 1) + + self._value_error_if( + not ( + check_contiguous_16B_alignment(self.ab_dtype, self.a_desc.stride_order, (tensor_m, k, l)) + and check_contiguous_16B_alignment(self.ab_dtype, b_stride_order_for_check, b_shape_for_check) + and check_contiguous_16B_alignment(self.d_dtype, self.d_desc.stride_order, (tensor_m, n_2, 1)) + ), + "Invalid tensor alignment: tensors must be 16B aligned", + ) + + # ---- Expert count limit ---- + self._value_error_if( + self.expert_cnt > 1024, + f"expert_cnt must be <= 1024, got {self.expert_cnt}", + ) + + # ---- Disabled configurations ---- + self._not_implemented_error_if( + (self._is_fp8(self.ab_dtype)) and (self.mma_tiler_mn[1] == 128) and (self._is_fp8(self.d_dtype)), + "Invalid configuration: fp8 ab_dtype with mma_tiler_mn[1] == 128 and fp8 d_dtype is not supported. " "Please use mma_tiler_mn[1] == 256 instead", + ) + self._not_implemented_error_if( + self._is_fp4x2(self.ab_dtype) and (self.c_dtype not in [torch.float16, torch.bfloat16]), + f"Invalid configuration: for fp4 ab_dtype, c_dtype must be float16 or bfloat16, got {self.c_dtype}", + ) + self._not_implemented_error_if(self._has_bias and self.mma_tiler_mn[1] != 256, "Bias fusion currently requires mma_tiler_mn[1] == 256") + + # ---- SM100+ check ---- + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + device = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(device) + compute_capability = major * 10 + minor + if compute_capability < 100: + raise RuntimeError(f"GroupedGemmGlu requires SM100+ compute capability, " f"but found SM{compute_capability} on device {device}") + + self._is_supported = True + self._logger.debug("check_support completed successfully") + return True + + # --------------------------------------------------------------------- # + # compile + # --------------------------------------------------------------------- # + + def compile(self) -> None: + """Compile the kernel.""" + self._logger.debug("Entering compile") + self._ensure_support_checked() + if self._compiled_kernel is not None: + self._logger.debug("Kernel already compiled; skipping recompilation") + return + if self.a_desc.shape[0] == 0: + self._logger.debug("sample valid_m is zero, skipping kernel compilation") + return + + # ---- Instantiate the unified kernel ---- + gemm_glu = self._kernel( + sf_vec_size=self.sf_vec_size, + acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), + use_2cta_instrs=self.use_2cta_instrs, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + vectorized_f32=self.vector_f32, + generate_sfd=self.generate_sfd, + discrete_col_sfd=self.discrete_col_sfd, + expert_cnt=self.expert_cnt, + weight_mode=self.weight_mode, + act_func=self.act_func, + enable_bias=self._has_bias, + use_dynamic_sched=self.use_dynamic_sched, + ) + + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) + max_active_clusters -= self.num_cluster_overlap_margin + self._value_error_if( + max_active_clusters <= 0, + "max_active_clusters must be > 0 after applying overlap margin; reduce CUDNNFE_CLUSTER_OVERLAP_MARGIN", + ) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + + # ---- Allocate workspace ---- + workspace_bytes = gemm_glu.get_workspace_bytes() + self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device="cuda") + + if self.weight_mode == MoEWeightMode.DENSE: + self._compile_dense(gemm_glu, max_active_clusters, fake_stream) + else: + self._compile_discrete(gemm_glu, max_active_clusters, fake_stream) + + self._logger.debug("Kernel compiled successfully") + + # -- Dense compile path ------------------------------------------------- # + + def _compile_dense(self, gemm_glu, max_active_clusters, fake_stream) -> None: + """Compile for dense (contiguous) weight mode.""" + use_full_dynamic = os.environ.get("CUDNN_FE_GROUPED_GEMM_DYNAMIC_MNKL", "1") != "0" + + fake_workspace_ptr = cute.runtime.nullptr( + dtype=cutlass.Uint8, + assumed_align=128, + ) + + if not use_full_dynamic: + valid_m = cute.sym_int(divisibility=256) + + a_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.a_desc.dtype, + shape=(valid_m, *self.a_desc.shape[1:]), + stride_order=self.a_desc.stride_order, + ) + b_cute_fake = self._make_fake_cute_tensor_from_desc(self.b_desc, assumed_align=16) + c_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.c_desc.dtype, + shape=(valid_m, *self.c_desc.shape[1:]), + stride_order=self.c_desc.stride_order, + ) + d_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.d_desc.dtype, + shape=(valid_m, *self.d_desc.shape[1:]), + stride_order=self.d_desc.stride_order, + ) + d_col_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.d_col_desc.dtype, + shape=(valid_m, *self.d_col_desc.shape[1:]), + stride_order=self.d_col_desc.stride_order, + ) + + tensor_m_128 = cute.sym_int() + stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfa_cute_fake = self._make_fake_cute_tensor( + dtype=self.sfa_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfa_desc.shape[4], 1), + stride=(16, 4, self.sfa_desc.stride[2], 1, 512, stride_tensor_m_128), + ) + + sfb_cute_fake = self._make_fake_cute_tensor_from_desc(self.sfb_desc, assumed_align=16) + + prob_cute_fake = None + if self.prob_desc is not None: + prob_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.prob_desc.dtype, + shape=(valid_m, 1, 1), + stride_order=self.prob_desc.stride_order, + ) + + sfd_row_fake = None + sfd_col_fake = None + if self.sfd_row_desc is not None: + stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) + sfd_row_fake = self._make_fake_cute_tensor( + dtype=self.sfd_row_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), + stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), + ) + if self.sfd_col_desc is not None: + rest_m = cute.sym_int(divisibility=1) + stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) + stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) + sfd_col_fake = self._make_fake_cute_tensor( + dtype=self.sfd_col_desc.dtype, + shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), + stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), + ) + bias_cute_fake = self._make_fake_cute_tensor_from_desc(self.bias_desc, assumed_align=16) + else: + valid_m = cute.sym_int(divisibility=256) + n_sym = cute.sym_int() + n_2_sym = cute.sym_int() + k_sym = cute.sym_int() + l_sym = cute.sym_int() + + a_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.a_desc.dtype, + shape=(valid_m, k_sym, 1), + stride_order=self.a_desc.stride_order, + dynamic_mode=self.a_desc.stride_order[0], + divisibility=32 if self._is_fp4x2(self.ab_dtype) else 16, + ) + b_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.b_desc.dtype, + shape=(n_sym, k_sym, l_sym), + stride_order=self.b_desc.stride_order, + dynamic_mode=self.b_desc.stride_order[0], + divisibility=32 if self._is_fp4x2(self.ab_dtype) else 16, + ) + c_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.c_desc.dtype, + shape=(valid_m, n_2_sym, 1), + stride_order=self.c_desc.stride_order, + dynamic_mode=self.c_desc.stride_order[0], + divisibility=8 if self._is_f16(self.c_desc.dtype) else 16, + ) + d_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.d_desc.dtype, + shape=(valid_m, n_2_sym, 1), + stride_order=self.d_desc.stride_order, + dynamic_mode=self.d_desc.stride_order[0], + divisibility=8 if self._is_f16(self.d_desc.dtype) else 16, + ) + d_col_cute_fake = self._make_fake_cute_compact_tensor( + dtype=self.d_col_desc.dtype, + shape=(valid_m, n_2_sym, 1), + stride_order=self.d_col_desc.stride_order, + dynamic_mode=self.d_col_desc.stride_order[0], + divisibility=8 if self._is_f16(self.d_col_desc.dtype) else 16, + ) + + tensor_m_128 = cute.sym_int() + rest_k = cute.sym_int() + stride_rest_k = cute.sym_int(divisibility=32 * 4 * 4) + stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfa_shape = list(self.sfa_desc.shape) + sfa_shape[2] = tensor_m_128 + sfa_shape[4] = rest_k + sfa_stride = list(self.sfa_desc.stride) + sfa_stride[2] = stride_rest_k + sfa_stride[5] = stride_tensor_m_128 + sfa_cute_fake = self._make_fake_cute_tensor( + dtype=self.sfa_desc.dtype, + shape=tuple(sfa_shape), + stride=tuple(sfa_stride), + ) + + tensor_n_128 = cute.sym_int() + stride_sfb_rest_k = cute.sym_int(divisibility=32 * 4 * 4) + stride_sfb_tensor_n_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfb_cute_fake = self._make_fake_cute_tensor( + dtype=self.sfb_desc.dtype, + shape=(32, 4, tensor_n_128, 4, rest_k, l_sym), + stride=(16, 4, stride_sfb_tensor_n_128, 1, 512, stride_sfb_rest_k), + ) + + prob_cute_fake = None + if self.prob_desc is not None: + prob_cute_fake = self._make_fake_cute_tensor( + dtype=self.prob_desc.dtype, + shape=(valid_m, *self.prob_desc.shape[1:]), + stride=self.prob_desc.stride, + ) + + sfd_row_fake = None + sfd_col_fake = None + if self.sfd_row_desc is not None: + rest_n2 = cute.sym_int() + stride_sfd_rest_n2 = cute.sym_int(divisibility=32 * 4 * 4) + stride_sfd_rest_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfd_row_fake = self._make_fake_cute_tensor( + dtype=self.sfd_row_desc.dtype, + shape=(32, 4, tensor_m_128, 4, rest_n2, 1), + stride=(16, 4, stride_sfd_rest_n2, 1, 512, stride_sfd_rest_tensor_m_128), + ) + if self.sfd_col_desc is not None: + tensor_n2_128 = cute.sym_int() + rest_m_dyn = cute.sym_int() + stride_sfd_rest_m = cute.sym_int(divisibility=32 * 4 * 4) + stride_sfd_n2 = cute.sym_int(divisibility=32 * 4 * 4) + sfd_col_fake = self._make_fake_cute_tensor( + dtype=self.sfd_col_desc.dtype, + shape=(32, 4, tensor_n2_128, 4, rest_m_dyn, 1), + stride=(16, 4, stride_sfd_rest_m, 1, 512, stride_sfd_n2), + ) + bias_cute_fake = None + if self.bias_desc is not None: + bias_cute_fake = self._make_fake_cute_tensor( + dtype=self.bias_desc.dtype, + shape=(n_sym, l_sym), + stride=(1, n_sym), + ) + + # Compile with keyword args (dense mode uses the unified __call__ positional order). + # linear_offset, geglu_alpha, glu_clamp_max, and glu_clamp_min are runtime + # cutlass.Float32 (not Constexpr), so the compile-time placeholders below are + # irrelevant -- the values passed through tensor_api() at execute() time are + # what the kernel actually uses. + _compiled_kernel = cute.compile( + gemm_glu, + a=a_cute_fake, + b=b_cute_fake, + sfb=sfb_cute_fake, + n=cutlass.Int32(0), + k=cutlass.Int32(0), + b_stride_size=cutlass.Int64(0), + b_major_mode=OperandMajorMode.K, + workspace_ptr=fake_workspace_ptr, + c=c_cute_fake, + d=d_cute_fake, + d_col=d_col_cute_fake, + sfa=sfa_cute_fake, + sfd_row_tensor=sfd_row_fake, + sfd_col_tensor=sfd_col_fake, + amax_tensor=self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16), + norm_const_tensor=self._make_fake_cute_tensor_from_desc(self.norm_const_desc, assumed_align=16), + padded_offsets=self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16), + alpha=self._make_fake_cute_tensor_from_desc(self.alpha_desc, assumed_align=16), + prob=prob_cute_fake, + bias=bias_cute_fake, + max_active_clusters=max_active_clusters, + stream=fake_stream, + epilogue_op=lambda x: x, + linear_offset=cutlass.Float32(0.0), + geglu_alpha=cutlass.Float32(1.702), + glu_clamp_max=cutlass.Float32(7.0), + glu_clamp_min=cutlass.Float32(-7.0), + options="--enable-tvm-ffi", + ) + + # Cache workspace pointer for the tensor_api closure + cached_workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator + + def tensor_api( + a_tensor: torch.Tensor, + b_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + d_col_tensor: Optional[torch.Tensor], + sfa_tensor: torch.Tensor, + sfb_tensor: torch.Tensor, + sfd_row_tensor: Optional[torch.Tensor], + sfd_col_tensor: Optional[torch.Tensor], + amax_tensor: Optional[torch.Tensor], + norm_const_tensor: Optional[torch.Tensor], + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + prob_tensor: Optional[torch.Tensor], + bias_tensor: Optional[torch.Tensor], + stream: cuda.CUstream, + linear_offset: float = 0.0, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + ) -> None: + norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") + _compiled_kernel( + a_tensor, + b_tensor, + sfb_tensor, + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int64(0), + cached_workspace_ptr, + c_tensor, + d_tensor, + d_col_tensor, + sfa_tensor, + sfd_row_tensor, + sfd_col_tensor, + amax_tensor, + norm_const_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + bias_tensor, + stream, + cutlass.Float32(linear_offset), + cutlass.Float32(geglu_alpha), + cutlass.Float32(glu_clamp_max), + cutlass.Float32(glu_clamp_min), + ) + + self._compiled_kernel = tensor_api + + # -- Discrete compile path ---------------------------------------------- # + + def _compile_discrete(self, gemm_glu, max_active_clusters, fake_stream) -> None: + """Compile for discrete (per-expert pointer) weight mode.""" + if len(self.b_shape) == 2: + n, k = self.b_shape + else: + n, k, _ = self.b_shape + + b_major_mode = OperandMajorMode.K if self.b_major == "k" else OperandMajorMode.MN + if self.b_major == "k": + b_stride_size = k + else: + b_stride_size = n + + ab_cutlass_dtype = _convert_to_cutlass_data_type(self.a_desc.dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2) + align = 32 if ab_cutlass_dtype.width == 4 else 16 + + valid_m = cute.sym_int(divisibility=256) + a_tensor = self._make_fake_cute_compact_tensor( + dtype=self.a_desc.dtype, + shape=(valid_m, *self.a_desc.shape[1:]), + stride_order=self.a_desc.stride_order, + assumed_align=align, + ) + c_tensor = self._make_fake_cute_compact_tensor( + dtype=self.c_desc.dtype, + shape=(valid_m, *self.c_desc.shape[1:]), + stride_order=self.c_desc.stride_order, + ) + d_tensor = self._make_fake_cute_compact_tensor( + dtype=self.d_desc.dtype, + shape=(valid_m, *self.d_desc.shape[1:]), + stride_order=self.d_desc.stride_order, + ) + d_col_tensor = self._make_fake_cute_compact_tensor( + dtype=self.d_col_desc.dtype, + shape=(valid_m, *self.d_col_desc.shape[1:]), + stride_order=self.d_col_desc.stride_order, + ) + + tensor_m_128 = cute.sym_int() + stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) + sfa_shape = list(self.sfa_desc.shape) + sfa_shape[2] = tensor_m_128 + sfa_stride = list(self.sfa_desc.stride) + sfa_stride[5] = stride_tensor_m_128 + sfa_tensor = self._make_fake_cute_tensor( + dtype=self.sfa_desc.dtype, + shape=tuple(sfa_shape), + stride=tuple(sfa_stride), + assumed_align=16, + ) + sfd_row_tensor = None + if self.sfd_row_desc is not None: + stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) + sfd_row_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_row_desc.dtype, + shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), + stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), + assumed_align=16, + ) + sfd_col_tensor = None + if self.sfd_col_desc is not None: + rest_m = cute.sym_int(divisibility=1) + stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) + stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) + sfd_col_tensor = self._make_fake_cute_tensor( + dtype=self.sfd_col_desc.dtype, + shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), + stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), + assumed_align=16, + ) + amax_tensor = self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16) + norm_const_tensor_cute = self._make_fake_cute_tensor_from_desc(self.norm_const_desc, assumed_align=16) + padded_offsets_tensor = self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16) + alpha_tensor = self._make_fake_cute_tensor_from_desc(self.alpha_desc, assumed_align=16) + prob_tensor = None + if self.prob_desc is not None: + prob_tensor = self._make_fake_cute_tensor( + dtype=self.prob_desc.dtype, + shape=(valid_m, *self.prob_desc.shape[1:]), + stride=self.prob_desc.stride, + assumed_align=16, + ) + bias_tensor = self._make_fake_cute_tensor_from_desc(self.bias_desc, assumed_align=16) + + # Compile-time pointer placeholders + b_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") + sfb_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") + b_ptrs_cute = from_dlpack(b_ptrs_placeholder, assumed_align=8).iterator + sfb_ptrs_cute = from_dlpack(sfb_ptrs_placeholder, assumed_align=8).iterator + + workspace_ptr_cute = from_dlpack(self._workspace, assumed_align=128).iterator + + # linear_offset, geglu_alpha, glu_clamp_max, and glu_clamp_min are runtime + # cutlass.Float32 (not Constexpr), so the compile-time placeholders below are + # irrelevant -- the values passed through tensor_api() at execute() time are + # what the kernel actually uses. + self._logger.debug("Compiling discrete grouped GEMM GLU kernel") + _compiled_kernel = cute.compile( + gemm_glu, + a_tensor, + b_ptrs_cute, + sfb_ptrs_cute, + cutlass.Int32(n), + cutlass.Int32(k), + cutlass.Int64(b_stride_size), + b_major_mode, + workspace_ptr_cute, + c_tensor, + d_tensor, + d_col_tensor, + sfa_tensor, + sfd_row_tensor, + sfd_col_tensor, + amax_tensor, + norm_const_tensor_cute, + padded_offsets_tensor, + alpha_tensor, + prob_tensor, + bias_tensor, + max_active_clusters, + fake_stream, + lambda x: x, # epilogue_op (Constexpr, baked in) + cutlass.Float32(0.0), + cutlass.Float32(1.702), + cutlass.Float32(7.0), + cutlass.Float32(-7.0), + options="--enable-tvm-ffi", + ) + + self._n = n + self._k = k + self._b_stride_size = b_stride_size + + # Cache constant values for execute() closure + cached_workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator + cached_n = cutlass.Int32(self._n) + cached_k = cutlass.Int32(self._k) + cached_b_stride = cutlass.Int64(self._b_stride_size) + + def tensor_api( + a_tensor: torch.Tensor, + b_ptrs_device: torch.Tensor, + sfb_ptrs_device: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + d_col_tensor: Optional[torch.Tensor], + sfa_tensor: torch.Tensor, + sfd_row_tensor: Optional[torch.Tensor], + sfd_col_tensor: Optional[torch.Tensor], + amax_tensor: Optional[torch.Tensor], + norm_const_tensor: Optional[torch.Tensor], + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + prob_tensor: Optional[torch.Tensor], + bias_tensor: Optional[torch.Tensor], + stream: cuda.CUstream, + linear_offset: float = 0.0, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + ) -> None: + norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") + b_ptrs_addr = int(b_ptrs_device.data_ptr()) + sfb_ptrs_addr = int(sfb_ptrs_device.data_ptr()) + + _compiled_kernel( + a_tensor, + b_ptrs_addr, + sfb_ptrs_addr, + cached_n, + cached_k, + cached_b_stride, + cached_workspace_ptr, + c_tensor, + d_tensor, + d_col_tensor, + sfa_tensor, + sfd_row_tensor, + sfd_col_tensor, + amax_tensor, + norm_const_tensor, + padded_offsets, + alpha_tensor, + prob_tensor, + bias_tensor, + stream, + cutlass.Float32(linear_offset), + cutlass.Float32(geglu_alpha), + cutlass.Float32(glu_clamp_max), + cutlass.Float32(glu_clamp_min), + ) + + self._compiled_kernel = tensor_api + + # --------------------------------------------------------------------- # + # execute + # --------------------------------------------------------------------- # + + def execute( + self, + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + sfa_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + # Dense mode: + b_tensor: Optional[torch.Tensor] = None, + sfb_tensor: Optional[torch.Tensor] = None, + bias_tensor: Optional[torch.Tensor] = None, + # Discrete mode: + b_ptrs: Optional[torch.Tensor] = None, + sfb_ptrs: Optional[torch.Tensor] = None, + # Optional: + d_col_tensor: Optional[torch.Tensor] = None, + sfd_row_tensor: Optional[torch.Tensor] = None, + sfd_col_tensor: Optional[torch.Tensor] = None, + amax_tensor: Optional[torch.Tensor] = None, + norm_const_tensor: Optional[torch.Tensor] = None, + prob_tensor: Optional[torch.Tensor] = None, + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + """Execute the compiled kernel. + + For dense mode, supply ``b_tensor`` and ``sfb_tensor``. + For discrete mode, supply ``b_ptrs`` and ``sfb_ptrs``. + + :param a_tensor: Input A tensor + :param c_tensor: Intermediate C tensor + :param d_tensor: Output D tensor + :param sfa_tensor: Scale factor A + :param padded_offsets: End offset per expert after padding + :param alpha_tensor: Per-group scaling factors + :param b_tensor: (Dense) Input B tensor (weights) + :param sfb_tensor: (Dense) Scale factor B + :param bias_tensor: Optional bias tensor with shape (n, l) and stride (1, n). + Bias fusion is specialized at compile time: if ``sample_bias`` was omitted + at construction, ``bias_tensor`` must also be omitted at execute time. + :param b_ptrs: (Discrete) 1-D int64 device tensor of per-expert B data pointers + :param sfb_ptrs: (Discrete) 1-D int64 device tensor of per-expert SFB data pointers + :param d_col_tensor: Optional column-quantized output + :param sfd_row_tensor: Optional row scale factor D + :param sfd_col_tensor: Optional column scale factor D + :param amax_tensor: Optional amax tensor + :param norm_const_tensor: Optional normalization constant + :param prob_tensor: Optional probability tensor + :param linear_offset: Linear offset applied to the up branch in the + ``act_func == "geglu"`` activation. Ignored when + ``act_func == "swiglu"``. When ``None`` (default), the offset is + chosen based on ``act_func`` for backwards compatibility: + ``1.0`` for ``"geglu"`` and ``0.0`` for ``"swiglu"``. + :param geglu_alpha: Pre-sigmoid scaling factor for the GeGLU activation. + The fused activation is + ``out = (clamp(up, glu_clamp_min, glu_clamp_max) + linear_offset) + * silu(geglu_alpha * clamp(gate, max=glu_clamp_max))``. + Defaults to ``1.702`` (GPT-OSS / scaled-GeGLU). Ignored when + ``act_func == "swiglu"``. + :param glu_clamp_max: Upper clamp limit applied to both ``gate`` and + ``up`` before the activation. Default ``7.0``. Ignored when + ``act_func == "swiglu"``. + :param glu_clamp_min: Lower clamp limit applied only to ``up`` (the + kernel never lower-clamps ``gate``). Default ``-7.0``. Ignored + when ``act_func == "swiglu"``. + :param current_stream: CUDA stream + """ + self._logger.debug("Entering execute") + current_stream = self._get_default_stream(current_stream) + + if a_tensor.shape[0] == 0: + self._logger.debug("execute: valid_m is zero, skipping kernel execution") + return + self._runtime_error_if( + self._compiled_kernel is None, + "Kernel not compiled; call compile() first", + ) + + # Resolve linear_offset default: None -> activation-derived legacy value + # (1.0 for geglu, 0.0 for swiglu) for backwards compatibility with callers + # that pre-date the explicit linear_offset kwarg. + if linear_offset is None: + linear_offset = 1.0 if self.act_func == "geglu" else 0.0 + + self._logger.debug("Executing grouped GEMM GLU kernel") + if self._has_bias: + self._value_error_if( + bias_tensor is None, + "bias_tensor must be provided at execute() when the API was compiled with sample_bias", + ) + + if self.weight_mode == MoEWeightMode.DENSE: + self._compiled_kernel( + a_tensor=a_tensor, + b_tensor=b_tensor, + c_tensor=c_tensor, + d_tensor=d_tensor, + d_col_tensor=d_col_tensor, + sfa_tensor=sfa_tensor, + sfb_tensor=sfb_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + amax_tensor=amax_tensor, + norm_const_tensor=norm_const_tensor, + padded_offsets=padded_offsets, + alpha_tensor=alpha_tensor, + bias_tensor=bias_tensor, + prob_tensor=prob_tensor, + stream=current_stream, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + ) + else: + self._compiled_kernel( + a_tensor=a_tensor, + b_ptrs_device=b_ptrs, + sfb_ptrs_device=sfb_ptrs, + c_tensor=c_tensor, + d_tensor=d_tensor, + d_col_tensor=d_col_tensor, + sfa_tensor=sfa_tensor, + sfd_row_tensor=sfd_row_tensor, + sfd_col_tensor=sfd_col_tensor, + amax_tensor=amax_tensor, + norm_const_tensor=norm_const_tensor, + padded_offsets=padded_offsets, + alpha_tensor=alpha_tensor, + prob_tensor=prob_tensor, + bias_tensor=bias_tensor, + stream=current_stream, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + ) + + self._logger.debug("Execute completed") + + +__all__ = ["GroupedGemmGluBlockScaledAPI"] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu/api.py b/python/cudnn/grouped_gemm/grouped_gemm_glu/api.py index b814b046f..f6b50567c 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_glu/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/api.py @@ -43,59 +43,98 @@ at execution time. """ -from .moe_blockscaled_grouped_gemm_glu_bias import BlockScaledMoEGroupedGemmGluBiasKernel +from dataclasses import dataclass, replace + +from ..grouped_gemm_utils import ( + GroupedGemmBackend, + backend_cache_key, + select_grouped_gemm_backend, +) from ..moe_utils import MoEWeightMode from cuda.bindings import driver as cuda import logging import os import torch -from typing import Tuple, Optional +from typing import Any, Tuple, Optional, overload -import cutlass -import cutlass.cute as cute -from cutlass.cute.nvgpu import OperandMajorMode -from cutlass.cute.runtime import from_dlpack, make_fake_stream +from cudnn.api_base import APIBase, TupleDict, ceil_div -from cudnn.datatypes import _convert_to_cutlass_data_type -from cudnn.api_base import APIBase, TupleDict, ceil_div, is_power_of_2 +_BLOCK_SCALED_DTYPE_PAIRS = { + (dtype, dtype) + for dtype in ( + torch.float4_e2m1fn_x2, + torch.uint8, + torch.float8_e5m2, + torch.float8_e4m3fn, + ) +} + + +from ._bf16_api import GroupedGemmGluBf16API +from ._blockscaled_api import GroupedGemmGluBlockScaledAPI + + +@dataclass(frozen=True) +class GluCall: + """Immutable normalized input for GLU dispatch, allocation, and caching.""" + + a_tensor: torch.Tensor + sfa_tensor: Optional[torch.Tensor] + padded_offsets: torch.Tensor + alpha_tensor: torch.Tensor + b_tensor: Optional[torch.Tensor] = None + sfb_tensor: Optional[torch.Tensor] = None + bias_tensor: Optional[torch.Tensor] = None + b_ptrs: Optional[torch.Tensor] = None + sfb_ptrs: Optional[torch.Tensor] = None + n: Optional[int] = None + b_dtype: Optional[torch.dtype] = None + b_major: str = "k" + norm_const_tensor: Optional[torch.Tensor] = None + prob_tensor: Optional[torch.Tensor] = None + acc_dtype: torch.dtype = torch.float32 + c_dtype: torch.dtype = torch.bfloat16 + d_dtype: torch.dtype = torch.bfloat16 + cd_major: str = "n" + mma_tiler_mn: Tuple[int, int] = (256, 256) + cluster_shape_mn: Optional[Tuple[int, int]] = None + sf_vec_size: int = 16 + vector_f32: bool = False + m_aligned: int = 256 + discrete_col_sfd: bool = False + act_func: str = "swiglu" + linear_offset: Optional[float] = None + geglu_alpha: float = 1.702 + glu_clamp_max: float = 7.0 + glu_clamp_min: float = -7.0 + use_dynamic_sched: bool = False + current_stream: Optional[cuda.CUstream] = None + generate_c: bool = False + weight_mode: Optional[MoEWeightMode] = None + b_shape: Optional[Tuple[int, ...]] = None + num_experts: Optional[int] = None class GroupedGemmGluSm100(APIBase): - """Unified API for grouped GEMM GLU forward operation on SM100+ GPUs. - - This kernel performs block-scaled grouped GEMM with GLU activation - (SwiGLU or GeGLU), designed for MoE workloads. It supports both dense - (contiguous) and discrete (per-expert pointer) weight layouts through - the ``BlockScaledMoEGroupedGemmGluBiasKernel``. - - Weight mode is auto-detected from the constructor arguments: + """Stable public facade that selects the GLU backend during support checking.""" - - **Dense**: provide ``sample_b`` and ``sample_sfb``. - - **Discrete**: provide ``num_experts``, ``b_shape``, and ``b_dtype``. - - Example:: - - # Dense mode - api = GroupedGemmGluSm100( - sample_a=a, sample_b=b, sample_c=c, sample_d=d, - sample_sfa=sfa, sample_sfb=sfb, - sample_padded_offsets=offsets, sample_alpha=alpha, - sample_d_col=d_col, - ) - - # Discrete mode - api = GroupedGemmGluSm100( - sample_a=a, num_experts=8, b_shape=(n, k), b_dtype=torch.uint8, - sample_c=c, sample_d=d, sample_sfa=sfa, - sample_padded_offsets=offsets, sample_alpha=alpha, - sample_d_col=d_col, - ) - - api.check_support() - api.compile() - api.execute(...) - """ + # BF16 implementation + @overload + def __init__( + self, + sample_a: torch.Tensor, + sample_c: torch.Tensor, + sample_d: torch.Tensor, + sample_sfa: None, + sample_padded_offsets: torch.Tensor, + sample_alpha: torch.Tensor, + sample_d_col: None, + *args: Any, + **kwargs: Any, + ) -> None: ... + # Block-scaled implementation + @overload def __init__( self, sample_a: torch.Tensor, @@ -104,22 +143,31 @@ def __init__( sample_sfa: torch.Tensor, sample_padded_offsets: torch.Tensor, sample_alpha: torch.Tensor, - sample_d_col: torch.Tensor, - # Dense mode (contiguous) -- provide these. sample_bias is optional: + sample_d_col: Optional[torch.Tensor], + *args: Any, + **kwargs: Any, + ) -> None: ... + + def __init__( + self, + sample_a: torch.Tensor, + sample_c: torch.Tensor, + sample_d: torch.Tensor, + sample_sfa: Optional[torch.Tensor], + sample_padded_offsets: torch.Tensor, + sample_alpha: torch.Tensor, + sample_d_col: Optional[torch.Tensor], sample_b: Optional[torch.Tensor] = None, sample_sfb: Optional[torch.Tensor] = None, sample_bias: Optional[torch.Tensor] = None, - # Discrete mode -- provide these instead: num_experts: Optional[int] = None, b_shape: Optional[Tuple[int, ...]] = None, b_dtype: Optional[torch.dtype] = None, - # Optional quantization output arguments sample_sfd_row: Optional[torch.Tensor] = None, sample_sfd_col: Optional[torch.Tensor] = None, sample_amax: Optional[torch.Tensor] = None, sample_norm_const: Optional[torch.Tensor] = None, sample_prob: Optional[torch.Tensor] = None, - # Configuration acc_dtype: torch.dtype = torch.float32, mma_tiler_mn: Tuple[int, int] = (256, 256), cluster_shape_mn: Optional[Tuple[int, int]] = None, @@ -130,991 +178,103 @@ def __init__( act_func: str = "swiglu", b_major: str = "k", use_dynamic_sched: bool = False, - ): - """Initialize the GroupedGemmGluSm100 API. - - :param sample_a: Sample A tensor (valid_m, k, 1) - :param sample_c: Sample C tensor for intermediate storage - :param sample_d: Sample D output tensor (valid_m, n/2, 1) after GLU - :param sample_sfa: Sample scale factor A tensor - :param sample_padded_offsets: End offset for each expert after padding - :param sample_alpha: Per-group alpha scaling factors - :param sample_d_col: Column-quantized D tensor - :param sample_b: (Dense) Sample B tensor (n, k, l) - :param sample_sfb: (Dense) Sample scale factor B tensor - :param sample_bias: Optional bias tensor with shape (n, l) or (n, expert_cnt), stride (1, n). - Dense mode supports fp16/bfloat16/float32 bias; discrete mode supports fp16/bfloat16 bias. - :param num_experts: (Discrete) Number of experts - :param b_shape: (Discrete) Shape of a single expert B tensor, e.g. (n, k) - :param b_dtype: (Discrete) Data type of B tensors - :param sample_sfd_row: Optional row scale factor for D - :param sample_sfd_col: Optional column scale factor for D - :param sample_amax: Optional amax tensor for quantization - :param sample_norm_const: Optional normalization constant - :param sample_prob: Optional probability tensor for gating - :param acc_dtype: Accumulator data type - :param mma_tiler_mn: MMA tiler shape (M, N) - :param cluster_shape_mn: Cluster shape (M, N) - :param sf_vec_size: Scale factor vector size - :param vector_f32: Use vectorized f32 operations - :param m_aligned: Alignment for group M dimension - :param discrete_col_sfd: Generate discrete col-major scale factor tensor - :param act_func: Activation function, one of "swiglu" or "geglu" - :param b_major: Major dimension for B tensor, one of "k" or "n" - :param use_dynamic_sched: Enable dynamic tile scheduling for load balancing - """ + generate_c: bool = False, + ) -> None: super().__init__() - - self._warn_experimental_api() - self._logger.debug("Entering __init__") - - # ---- Weight mode auto-detection ---- - if sample_b is not None and num_experts is None: - self.weight_mode = MoEWeightMode.DENSE - if sample_sfb is None: - raise ValueError("sample_sfb is required when sample_b is provided (dense mode)") - elif num_experts is not None and sample_b is None: - self.weight_mode = MoEWeightMode.DISCRETE - if b_shape is None or b_dtype is None: - raise ValueError("b_shape and b_dtype are required in discrete mode") - else: - raise ValueError("Provide either (sample_b, sample_sfb) for dense mode " "or (num_experts, b_shape, b_dtype) for discrete mode, but not both.") - - # ---- Common tensor descriptors ---- - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") - self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") - self.d_desc = self._make_tensor_desc(sample_d, name="sample_d") - self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") - self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") - self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") - - self.d_col_desc = self._make_tensor_desc(sample_d_col, name="sample_d_col") - self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias") - self.sfd_row_desc = self._make_tensor_desc(sample_sfd_row, name="sample_sfd_row") - self.sfd_col_desc = self._make_tensor_desc(sample_sfd_col, name="sample_sfd_col") - self.amax_desc = self._make_tensor_desc(sample_amax, name="sample_amax") - self.norm_const_desc = self._unpad_tensor_to_ndim( - self._make_tensor_desc(sample_norm_const, name="sample_norm_const"), - 1, - "norm_const", - ) - self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") - - # ---- Mode-specific state ---- - if self.weight_mode == MoEWeightMode.DENSE: - self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") - self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb") - self.expert_cnt = self.padded_offsets_desc.shape[0] - else: - self._value_error_if(num_experts == 0, "num_experts must be > 0") - self.expert_cnt = num_experts - self.b_shape = b_shape - self.b_dtype = b_dtype - self.b_major = b_major - self._value_error_if( - self.padded_offsets_desc.shape[0] != self.expert_cnt, - f"padded_offsets length ({self.padded_offsets_desc.shape[0]}) " f"must equal num_experts ({self.expert_cnt})", - ) - - # ---- Configuration ---- - self.acc_dtype = acc_dtype - self.mma_tiler_mn = mma_tiler_mn - self.use_2cta_instrs = mma_tiler_mn[0] == 256 - if cluster_shape_mn is None: - self.cluster_shape_mn = (2, 1) if self.use_2cta_instrs else (1, 1) - else: - self.cluster_shape_mn = cluster_shape_mn - self.sf_vec_size = sf_vec_size - self.vector_f32 = vector_f32 - self.m_aligned = m_aligned - self.discrete_col_sfd = discrete_col_sfd - self.act_func = act_func - if self.weight_mode == MoEWeightMode.DENSE: - self.b_major = b_major # stored for both modes - - self.use_dynamic_sched = use_dynamic_sched - - self._interpret_uint8_as_fp4x2 = True - self._has_bias = self.bias_desc is not None - self._kernel = BlockScaledMoEGroupedGemmGluBiasKernel - - self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) - self._logger.debug(f"setting num_cluster_overlap_margin: {self.num_cluster_overlap_margin}") - - self._workspace = None - - self._logger.debug("__init__ completed") - - # --------------------------------------------------------------------- # - # check_support - # --------------------------------------------------------------------- # + self._pending_init_kwargs = dict(locals()) + self._pending_init_kwargs.pop("self") + self._pending_init_kwargs.pop("__class__", None) + self._implementation = None def check_support(self) -> bool: - """Check if the kernel configuration is supported. - - :return: True if supported, raises exception otherwise - """ - self._logger.debug("Entering check_support") - - # ---- SFD group validation ---- - all_none = all(x is None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) - all_provided = all(x is not None for x in [self.sfd_row_desc, self.sfd_col_desc, self.norm_const_desc]) - self._value_error_if( - not (all_none or all_provided), - "sfd_row_desc, sfd_col_desc, and norm_const_desc must be all None or all not None", - ) - self.generate_sfd = all_provided - if self.discrete_col_sfd and not self.generate_sfd: - self._logger.warning("discrete_col_sfd is True but generate_sfd is False, discrete_col_sfd will be ignored") - self.discrete_col_sfd = False - - # ---- Shapes and strides ---- - self._logger.debug("Checking tensor shapes and strides") - tensor_m, k, _one = self._tensor_shape(self.a_desc, name="sample_a") - - if self.weight_mode == MoEWeightMode.DENSE: - n, _, l = self._tensor_shape(self.b_desc, name="sample_b") - else: - # Discrete: extract n, k from b_shape - if len(self.b_shape) == 2: - n, b_k = self.b_shape - else: - n, b_k, _ = self.b_shape - self._value_error_if(b_k != k, f"B K dimension ({b_k}) must match A K dimension ({k})") - l = self.expert_cnt # for shape checks that use l - - _, n_2, _one = self._tensor_shape(self.d_desc, name="sample_d") - - self._value_error_if( - n % 64 != 0, - f"N must be divisible by 64 for GLU (two consecutive 32-column blocks), got N={n}", - ) - - self._check_tensor_shape(self.a_desc, (tensor_m, k, 1), "A") - if self.weight_mode == MoEWeightMode.DENSE: - self._check_tensor_shape(self.b_desc, (n, k, l), "B") - self._check_tensor_shape(self.c_desc, (tensor_m, n, 1), "C") - self._check_tensor_shape(self.d_desc, (tensor_m, n // 2, 1), "D") - self._check_tensor_shape(self.d_col_desc, (tensor_m, n // 2, 1), "D_col") - self._check_tensor_shape(self.bias_desc, (n, l), "bias") - - rest_k = ceil_div(ceil_div(k, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfa_desc, (32, 4, ceil_div(tensor_m, 128), 4, rest_k, 1), "SFA") - if self.weight_mode == MoEWeightMode.DENSE: - self._check_tensor_shape(self.sfb_desc, (32, 4, ceil_div(n, 128), 4, rest_k, l), "SFB") - - rest_n2 = ceil_div(ceil_div(n // 2, self.sf_vec_size), 4) - self._check_tensor_shape( - self.sfd_row_desc, - (32, 4, ceil_div(tensor_m, 128), 4, rest_n2, 1), - "SFD_row", - ) - rest_m = ceil_div(ceil_div(tensor_m, self.sf_vec_size), 4) - self._check_tensor_shape(self.sfd_col_desc, (32, 4, ceil_div(n // 2, 128), 4, rest_m, 1), "SFD_col") - - self._check_tensor_shape(self.alpha_desc, (self.expert_cnt,), "alpha") - self._check_tensor_shape(self.prob_desc, (tensor_m, 1, 1), "prob") - self._check_tensor_shape(self.amax_desc, (self.expert_cnt, 1), "amax") - self._check_tensor_shape(self.norm_const_desc, (1,), "norm_const") - self._check_tensor_shape(self.padded_offsets_desc, (self.expert_cnt,), "padded_offsets") - - # Strides - _ = self._check_tensor_stride( - self.a_desc, - stride=[(k, 1, tensor_m * k)], - extra_error_msg="A must have k-major layout", - ) - if self.weight_mode == MoEWeightMode.DENSE: - _ = self._check_tensor_stride( - self.b_desc, - stride=[(k, 1, n * k)], - extra_error_msg="B must have k-major layout", - ) - _ = self._check_tensor_stride( - self.c_desc, - stride=[(n, 1, tensor_m * n)], - extra_error_msg="C must have n-major layout", - ) - _ = self._check_tensor_stride( - self.d_desc, - stride=[(n_2, 1, tensor_m * n_2)], - extra_error_msg="D must have n-major layout", - ) - _ = self._check_tensor_stride( - self.d_col_desc, - stride=[(n_2, 1, tensor_m * n_2)], - extra_error_msg="D_col must have n-major layout", - ) - _ = self._check_tensor_stride( - self.bias_desc, - stride=[(1, n)], - ) - - # ---- Data types ---- - self._logger.debug("Checking data types") - self.ab_dtype = self._check_dtype( - self.a_desc, - dtype=[ - torch.float4_e2m1fn_x2, - torch.uint8, - torch.float8_e5m2, - torch.float8_e4m3fn, - ], - name="A/B", - ) - if self.weight_mode == MoEWeightMode.DENSE: - self._check_dtype( - self.b_desc, - dtype=self.ab_dtype, - name="B", - extra_error_msg="B must have the same dtype as A", - ) - self._check_dtype( - self.bias_desc, - dtype=[torch.bfloat16, torch.float16, torch.float32], - name="bias", - extra_error_msg="bias must be fp16, bfloat16, or float32", - ) - else: - self._value_error_if( - self.b_dtype != self.ab_dtype, - f"b_dtype ({self.b_dtype}) must match A dtype ({self.ab_dtype})", - ) - self._check_dtype( - self.bias_desc, - dtype=[torch.bfloat16, torch.float16], - name="bias", - extra_error_msg="bias must be fp16 or bfloat16 in discrete mode", - ) - - self.sf_dtype = self._check_dtype( - self.sfa_desc, - dtype=[torch.float8_e8m0fnu, torch.float8_e4m3fn], - name="SFA/SFB/SFD", - ) - if self.weight_mode == MoEWeightMode.DENSE: - self._check_dtype( - self.sfb_desc, - dtype=self.sf_dtype, - name="SFB", - extra_error_msg="SFB must have the same dtype as SFA", - ) - self._check_dtype( - self.sfd_row_desc, - dtype=self.sf_dtype, - name="SFD_row", - extra_error_msg="SFD_row must have the same dtype as SFA", - ) - self._check_dtype( - self.sfd_col_desc, - dtype=self.sf_dtype, - name="SFD_col", - extra_error_msg="SFD_col must have the same dtype as SFA", - ) - - self._value_error_if( - self.sf_vec_size not in [16, 32], - f"sf_vec_size must be 16 or 32, got {self.sf_vec_size}", - ) - self._value_error_if( - self.sf_dtype in [torch.float8_e4m3fn] and self.sf_vec_size == 32, - f"sf_dtype {self.sf_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", - ) - self._value_error_if( - self._is_fp8(self.ab_dtype) and self.sf_vec_size == 16, - f"ab_dtype {self.ab_dtype} and sf_vec_size {self.sf_vec_size} combination is not supported", - ) - - self._check_dtype( - self.acc_dtype, - dtype=torch.float32, - name="Accumulator", - extra_error_msg="Accumulator must be float32", - ) - self.c_dtype = self._check_dtype( - self.c_desc, - dtype=[ - torch.float32, - torch.float16, - torch.bfloat16, - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float4_e2m1fn_x2, - ], - name="C", - ) - - if self._is_fp4x2(self.ab_dtype): - self.d_dtype = self._check_dtype( - self.d_desc, - dtype=[torch.float16, torch.bfloat16, torch.float32], - name="D", - extra_error_msg="D must be fp16, bf16, or float32 when ab_dtype is fp4", - ) - else: - self.d_dtype = self._check_dtype( - self.d_desc, - dtype=[ - torch.float16, - torch.bfloat16, - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.float4_e2m1fn_x2, - ], - name="D", - ) - self._check_dtype( - self.d_col_desc, - dtype=self.d_dtype, - name="D_col", - extra_error_msg="D_col must have the same dtype as D", - ) - - self._not_implemented_error_if( - self.bias_desc is None and self._is_fp4x2(self.ab_dtype) and self.sf_vec_size == 16 and self.d_dtype == torch.float32, - "Invalid configuration: fp4 ab_dtype, sf_vec_size 16, d_dtype float32 is not supported. " "Please use sf_vec_size 32 or d_dtype bf16 instead", - ) - - # ---- Activation function validation (both modes) ---- - self._value_error_if( - self.act_func not in ["swiglu", "geglu"], - f"act_func must be 'swiglu' or 'geglu', got {self.act_func}", - ) - - # ---- Discrete-mode-specific validation ---- - if self.weight_mode == MoEWeightMode.DISCRETE: - self._value_error_if( - self.b_major not in ["k", "n"], - f"b_major must be 'k' or 'n', got {self.b_major}", + if self._implementation is None: + kwargs = self._pending_init_kwargs + defining_b_dtype = kwargs["sample_b"].dtype if kwargs["sample_b"] is not None else kwargs["b_dtype"] + backend = select_grouped_gemm_backend( + operation="grouped_gemm_glu_sm100", + a_dtype=kwargs["sample_a"].dtype, + b_dtype=defining_b_dtype, + scale_controls=( + ("sample_sfa", kwargs["sample_sfa"]), + ("sample_sfb", kwargs["sample_sfb"]), + ("sample_d_col", kwargs["sample_d_col"]), + ("sample_sfd_row", kwargs["sample_sfd_row"]), + ("sample_sfd_col", kwargs["sample_sfd_col"]), + ("sample_amax", kwargs["sample_amax"]), + ("sample_norm_const", kwargs["sample_norm_const"]), + ("sf_vec_size", kwargs["sf_vec_size"] if kwargs["sf_vec_size"] != 16 else None), + ("discrete_col_sfd", kwargs["discrete_col_sfd"] if kwargs["discrete_col_sfd"] else None), + ), + block_scaled_dtype_pairs=_BLOCK_SCALED_DTYPE_PAIRS, ) - self._value_error_if( - self._is_fp4x2(self.ab_dtype) and self.b_major != "k", - "b_major must be 'k' when ab_dtype is fp4", - ) - - # ---- MMA tile / cluster shape ---- - self._logger.debug("Checking MMA tile shape and cluster shape") - self._value_error_if( - not self.use_2cta_instrs and self.mma_tiler_mn[0] != 128, - f"MMA tiler M must be 128 when use_2cta_instrs=False, got {self.mma_tiler_mn[0]}", - ) - self._value_error_if( - self.use_2cta_instrs and self.mma_tiler_mn[0] != 256, - f"MMA tiler M must be 256 when use_2cta_instrs=True, got {self.mma_tiler_mn[0]}", - ) - self._value_error_if( - self.mma_tiler_mn[1] != 256, - f"MMA tiler N must be 256, got {self.mma_tiler_mn[1]}", - ) - self._value_error_if( - self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0, - f"cluster_shape_mn[0] must be divisible by 2 when use_2cta_instrs=True, got {self.cluster_shape_mn[0]}", - ) - self._value_error_if( - not ( - self.cluster_shape_mn[0] * self.cluster_shape_mn[1] <= 16 - and self.cluster_shape_mn[0] > 0 - and self.cluster_shape_mn[1] > 0 - and self.cluster_shape_mn[0] <= 4 - and self.cluster_shape_mn[1] <= 4 - and is_power_of_2(self.cluster_shape_mn[0]) - and is_power_of_2(self.cluster_shape_mn[1]) - ), - f"Invalid cluster shape: expected values to be powers of 2 and product <= 16, got {self.cluster_shape_mn}", - ) - cluster_tiler_m = (self.cluster_shape_mn[0] // (2 if self.use_2cta_instrs else 1)) * self.mma_tiler_mn[0] - self._value_error_if( - cluster_tiler_m not in [128, 256], - f"Invalid cluster tiler shape: expected cluster_tiler_m in {{128, 256}}, got {cluster_tiler_m}", - ) - self._value_error_if( - self.m_aligned % self.mma_tiler_mn[0] != 0, - f"m_aligned must be divisible by mma_tiler_mn[0], got {self.m_aligned} % {self.mma_tiler_mn[0]} != 0", - ) - self._value_error_if( - self.m_aligned != BlockScaledMoEGroupedGemmGluBiasKernel.FIX_PAD_SIZE, - f"m_aligned must be {BlockScaledMoEGroupedGemmGluBiasKernel.FIX_PAD_SIZE} (FIX_PAD_SIZE), got {self.m_aligned}", - ) - - # ---- Tensor alignment ---- - self._logger.debug("Checking tensor alignment") - - def check_contiguous_16B_alignment(dtype, stride_order, tensor_shape): - is_mode0_major = stride_order == (0, 1, 2) - major_mode_idx = 0 if is_mode0_major else 1 - num_major_elements = tensor_shape[major_mode_idx] - num_contiguous_elements = 16 * 8 // (_convert_to_cutlass_data_type(dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2).width) - return num_major_elements % num_contiguous_elements == 0 - - if self.weight_mode == MoEWeightMode.DENSE: - b_stride_order_for_check = self.b_desc.stride_order - b_shape_for_check = (n, k, l) - else: - b_stride_order_for_check = (0, 1, 2) if self.b_major == "n" else (1, 0, 2) - b_shape_for_check = (n, k, 1) - - self._value_error_if( - not ( - check_contiguous_16B_alignment(self.ab_dtype, self.a_desc.stride_order, (tensor_m, k, l)) - and check_contiguous_16B_alignment(self.ab_dtype, b_stride_order_for_check, b_shape_for_check) - and check_contiguous_16B_alignment(self.d_dtype, self.d_desc.stride_order, (tensor_m, n_2, 1)) - ), - "Invalid tensor alignment: tensors must be 16B aligned", - ) - - # ---- Expert count limit ---- - self._value_error_if( - self.expert_cnt > 1024, - f"expert_cnt must be <= 1024, got {self.expert_cnt}", - ) - - # ---- Disabled configurations ---- - self._not_implemented_error_if( - (self._is_fp8(self.ab_dtype)) and (self.mma_tiler_mn[1] == 128) and (self._is_fp8(self.d_dtype)), - "Invalid configuration: fp8 ab_dtype with mma_tiler_mn[1] == 128 and fp8 d_dtype is not supported. " "Please use mma_tiler_mn[1] == 256 instead", - ) - self._not_implemented_error_if( - self._is_fp4x2(self.ab_dtype) and (self.c_dtype not in [torch.float16, torch.bfloat16]), - f"Invalid configuration: for fp4 ab_dtype, c_dtype must be float16 or bfloat16, got {self.c_dtype}", - ) - self._not_implemented_error_if(self._has_bias and self.mma_tiler_mn[1] != 256, "Bias fusion currently requires mma_tiler_mn[1] == 256") - - # ---- SM100+ check ---- - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is not available") - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) - compute_capability = major * 10 + minor - if compute_capability < 100: - raise RuntimeError(f"GroupedGemmGlu requires SM100+ compute capability, " f"but found SM{compute_capability} on device {device}") - - self._is_supported = True - self._logger.debug("check_support completed successfully") - return True - - # --------------------------------------------------------------------- # - # compile - # --------------------------------------------------------------------- # - - def compile(self) -> None: - """Compile the kernel.""" - self._logger.debug("Entering compile") - self._ensure_support_checked() - if self._compiled_kernel is not None: - self._logger.debug("Kernel already compiled; skipping recompilation") - return - if self.a_desc.shape[0] == 0: - self._logger.debug("sample valid_m is zero, skipping kernel compilation") - return - - # ---- Instantiate the unified kernel ---- - gemm_glu = self._kernel( - sf_vec_size=self.sf_vec_size, - acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), - use_2cta_instrs=self.use_2cta_instrs, - mma_tiler_mn=self.mma_tiler_mn, - cluster_shape_mn=self.cluster_shape_mn, - vectorized_f32=self.vector_f32, - generate_sfd=self.generate_sfd, - discrete_col_sfd=self.discrete_col_sfd, - expert_cnt=self.expert_cnt, - weight_mode=self.weight_mode, - act_func=self.act_func, - enable_bias=self._has_bias, - use_dynamic_sched=self.use_dynamic_sched, - ) - - hardware_info = cutlass.utils.HardwareInfo() - max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - max_active_clusters -= self.num_cluster_overlap_margin - self._value_error_if( - max_active_clusters <= 0, - "max_active_clusters must be > 0 after applying overlap margin; reduce CUDNNFE_CLUSTER_OVERLAP_MARGIN", - ) - fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) - - # ---- Allocate workspace ---- - workspace_bytes = gemm_glu.get_workspace_bytes() - self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device="cuda") - - if self.weight_mode == MoEWeightMode.DENSE: - self._compile_dense(gemm_glu, max_active_clusters, fake_stream) - else: - self._compile_discrete(gemm_glu, max_active_clusters, fake_stream) - - self._logger.debug("Kernel compiled successfully") - - # -- Dense compile path ------------------------------------------------- # - - def _compile_dense(self, gemm_glu, max_active_clusters, fake_stream) -> None: - """Compile for dense (contiguous) weight mode.""" - use_full_dynamic = os.environ.get("CUDNN_FE_GROUPED_GEMM_DYNAMIC_MNKL", "1") != "0" - - fake_workspace_ptr = cute.runtime.nullptr( - dtype=cutlass.Uint8, - assumed_align=128, - ) - - if not use_full_dynamic: - valid_m = cute.sym_int(divisibility=256) - - a_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.a_desc.dtype, - shape=(valid_m, *self.a_desc.shape[1:]), - stride_order=self.a_desc.stride_order, - ) - b_cute_fake = self._make_fake_cute_tensor_from_desc(self.b_desc, assumed_align=16) - c_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.c_desc.dtype, - shape=(valid_m, *self.c_desc.shape[1:]), - stride_order=self.c_desc.stride_order, - ) - d_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.d_desc.dtype, - shape=(valid_m, *self.d_desc.shape[1:]), - stride_order=self.d_desc.stride_order, - ) - d_col_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.d_col_desc.dtype, - shape=(valid_m, *self.d_col_desc.shape[1:]), - stride_order=self.d_col_desc.stride_order, - ) - - tensor_m_128 = cute.sym_int() - stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfa_cute_fake = self._make_fake_cute_tensor( - dtype=self.sfa_desc.dtype, - shape=(32, 4, tensor_m_128, 4, self.sfa_desc.shape[4], 1), - stride=(16, 4, self.sfa_desc.stride[2], 1, 512, stride_tensor_m_128), - ) - - sfb_cute_fake = self._make_fake_cute_tensor_from_desc(self.sfb_desc, assumed_align=16) - - prob_cute_fake = None - if self.prob_desc is not None: - prob_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.prob_desc.dtype, - shape=(valid_m, 1, 1), - stride_order=self.prob_desc.stride_order, - ) - - sfd_row_fake = None - sfd_col_fake = None - if self.sfd_row_desc is not None: - stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_row_fake = self._make_fake_cute_tensor( - dtype=self.sfd_row_desc.dtype, - shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), - stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), - ) - if self.sfd_col_desc is not None: - rest_m = cute.sym_int(divisibility=1) - stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) - stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_col_fake = self._make_fake_cute_tensor( - dtype=self.sfd_col_desc.dtype, - shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), - stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), + self.backend = backend + if backend is GroupedGemmBackend.BF16: + self._implementation = GroupedGemmGluBf16API( + sample_a=kwargs["sample_a"], + sample_c=kwargs["sample_c"], + sample_d=kwargs["sample_d"], + sample_padded_offsets=kwargs["sample_padded_offsets"], + sample_alpha=kwargs["sample_alpha"], + sample_b=kwargs["sample_b"], + sample_bias=kwargs["sample_bias"], + sample_prob=kwargs["sample_prob"], + num_experts=kwargs["num_experts"], + b_shape=kwargs["b_shape"], + b_dtype=kwargs["b_dtype"], + acc_dtype=kwargs["acc_dtype"], + mma_tiler_mn=kwargs["mma_tiler_mn"], + cluster_shape_mn=kwargs["cluster_shape_mn"], + vector_f32=kwargs["vector_f32"], + m_aligned=kwargs["m_aligned"], + generate_c=kwargs["generate_c"], + act_func=kwargs["act_func"], + b_major=kwargs["b_major"], + use_dynamic_sched=kwargs["use_dynamic_sched"], ) - bias_cute_fake = self._make_fake_cute_tensor_from_desc(self.bias_desc, assumed_align=16) - else: - valid_m = cute.sym_int(divisibility=256) - n_sym = cute.sym_int() - n_2_sym = cute.sym_int() - k_sym = cute.sym_int() - l_sym = cute.sym_int() - - a_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.a_desc.dtype, - shape=(valid_m, k_sym, 1), - stride_order=self.a_desc.stride_order, - dynamic_mode=self.a_desc.stride_order[0], - divisibility=32 if self._is_fp4x2(self.ab_dtype) else 16, - ) - b_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.b_desc.dtype, - shape=(n_sym, k_sym, l_sym), - stride_order=self.b_desc.stride_order, - dynamic_mode=self.b_desc.stride_order[0], - divisibility=32 if self._is_fp4x2(self.ab_dtype) else 16, - ) - c_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.c_desc.dtype, - shape=(valid_m, n_2_sym, 1), - stride_order=self.c_desc.stride_order, - dynamic_mode=self.c_desc.stride_order[0], - divisibility=8 if self._is_f16(self.c_desc.dtype) else 16, - ) - d_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.d_desc.dtype, - shape=(valid_m, n_2_sym, 1), - stride_order=self.d_desc.stride_order, - dynamic_mode=self.d_desc.stride_order[0], - divisibility=8 if self._is_f16(self.d_desc.dtype) else 16, - ) - d_col_cute_fake = self._make_fake_cute_compact_tensor( - dtype=self.d_col_desc.dtype, - shape=(valid_m, n_2_sym, 1), - stride_order=self.d_col_desc.stride_order, - dynamic_mode=self.d_col_desc.stride_order[0], - divisibility=8 if self._is_f16(self.d_col_desc.dtype) else 16, - ) - - tensor_m_128 = cute.sym_int() - rest_k = cute.sym_int() - stride_rest_k = cute.sym_int(divisibility=32 * 4 * 4) - stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfa_shape = list(self.sfa_desc.shape) - sfa_shape[2] = tensor_m_128 - sfa_shape[4] = rest_k - sfa_stride = list(self.sfa_desc.stride) - sfa_stride[2] = stride_rest_k - sfa_stride[5] = stride_tensor_m_128 - sfa_cute_fake = self._make_fake_cute_tensor( - dtype=self.sfa_desc.dtype, - shape=tuple(sfa_shape), - stride=tuple(sfa_stride), - ) - - tensor_n_128 = cute.sym_int() - stride_sfb_rest_k = cute.sym_int(divisibility=32 * 4 * 4) - stride_sfb_tensor_n_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfb_cute_fake = self._make_fake_cute_tensor( - dtype=self.sfb_desc.dtype, - shape=(32, 4, tensor_n_128, 4, rest_k, l_sym), - stride=(16, 4, stride_sfb_tensor_n_128, 1, 512, stride_sfb_rest_k), - ) - - prob_cute_fake = None - if self.prob_desc is not None: - prob_cute_fake = self._make_fake_cute_tensor( - dtype=self.prob_desc.dtype, - shape=(valid_m, *self.prob_desc.shape[1:]), - stride=self.prob_desc.stride, - ) - - sfd_row_fake = None - sfd_col_fake = None - if self.sfd_row_desc is not None: - rest_n2 = cute.sym_int() - stride_sfd_rest_n2 = cute.sym_int(divisibility=32 * 4 * 4) - stride_sfd_rest_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfd_row_fake = self._make_fake_cute_tensor( - dtype=self.sfd_row_desc.dtype, - shape=(32, 4, tensor_m_128, 4, rest_n2, 1), - stride=(16, 4, stride_sfd_rest_n2, 1, 512, stride_sfd_rest_tensor_m_128), - ) - if self.sfd_col_desc is not None: - tensor_n2_128 = cute.sym_int() - rest_m_dyn = cute.sym_int() - stride_sfd_rest_m = cute.sym_int(divisibility=32 * 4 * 4) - stride_sfd_n2 = cute.sym_int(divisibility=32 * 4 * 4) - sfd_col_fake = self._make_fake_cute_tensor( - dtype=self.sfd_col_desc.dtype, - shape=(32, 4, tensor_n2_128, 4, rest_m_dyn, 1), - stride=(16, 4, stride_sfd_rest_m, 1, 512, stride_sfd_n2), - ) - bias_cute_fake = None - if self.bias_desc is not None: - bias_cute_fake = self._make_fake_cute_tensor( - dtype=self.bias_desc.dtype, - shape=(n_sym, l_sym), - stride=(1, n_sym), - ) - - # Compile with keyword args (dense mode uses the unified __call__ positional order). - # linear_offset, geglu_alpha, glu_clamp_max, and glu_clamp_min are runtime - # cutlass.Float32 (not Constexpr), so the compile-time placeholders below are - # irrelevant -- the values passed through tensor_api() at execute() time are - # what the kernel actually uses. - _compiled_kernel = cute.compile( - gemm_glu, - a=a_cute_fake, - b=b_cute_fake, - sfb=sfb_cute_fake, - n=cutlass.Int32(0), - k=cutlass.Int32(0), - b_stride_size=cutlass.Int64(0), - b_major_mode=OperandMajorMode.K, - workspace_ptr=fake_workspace_ptr, - c=c_cute_fake, - d=d_cute_fake, - d_col=d_col_cute_fake, - sfa=sfa_cute_fake, - sfd_row_tensor=sfd_row_fake, - sfd_col_tensor=sfd_col_fake, - amax_tensor=self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16), - norm_const_tensor=self._make_fake_cute_tensor_from_desc(self.norm_const_desc, assumed_align=16), - padded_offsets=self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16), - alpha=self._make_fake_cute_tensor_from_desc(self.alpha_desc, assumed_align=16), - prob=prob_cute_fake, - bias=bias_cute_fake, - max_active_clusters=max_active_clusters, - stream=fake_stream, - epilogue_op=lambda x: x, - linear_offset=cutlass.Float32(0.0), - geglu_alpha=cutlass.Float32(1.702), - glu_clamp_max=cutlass.Float32(7.0), - glu_clamp_min=cutlass.Float32(-7.0), - options="--enable-tvm-ffi", - ) - - # Cache workspace pointer for the tensor_api closure - cached_workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator - - def tensor_api( - a_tensor: torch.Tensor, - b_tensor: torch.Tensor, - c_tensor: torch.Tensor, - d_tensor: torch.Tensor, - d_col_tensor: Optional[torch.Tensor], - sfa_tensor: torch.Tensor, - sfb_tensor: torch.Tensor, - sfd_row_tensor: Optional[torch.Tensor], - sfd_col_tensor: Optional[torch.Tensor], - amax_tensor: Optional[torch.Tensor], - norm_const_tensor: Optional[torch.Tensor], - padded_offsets: torch.Tensor, - alpha_tensor: torch.Tensor, - prob_tensor: Optional[torch.Tensor], - bias_tensor: Optional[torch.Tensor], - stream: cuda.CUstream, - linear_offset: float = 0.0, - geglu_alpha: float = 1.702, - glu_clamp_max: float = 7.0, - glu_clamp_min: float = -7.0, - ) -> None: - norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") - _compiled_kernel( - a_tensor, - b_tensor, - sfb_tensor, - cutlass.Int32(0), - cutlass.Int32(0), - cutlass.Int64(0), - cached_workspace_ptr, - c_tensor, - d_tensor, - d_col_tensor, - sfa_tensor, - sfd_row_tensor, - sfd_col_tensor, - amax_tensor, - norm_const_tensor, - padded_offsets, - alpha_tensor, - prob_tensor, - bias_tensor, - stream, - cutlass.Float32(linear_offset), - cutlass.Float32(geglu_alpha), - cutlass.Float32(glu_clamp_max), - cutlass.Float32(glu_clamp_min), - ) - - self._compiled_kernel = tensor_api - - # -- Discrete compile path ---------------------------------------------- # - - def _compile_discrete(self, gemm_glu, max_active_clusters, fake_stream) -> None: - """Compile for discrete (per-expert pointer) weight mode.""" - if len(self.b_shape) == 2: - n, k = self.b_shape - else: - n, k, _ = self.b_shape - - b_major_mode = OperandMajorMode.K if self.b_major == "k" else OperandMajorMode.MN - if self.b_major == "k": - b_stride_size = k - else: - b_stride_size = n - - ab_cutlass_dtype = _convert_to_cutlass_data_type(self.a_desc.dtype, interpret_uint8_as_fp4x2=self._interpret_uint8_as_fp4x2) - align = 32 if ab_cutlass_dtype.width == 4 else 16 - - valid_m = cute.sym_int(divisibility=256) - a_tensor = self._make_fake_cute_compact_tensor( - dtype=self.a_desc.dtype, - shape=(valid_m, *self.a_desc.shape[1:]), - stride_order=self.a_desc.stride_order, - assumed_align=align, - ) - c_tensor = self._make_fake_cute_compact_tensor( - dtype=self.c_desc.dtype, - shape=(valid_m, *self.c_desc.shape[1:]), - stride_order=self.c_desc.stride_order, - ) - d_tensor = self._make_fake_cute_compact_tensor( - dtype=self.d_desc.dtype, - shape=(valid_m, *self.d_desc.shape[1:]), - stride_order=self.d_desc.stride_order, - ) - d_col_tensor = self._make_fake_cute_compact_tensor( - dtype=self.d_col_desc.dtype, - shape=(valid_m, *self.d_col_desc.shape[1:]), - stride_order=self.d_col_desc.stride_order, - ) - - tensor_m_128 = cute.sym_int() - stride_tensor_m_128 = cute.sym_int(divisibility=32 * 4 * 4) - sfa_shape = list(self.sfa_desc.shape) - sfa_shape[2] = tensor_m_128 - sfa_stride = list(self.sfa_desc.stride) - sfa_stride[5] = stride_tensor_m_128 - sfa_tensor = self._make_fake_cute_tensor( - dtype=self.sfa_desc.dtype, - shape=tuple(sfa_shape), - stride=tuple(sfa_stride), - assumed_align=16, - ) - sfd_row_tensor = None - if self.sfd_row_desc is not None: - stride_sfd_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_row_tensor = self._make_fake_cute_tensor( - dtype=self.sfd_row_desc.dtype, - shape=(32, 4, tensor_m_128, 4, self.sfd_row_desc.shape[4], 1), - stride=(16, 4, self.sfd_row_desc.stride[2], 1, 512, stride_sfd_m), - assumed_align=16, - ) - sfd_col_tensor = None - if self.sfd_col_desc is not None: - rest_m = cute.sym_int(divisibility=1) - stride_sfd_n = cute.sym_int(divisibility=32 * 4 * 4) - stride_rest_m = cute.sym_int(divisibility=32 * 4 * 4) - sfd_col_tensor = self._make_fake_cute_tensor( - dtype=self.sfd_col_desc.dtype, - shape=(32, 4, self.sfd_col_desc.shape[2], 4, rest_m, 1), - stride=(16, 4, stride_rest_m, 1, 512, stride_sfd_n), - assumed_align=16, - ) - amax_tensor = self._make_fake_cute_tensor_from_desc(self.amax_desc, assumed_align=16) - norm_const_tensor_cute = self._make_fake_cute_tensor_from_desc(self.norm_const_desc, assumed_align=16) - padded_offsets_tensor = self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc, assumed_align=16) - alpha_tensor = self._make_fake_cute_tensor_from_desc(self.alpha_desc, assumed_align=16) - prob_tensor = None - if self.prob_desc is not None: - prob_tensor = self._make_fake_cute_tensor( - dtype=self.prob_desc.dtype, - shape=(valid_m, *self.prob_desc.shape[1:]), - stride=self.prob_desc.stride, - assumed_align=16, - ) - bias_tensor = self._make_fake_cute_tensor_from_desc(self.bias_desc, assumed_align=16) - - # Compile-time pointer placeholders - b_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - sfb_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device="cuda") - b_ptrs_cute = from_dlpack(b_ptrs_placeholder, assumed_align=8).iterator - sfb_ptrs_cute = from_dlpack(sfb_ptrs_placeholder, assumed_align=8).iterator - - workspace_ptr_cute = from_dlpack(self._workspace, assumed_align=128).iterator - - # linear_offset, geglu_alpha, glu_clamp_max, and glu_clamp_min are runtime - # cutlass.Float32 (not Constexpr), so the compile-time placeholders below are - # irrelevant -- the values passed through tensor_api() at execute() time are - # what the kernel actually uses. - self._logger.debug("Compiling discrete grouped GEMM GLU kernel") - _compiled_kernel = cute.compile( - gemm_glu, - a_tensor, - b_ptrs_cute, - sfb_ptrs_cute, - cutlass.Int32(n), - cutlass.Int32(k), - cutlass.Int64(b_stride_size), - b_major_mode, - workspace_ptr_cute, - c_tensor, - d_tensor, - d_col_tensor, - sfa_tensor, - sfd_row_tensor, - sfd_col_tensor, - amax_tensor, - norm_const_tensor_cute, - padded_offsets_tensor, - alpha_tensor, - prob_tensor, - bias_tensor, - max_active_clusters, - fake_stream, - lambda x: x, # epilogue_op (Constexpr, baked in) - cutlass.Float32(0.0), - cutlass.Float32(1.702), - cutlass.Float32(7.0), - cutlass.Float32(-7.0), - options="--enable-tvm-ffi", - ) - - self._n = n - self._k = k - self._b_stride_size = b_stride_size - - # Cache constant values for execute() closure - cached_workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator - cached_n = cutlass.Int32(self._n) - cached_k = cutlass.Int32(self._k) - cached_b_stride = cutlass.Int64(self._b_stride_size) - - def tensor_api( - a_tensor: torch.Tensor, - b_ptrs_device: torch.Tensor, - sfb_ptrs_device: torch.Tensor, - c_tensor: torch.Tensor, - d_tensor: torch.Tensor, - d_col_tensor: Optional[torch.Tensor], - sfa_tensor: torch.Tensor, - sfd_row_tensor: Optional[torch.Tensor], - sfd_col_tensor: Optional[torch.Tensor], - amax_tensor: Optional[torch.Tensor], - norm_const_tensor: Optional[torch.Tensor], - padded_offsets: torch.Tensor, - alpha_tensor: torch.Tensor, - prob_tensor: Optional[torch.Tensor], - bias_tensor: Optional[torch.Tensor], - stream: cuda.CUstream, - linear_offset: float = 0.0, - geglu_alpha: float = 1.702, - glu_clamp_max: float = 7.0, - glu_clamp_min: float = -7.0, - ) -> None: - norm_const_tensor = self._unpad_tensor_to_ndim(norm_const_tensor, 1, "norm_const") - b_ptrs_addr = int(b_ptrs_device.data_ptr()) - sfb_ptrs_addr = int(sfb_ptrs_device.data_ptr()) - - _compiled_kernel( - a_tensor, - b_ptrs_addr, - sfb_ptrs_addr, - cached_n, - cached_k, - cached_b_stride, - cached_workspace_ptr, - c_tensor, - d_tensor, - d_col_tensor, - sfa_tensor, - sfd_row_tensor, - sfd_col_tensor, - amax_tensor, - norm_const_tensor, - padded_offsets, - alpha_tensor, - prob_tensor, - bias_tensor, - stream, - cutlass.Float32(linear_offset), - cutlass.Float32(geglu_alpha), - cutlass.Float32(glu_clamp_max), - cutlass.Float32(glu_clamp_min), - ) - - self._compiled_kernel = tensor_api - - # --------------------------------------------------------------------- # - # execute - # --------------------------------------------------------------------- # + else: + block_kwargs = dict(kwargs) + block_kwargs.pop("generate_c", None) + self._implementation = GroupedGemmGluBlockScaledAPI(**block_kwargs) + self._kernel = self._implementation._kernel + self.weight_mode = self._implementation.weight_mode + supported = self._implementation.check_support() + self._is_supported = self._implementation._is_supported + if supported: + self._pending_init_kwargs = None + return supported + def compile(self) -> None: + if self._implementation is None: + self.check_support() + if self._is_supported: + self._implementation._is_supported = True + self._implementation.compile() + self._is_supported = self._implementation._is_supported + self._compiled_kernel = self._implementation._compiled_kernel + + # BF16 implementation + @overload + def execute( + self, + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + sfa_tensor: None, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor] = None, + *, + sfb_tensor: None = None, + sfb_ptrs: None = None, + d_col_tensor: None = None, + sfd_row_tensor: None = None, + sfd_col_tensor: None = None, + amax_tensor: None = None, + norm_const_tensor: None = None, + ) -> None: ... + + # Block-scaled implementation + @overload def execute( self, a_tensor: torch.Tensor, @@ -1123,14 +283,30 @@ def execute( sfa_tensor: torch.Tensor, padded_offsets: torch.Tensor, alpha_tensor: torch.Tensor, - # Dense mode: + b_tensor: Optional[torch.Tensor] = None, + *, + sfb_tensor: Optional[torch.Tensor] = None, + sfb_ptrs: Optional[torch.Tensor] = None, + d_col_tensor: Optional[torch.Tensor] = None, + sfd_row_tensor: Optional[torch.Tensor] = None, + sfd_col_tensor: Optional[torch.Tensor] = None, + amax_tensor: Optional[torch.Tensor] = None, + norm_const_tensor: Optional[torch.Tensor] = None, + ) -> None: ... + + def execute( + self, + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + sfa_tensor: Optional[torch.Tensor], + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, b_tensor: Optional[torch.Tensor] = None, sfb_tensor: Optional[torch.Tensor] = None, bias_tensor: Optional[torch.Tensor] = None, - # Discrete mode: b_ptrs: Optional[torch.Tensor] = None, sfb_ptrs: Optional[torch.Tensor] = None, - # Optional: d_col_tensor: Optional[torch.Tensor] = None, sfd_row_tensor: Optional[torch.Tensor] = None, sfd_col_tensor: Optional[torch.Tensor] = None, @@ -1143,121 +319,73 @@ def execute( glu_clamp_min: float = -7.0, current_stream: Optional[cuda.CUstream] = None, ) -> None: - """Execute the compiled kernel. - - For dense mode, supply ``b_tensor`` and ``sfb_tensor``. - For discrete mode, supply ``b_ptrs`` and ``sfb_ptrs``. - - :param a_tensor: Input A tensor - :param c_tensor: Intermediate C tensor - :param d_tensor: Output D tensor - :param sfa_tensor: Scale factor A - :param padded_offsets: End offset per expert after padding - :param alpha_tensor: Per-group scaling factors - :param b_tensor: (Dense) Input B tensor (weights) - :param sfb_tensor: (Dense) Scale factor B - :param bias_tensor: Optional bias tensor with shape (n, l) and stride (1, n). - Bias fusion is specialized at compile time: if ``sample_bias`` was omitted - at construction, ``bias_tensor`` must also be omitted at execute time. - :param b_ptrs: (Discrete) 1-D int64 device tensor of per-expert B data pointers - :param sfb_ptrs: (Discrete) 1-D int64 device tensor of per-expert SFB data pointers - :param d_col_tensor: Optional column-quantized output - :param sfd_row_tensor: Optional row scale factor D - :param sfd_col_tensor: Optional column scale factor D - :param amax_tensor: Optional amax tensor - :param norm_const_tensor: Optional normalization constant - :param prob_tensor: Optional probability tensor - :param linear_offset: Linear offset applied to the up branch in the - ``act_func == "geglu"`` activation. Ignored when - ``act_func == "swiglu"``. When ``None`` (default), the offset is - chosen based on ``act_func`` for backwards compatibility: - ``1.0`` for ``"geglu"`` and ``0.0`` for ``"swiglu"``. - :param geglu_alpha: Pre-sigmoid scaling factor for the GeGLU activation. - The fused activation is - ``out = (clamp(up, glu_clamp_min, glu_clamp_max) + linear_offset) - * silu(geglu_alpha * clamp(gate, max=glu_clamp_max))``. - Defaults to ``1.702`` (GPT-OSS / scaled-GeGLU). Ignored when - ``act_func == "swiglu"``. - :param glu_clamp_max: Upper clamp limit applied to both ``gate`` and - ``up`` before the activation. Default ``7.0``. Ignored when - ``act_func == "swiglu"``. - :param glu_clamp_min: Lower clamp limit applied only to ``up`` (the - kernel never lower-clamps ``gate``). Default ``-7.0``. Ignored - when ``act_func == "swiglu"``. - :param current_stream: CUDA stream - """ - self._logger.debug("Entering execute") - current_stream = self._get_default_stream(current_stream) - - if a_tensor.shape[0] == 0: - self._logger.debug("execute: valid_m is zero, skipping kernel execution") - return - self._runtime_error_if( - self._compiled_kernel is None, - "Kernel not compiled; call compile() first", - ) - - # Resolve linear_offset default: None -> activation-derived legacy value - # (1.0 for geglu, 0.0 for swiglu) for backwards compatibility with callers - # that pre-date the explicit linear_offset kwarg. - if linear_offset is None: - linear_offset = 1.0 if self.act_func == "geglu" else 0.0 - - self._logger.debug("Executing grouped GEMM GLU kernel") - if self._has_bias: - self._value_error_if( - bias_tensor is None, - "bias_tensor must be provided at execute() when the API was compiled with sample_bias", + if self._implementation is None: + raise RuntimeError("Kernel not compiled; call compile() first") + if self.backend is GroupedGemmBackend.BF16: + scale_controls = ( + ("sfa_tensor", sfa_tensor), + ("sfb_tensor", sfb_tensor), + ("sfb_ptrs", sfb_ptrs), + ("d_col_tensor", d_col_tensor), + ("sfd_row_tensor", sfd_row_tensor), + ("sfd_col_tensor", sfd_col_tensor), + ("amax_tensor", amax_tensor), + ("norm_const_tensor", norm_const_tensor), + ("geglu_alpha", geglu_alpha if geglu_alpha != 1.702 else None), + ( + "glu_clamp_max", + glu_clamp_max if glu_clamp_max != 7.0 else None, + ), + ( + "glu_clamp_min", + glu_clamp_min if glu_clamp_min != -7.0 else None, + ), ) - - if self.weight_mode == MoEWeightMode.DENSE: - self._compiled_kernel( + forbidden = [name for name, value in scale_controls if value is not None] + if forbidden: + raise ValueError(f"grouped_gemm_glu_sm100: BF16 forbids scale control " f"{forbidden[0]}") + if linear_offset is None: + linear_offset = 1.0 if self._implementation.act_func == "geglu" else 0.0 + self._implementation.execute( a_tensor=a_tensor, - b_tensor=b_tensor, c_tensor=c_tensor, d_tensor=d_tensor, - d_col_tensor=d_col_tensor, - sfa_tensor=sfa_tensor, - sfb_tensor=sfb_tensor, - sfd_row_tensor=sfd_row_tensor, - sfd_col_tensor=sfd_col_tensor, - amax_tensor=amax_tensor, - norm_const_tensor=norm_const_tensor, padded_offsets=padded_offsets, alpha_tensor=alpha_tensor, + b_tensor=b_tensor, + b_ptrs=b_ptrs, bias_tensor=bias_tensor, prob_tensor=prob_tensor, - stream=current_stream, linear_offset=linear_offset, - geglu_alpha=geglu_alpha, - glu_clamp_max=glu_clamp_max, - glu_clamp_min=glu_clamp_min, + current_stream=current_stream, ) else: - self._compiled_kernel( + self._implementation.execute( a_tensor=a_tensor, - b_ptrs_device=b_ptrs, - sfb_ptrs_device=sfb_ptrs, c_tensor=c_tensor, d_tensor=d_tensor, - d_col_tensor=d_col_tensor, sfa_tensor=sfa_tensor, + padded_offsets=padded_offsets, + alpha_tensor=alpha_tensor, + b_tensor=b_tensor, + sfb_tensor=sfb_tensor, + bias_tensor=bias_tensor, + b_ptrs=b_ptrs, + sfb_ptrs=sfb_ptrs, + d_col_tensor=d_col_tensor, sfd_row_tensor=sfd_row_tensor, sfd_col_tensor=sfd_col_tensor, amax_tensor=amax_tensor, norm_const_tensor=norm_const_tensor, - padded_offsets=padded_offsets, - alpha_tensor=alpha_tensor, prob_tensor=prob_tensor, - bias_tensor=bias_tensor, - stream=current_stream, linear_offset=linear_offset, geglu_alpha=geglu_alpha, glu_clamp_max=glu_clamp_max, glu_clamp_min=glu_clamp_min, + current_stream=current_stream, ) - - self._logger.debug("Execute completed") + self._is_supported = self._implementation._is_supported + self._compiled_kernel = self._implementation._compiled_kernel # --------------------------------------------------------------------------- # @@ -1268,42 +396,7 @@ def execute( _cache_of_GroupedGemmGluSm100Objects = {} -def grouped_gemm_glu_wrapper_sm100( - a_tensor: torch.Tensor, - sfa_tensor: torch.Tensor, - padded_offsets: torch.Tensor, - alpha_tensor: torch.Tensor, - # bias_tensor is optional in both modes: - b_tensor: Optional[torch.Tensor] = None, - sfb_tensor: Optional[torch.Tensor] = None, - bias_tensor: Optional[torch.Tensor] = None, - # Discrete mode: - b_ptrs: Optional[torch.Tensor] = None, - sfb_ptrs: Optional[torch.Tensor] = None, - n: Optional[int] = None, - b_dtype: Optional[torch.dtype] = None, - b_major: str = "k", - # Common: - norm_const_tensor: Optional[torch.Tensor] = None, - prob_tensor: Optional[torch.Tensor] = None, - acc_dtype: torch.dtype = torch.float32, - c_dtype: torch.dtype = torch.bfloat16, - d_dtype: torch.dtype = torch.bfloat16, - cd_major: str = "n", - mma_tiler_mn: Tuple[int, int] = (256, 256), - cluster_shape_mn: Optional[Tuple[int, int]] = None, - sf_vec_size: int = 16, - vector_f32: bool = False, - m_aligned: int = 256, - discrete_col_sfd: bool = False, - act_func: str = "swiglu", - linear_offset: Optional[float] = None, - geglu_alpha: float = 1.702, - glu_clamp_max: float = 7.0, - glu_clamp_min: float = -7.0, - use_dynamic_sched: bool = False, - current_stream: Optional[cuda.CUstream] = None, -) -> TupleDict: +def _grouped_gemm_glu_block_scaled_call(call: GluCall) -> TupleDict: """Convenience wrapper for grouped GEMM GLU forward operation. Auto-detects dense vs. discrete mode based on which weight arguments @@ -1372,6 +465,38 @@ def grouped_gemm_glu_wrapper_sm100( """ from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor + a_tensor = call.a_tensor + sfa_tensor = call.sfa_tensor + padded_offsets = call.padded_offsets + alpha_tensor = call.alpha_tensor + b_tensor = call.b_tensor + sfb_tensor = call.sfb_tensor + bias_tensor = call.bias_tensor + b_ptrs = call.b_ptrs + sfb_ptrs = call.sfb_ptrs + n = call.n + b_dtype = call.b_dtype + b_major = call.b_major + norm_const_tensor = call.norm_const_tensor + prob_tensor = call.prob_tensor + acc_dtype = call.acc_dtype + c_dtype = call.c_dtype + d_dtype = call.d_dtype + cd_major = call.cd_major + mma_tiler_mn = call.mma_tiler_mn + cluster_shape_mn = call.cluster_shape_mn + sf_vec_size = call.sf_vec_size + vector_f32 = call.vector_f32 + m_aligned = call.m_aligned + discrete_col_sfd = call.discrete_col_sfd + act_func = call.act_func + linear_offset = call.linear_offset + geglu_alpha = call.geglu_alpha + glu_clamp_max = call.glu_clamp_max + glu_clamp_min = call.glu_clamp_min + use_dynamic_sched = call.use_dynamic_sched + current_stream = call.current_stream + # Resolve linear_offset default: None means "use the activation-derived legacy # default" (1.0 for geglu, 0.0 for swiglu) for backwards compatibility with # callers that have not been updated to pass linear_offset explicitly. @@ -1561,6 +686,8 @@ def dynamic_m_tensor_signature( num_experts, ) + cache_key = backend_cache_key(GroupedGemmBackend.BLOCK_SCALED, *cache_key) + # ---- Cache lookup or create + compile ---- if cache_key in _cache_of_GroupedGemmGluSm100Objects: _logger.debug("grouped_gemm_glu_wrapper_sm100: Using cached object") @@ -1685,3 +812,353 @@ def dynamic_m_tensor_signature( sfd_row_tensor=sfd_row_tensor, sfd_col_tensor=sfd_col_tensor, ) + + +def _normalize_glu_call(call: GluCall) -> tuple[GluCall, GroupedGemmBackend]: + from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor + + is_dense = call.b_tensor is not None + is_discrete = call.b_ptrs is not None + if is_dense and is_discrete: + raise ValueError("Provide either (b_tensor, sfb_tensor) or (b_ptrs, sfb_ptrs), not both") + if not is_dense and not is_discrete: + raise ValueError("Must provide either (b_tensor, sfb_tensor) or (b_ptrs, sfb_ptrs)") + if call.a_tensor.ndim != 3 or call.a_tensor.shape[2] != 1: + raise ValueError(f"a_tensor must have shape (m, k, 1), got {tuple(call.a_tensor.shape)}") + + valid_m, k, _ = call.a_tensor.shape + if is_dense: + if call.b_tensor.ndim != 3: + raise ValueError(f"b_tensor must have shape (n, k, experts), got " f"{tuple(call.b_tensor.shape)}") + n_full, b_k, num_experts = call.b_tensor.shape + if b_k != k: + raise ValueError(f"b_tensor K dimension ({b_k}) must match a_tensor ({k})") + defining_b_dtype = call.b_tensor.dtype + b_shape = None + weight_mode = MoEWeightMode.DENSE + if call.n is not None or call.b_dtype is not None: + raise ValueError("Dense mode forbids n and b_dtype") + else: + _require_pointer_tensor(call.b_ptrs, "b_ptrs") + num_experts = call.b_ptrs.numel() + if call.n is None or call.b_dtype is None: + raise ValueError("n and b_dtype are required for discrete mode") + n_full = call.n + defining_b_dtype = call.b_dtype + b_shape = (n_full, k) + weight_mode = MoEWeightMode.DISCRETE + + backend = select_grouped_gemm_backend( + operation="grouped_gemm_glu_sm100", + a_dtype=call.a_tensor.dtype, + b_dtype=defining_b_dtype, + scale_controls=( + ("sfa_tensor", call.sfa_tensor), + ("sfb_tensor", call.sfb_tensor), + ("sfb_ptrs", call.sfb_ptrs), + ("norm_const_tensor", call.norm_const_tensor), + ("sf_vec_size", call.sf_vec_size if call.sf_vec_size != 16 else None), + ( + "discrete_col_sfd", + call.discrete_col_sfd if call.discrete_col_sfd else None, + ), + ( + "geglu_alpha", + call.geglu_alpha if call.geglu_alpha != 1.702 else None, + ), + ( + "glu_clamp_max", + call.glu_clamp_max if call.glu_clamp_max != 7.0 else None, + ), + ( + "glu_clamp_min", + call.glu_clamp_min if call.glu_clamp_min != -7.0 else None, + ), + ), + block_scaled_dtype_pairs=_BLOCK_SCALED_DTYPE_PAIRS, + ) + + linear_offset = call.linear_offset + if linear_offset is None: + linear_offset = 1.0 if call.act_func == "geglu" else 0.0 + + normalized = replace( + call, + linear_offset=linear_offset, + weight_mode=weight_mode, + b_shape=b_shape, + num_experts=num_experts, + ) + if backend is GroupedGemmBackend.BLOCK_SCALED: + return normalized, backend + + if call.prob_tensor is None: + raise ValueError("prob_tensor is required for BF16") + if call.cd_major != "n": + raise ValueError(f"cd_major must be 'n', got {call.cd_major}") + if call.act_func not in ("swiglu", "geglu"): + raise ValueError(f"act_func must be 'swiglu' or 'geglu', got {call.act_func}") + if call.c_dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise ValueError(f"c_dtype must be BF16, FP16, or FP32, got {call.c_dtype}") + if call.d_dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise ValueError(f"d_dtype must be BF16, FP16, or FP32, got {call.d_dtype}") + if call.m_aligned != 256: + raise ValueError(f"m_aligned must be 256, got {call.m_aligned}") + if valid_m % 256 != 0: + raise ValueError(f"a_tensor M dimension must be 256-aligned, got {valid_m}") + if n_full <= 0 or n_full % 64 != 0: + raise ValueError(f"N must be positive and divisible by 64, got {n_full}") + if tuple(call.prob_tensor.shape) != (valid_m, 1, 1): + raise ValueError(f"prob_tensor must have shape {(valid_m, 1, 1)}, got " f"{tuple(call.prob_tensor.shape)}") + if call.bias_tensor is not None and tuple(call.bias_tensor.shape) != ( + n_full, + num_experts, + ): + raise ValueError(f"bias_tensor must have shape {(n_full, num_experts)}, got " f"{tuple(call.bias_tensor.shape)}") + if is_discrete: + if call.b_ptrs.device != call.a_tensor.device: + raise ValueError(f"b_ptrs must be on the same device as a_tensor " f"({call.a_tensor.device}), got {call.b_ptrs.device}") + if call.b_ptrs.numel() != call.padded_offsets.numel(): + raise ValueError(f"b_ptrs length mismatch: expected {call.padded_offsets.numel()}, " f"got {call.b_ptrs.numel()}") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + major, minor = torch.cuda.get_device_capability(call.a_tensor.device) + compute_capability = major * 10 + minor + if compute_capability < 100: + raise RuntimeError(f"GroupedGemmGluSm100 requires SM100+, found SM{compute_capability}") + return normalized, backend + + +def _glu_stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: + return tuple( + index + for index, _ in sorted( + enumerate(tensor.stride()), + key=lambda item: (item[1], tensor.shape[item[0]]), + ) + ) + + +def _glu_tensor_signature(tensor: Optional[torch.Tensor], *, dynamic_m: bool = False) -> tuple: + if tensor is None: + return (None, None, None, None) + shape = (None, *tuple(tensor.shape[1:])) if dynamic_m else tuple(tensor.shape) + return ( + shape, + _glu_stride_order(tensor), + tensor.dtype, + (tensor.device.type, tensor.device.index), + ) + + +def _grouped_gemm_glu_bf16_call(call: GluCall) -> TupleDict: + valid_m, k, _ = call.a_tensor.shape + if call.weight_mode == MoEWeightMode.DENSE: + n_full = call.b_tensor.shape[0] + else: + n_full = call.n + n_out = n_full // 2 + + c_tensor = torch.empty_strided( + (valid_m, n_full, 1), + (n_full, 1, valid_m * n_full), + dtype=call.c_dtype, + device=call.a_tensor.device, + ) + d_tensor = torch.empty_strided( + (valid_m, n_out, 1), + (n_out, 1, valid_m * n_out), + dtype=call.d_dtype, + device=call.a_tensor.device, + ) + + overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + workspace_bytes = (128 * call.num_experts if call.weight_mode == MoEWeightMode.DISCRETE else 0) + (4 if call.use_dynamic_sched else 0) + cache_key = backend_cache_key( + GroupedGemmBackend.BF16, + call.weight_mode, + call.act_func, + _glu_tensor_signature(call.a_tensor, dynamic_m=True), + _glu_tensor_signature(call.b_tensor), + call.b_shape, + call.b_dtype, + _glu_tensor_signature(c_tensor, dynamic_m=True), + _glu_tensor_signature(d_tensor, dynamic_m=True), + _glu_tensor_signature(call.alpha_tensor), + _glu_tensor_signature(call.bias_tensor), + _glu_tensor_signature(call.padded_offsets), + _glu_tensor_signature(call.prob_tensor, dynamic_m=True), + ( + ( + tuple(call.b_ptrs.shape), + tuple(call.b_ptrs.stride()), + call.b_ptrs.dtype, + (call.b_ptrs.device.type, call.b_ptrs.device.index), + ) + if call.b_ptrs is not None + else None + ), + call.acc_dtype, + call.c_dtype, + call.d_dtype, + call.mma_tiler_mn, + call.cluster_shape_mn, + call.vector_f32, + call.m_aligned, + call.generate_c, + call.b_major, + call.use_dynamic_sched, + workspace_bytes, + (call.a_tensor.device.type, call.a_tensor.device.index), + overlap_margin, + ) + + if cache_key in _cache_of_GroupedGemmGluSm100Objects: + api = _cache_of_GroupedGemmGluSm100Objects[cache_key] + else: + api = GroupedGemmGluSm100( + sample_a=call.a_tensor, + sample_c=c_tensor, + sample_d=d_tensor, + sample_sfa=None, + sample_padded_offsets=call.padded_offsets, + sample_alpha=call.alpha_tensor, + sample_d_col=None, + sample_b=call.b_tensor, + sample_sfb=None, + sample_bias=call.bias_tensor, + num_experts=(call.num_experts if call.weight_mode == MoEWeightMode.DISCRETE else None), + b_shape=call.b_shape, + b_dtype=call.b_dtype, + sample_sfd_row=None, + sample_sfd_col=None, + sample_amax=None, + sample_norm_const=None, + sample_prob=call.prob_tensor, + acc_dtype=call.acc_dtype, + mma_tiler_mn=call.mma_tiler_mn, + cluster_shape_mn=call.cluster_shape_mn, + sf_vec_size=16, + vector_f32=call.vector_f32, + m_aligned=call.m_aligned, + discrete_col_sfd=False, + act_func=call.act_func, + b_major=call.b_major, + use_dynamic_sched=call.use_dynamic_sched, + generate_c=call.generate_c, + ) + if not api.check_support(): + raise RuntimeError("Unsupported BF16 configuration") + api.compile() + _cache_of_GroupedGemmGluSm100Objects[cache_key] = api + + api.execute( + a_tensor=call.a_tensor, + c_tensor=c_tensor, + d_tensor=d_tensor, + sfa_tensor=None, + padded_offsets=call.padded_offsets, + alpha_tensor=call.alpha_tensor, + b_tensor=call.b_tensor, + sfb_tensor=None, + bias_tensor=call.bias_tensor, + b_ptrs=call.b_ptrs, + sfb_ptrs=None, + d_col_tensor=None, + sfd_row_tensor=None, + sfd_col_tensor=None, + amax_tensor=None, + norm_const_tensor=None, + prob_tensor=call.prob_tensor, + linear_offset=call.linear_offset, + geglu_alpha=call.geglu_alpha, + glu_clamp_max=call.glu_clamp_max, + glu_clamp_min=call.glu_clamp_min, + current_stream=call.current_stream, + ) + return TupleDict( + c_tensor=c_tensor if call.generate_c else None, + d_tensor=d_tensor, + d_col_tensor=None, + amax_tensor=None, + sfd_row_tensor=None, + sfd_col_tensor=None, + ) + + +def grouped_gemm_glu_wrapper_sm100( + a_tensor: torch.Tensor, + sfa_tensor: Optional[torch.Tensor], + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor] = None, + sfb_tensor: Optional[torch.Tensor] = None, + bias_tensor: Optional[torch.Tensor] = None, + b_ptrs: Optional[torch.Tensor] = None, + sfb_ptrs: Optional[torch.Tensor] = None, + n: Optional[int] = None, + b_dtype: Optional[torch.dtype] = None, + b_major: str = "k", + norm_const_tensor: Optional[torch.Tensor] = None, + prob_tensor: Optional[torch.Tensor] = None, + acc_dtype: torch.dtype = torch.float32, + c_dtype: torch.dtype = torch.bfloat16, + d_dtype: torch.dtype = torch.bfloat16, + cd_major: str = "n", + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + sf_vec_size: int = 16, + vector_f32: bool = False, + m_aligned: int = 256, + discrete_col_sfd: bool = False, + act_func: str = "swiglu", + linear_offset: Optional[float] = None, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, + use_dynamic_sched: bool = False, + current_stream: Optional[cuda.CUstream] = None, + generate_c: bool = False, +) -> TupleDict: + """Dispatch grouped GEMM GLU once from an immutable normalized call.""" + call = GluCall( + a_tensor=a_tensor, + sfa_tensor=sfa_tensor, + padded_offsets=padded_offsets, + alpha_tensor=alpha_tensor, + b_tensor=b_tensor, + sfb_tensor=sfb_tensor, + bias_tensor=bias_tensor, + b_ptrs=b_ptrs, + sfb_ptrs=sfb_ptrs, + n=n, + b_dtype=b_dtype, + b_major=b_major, + norm_const_tensor=norm_const_tensor, + prob_tensor=prob_tensor, + acc_dtype=acc_dtype, + c_dtype=c_dtype, + d_dtype=d_dtype, + cd_major=cd_major, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + vector_f32=vector_f32, + m_aligned=m_aligned, + discrete_col_sfd=discrete_col_sfd, + act_func=act_func, + linear_offset=linear_offset, + geglu_alpha=geglu_alpha, + glu_clamp_max=glu_clamp_max, + glu_clamp_min=glu_clamp_min, + use_dynamic_sched=use_dynamic_sched, + current_stream=current_stream, + generate_c=generate_c, + ) + call, backend = _normalize_glu_call(call) + if backend is GroupedGemmBackend.BF16: + return _grouped_gemm_glu_bf16_call(call) + return _grouped_gemm_glu_block_scaled_call(call) + + +__all__ = ["GluCall", "GroupedGemmGluSm100", "grouped_gemm_glu_wrapper_sm100"] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_glu/moe_grouped_gemm_glu_bias.py b/python/cudnn/grouped_gemm/grouped_gemm_glu/moe_grouped_gemm_glu_bias.py new file mode 100644 index 000000000..2e5d34da7 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_glu/moe_grouped_gemm_glu_bias.py @@ -0,0 +1,1934 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +MoE BF16 Grouped GEMM Kernel with GLU (SwiGLU/GeGLU) Fusion. + +Supports: + - Static / Dynamic persistent tile scheduling (MoEPersistentTileScheduler) + - Dense (contiguous 3-D B) / Discrete (per-expert pointer array B) weight layout + - Optional bias + - GLU activation fusion (SwiGLU / GeGLU) + +This module contains only the kernel class. +MoE scheduler components live in moe_persistent_scheduler.py / moe_sched_extension.py / moe_utils.py. +""" + +from typing import Type, Tuple, Union, Optional + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.typing import Int32, AddressSpace +from ..moe_persistent_scheduler import ( + MoEPersistentTileScheduler, + MoESchedulerParams, + MoEWorkTileInfo, +) +from ..moe_utils import ( + MoEWeightMode, + TensormapWorkspace, + store_tma_desc, +) +from ..moe_sched_extension import ( + DiscreteWeightGroupedGemmSchedExtension, + ContiguousGroupedGemmSchedExtension, +) +from ..moe_kernel_helpers import ( + fmin, + fmax, + silu_f32, + silu_f32_geglu_scaled, + can_implement_bf16_grouped_gemm, + compute_grid, + epilog_gmem_copy_and_partition, +) + + +class MoEGroupedGemmGluBiasBf16Kernel: + """Plain BF16 grouped GEMM kernel with MoE scheduling and GLU fusion. + + The kernel is organized as persistent scheduler, A/B TMA load, MMA, and + epilogue warps. The epilogue applies SwiGLU or GeGLU and optionally adds + bias before writing D. + + :note: Constraints: + - MMA tiler M must be 128 or 256 (use_2cta_instrs) + - MMA tiler N must be 64/128/192/256 + - Cluster shape M must be multiple of 2 if Mma tiler M is 256 + - Cluster shape M/N must be positive and power of 2, total cluster size <= 16 + - FIX_PAD_SIZE (256) must be divisible by mma_tiler_mn[0] + - m_aligned parameter in create_mask() MUST equal FIX_PAD_SIZE (256) + - Each padded_offsets[i] will be a multiple of FIX_PAD_SIZE (guaranteed by m_aligned == FIX_PAD_SIZE) + """ + + # Fixed pad size for user-side padding (decoupled from kernel tile size) + FIX_PAD_SIZE = 256 + + @staticmethod + def can_implement( + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m: int, + n: int, + k: int, + l: int, + a_major: str, + b_major: str, + cd_major: str, + m_aligned: int, + ) -> bool: + return can_implement_bf16_grouped_gemm( + ab_dtype, + c_dtype, + d_dtype, + acc_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + m, + n, + k, + l, + a_major, + b_major, + cd_major, + m_aligned, + fix_pad_size=MoEGroupedGemmGluBiasBf16Kernel.FIX_PAD_SIZE, + ) + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + vectorized_f32: bool, + expert_cnt: int, + weight_mode: MoEWeightMode = MoEWeightMode.DISCRETE, + use_dynamic_sched: bool = False, + act_func: str = "swiglu", + enable_bias: bool = False, + generate_c: bool = False, + ): + mma_tile_m = mma_tiler_mn[0] + if self.FIX_PAD_SIZE % mma_tile_m != 0: + raise ValueError(f"FIX_PAD_SIZE ({self.FIX_PAD_SIZE}) must be divisible by " f"mma_tiler_mn[0] ({mma_tile_m}).") + if expert_cnt > 1024: + raise ValueError("Expert count > 1024 is not supported.") + if not isinstance(weight_mode, MoEWeightMode): + raise TypeError(f"weight_mode must be a MoEWeightMode, got {type(weight_mode)}") + if act_func not in ["swiglu", "geglu"]: + raise ValueError(f"Invalid activation function: {act_func}") + + self.expert_cnt = expert_cnt + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + # K dimension is deferred in _setup_attributes + self.mma_tiler = (*mma_tiler_mn, 1) + + self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + + self.enable_bias = enable_bias + self.occupancy = 1 + self.epilog_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.sched_warp_id = 6 + self.bias_load_warp_id = 7 if enable_bias else None + self.threads_per_warp = 32 + + all_warps = [ + *self.epilog_warp_id, + self.mma_warp_id, + self.tma_warp_id, + self.sched_warp_id, + ] + warps_wo_sched = [ + *self.epilog_warp_id, + self.mma_warp_id, + self.tma_warp_id, + ] + if enable_bias: + all_warps.append(self.bias_load_warp_id) + warps_wo_sched.append(self.bias_load_warp_id) + self.threads_per_cta = self.threads_per_warp * len(all_warps) + self.threads_wo_sched = self.threads_per_warp * len(warps_wo_sched) + + # Set barrier for cta sync, epilogue sync and tmem ptr sync + self.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=32 * len(self.epilog_warp_id), + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), + ) + self.sched_sync_barrier = pipeline.NamedBarrier( + barrier_id=4, + num_threads=self.threads_per_warp, + ) + self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + + self.vectorized_f32 = vectorized_f32 + self.generate_c = generate_c + + self.weight_mode = weight_mode + self.use_dynamic_sched = use_dynamic_sched + + self.act_func = act_func + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue subtile + - Setting up A/B/D stage counts in shared memory + - Computing A/B/D shared memory layout + """ + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + self.mma_tiler_d = ( + self.mma_tiler[0], + self.mma_tiler[1] // 2, + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_d = ( + self.mma_tiler_d[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_d[1], + self.mma_tiler_d[2], + ) + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + # Set epilogue subtile + self.epi_tile = (128, 32) + self.epi_tile_cnt = ( + self.cta_tile_shape_mnk_d[0] // self.epi_tile[0], + self.cta_tile_shape_mnk_d[1] // self.epi_tile[1], + ) + self.epi_tile_c = (128, 64) + + # Setup A/B/C/D stage count in shared memory and ACC stage count in tensor memory + ( + self.num_acc_stage, + self.num_ab_stage, + self.num_c_stage, + self.num_d_stage, + self.num_tile_stage, + self.num_bias_stage, + ) = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.epi_tile, + self.epi_tile_c, + self.c_dtype, + self.c_layout, + self.d_dtype, + self.d_layout, + self.num_smem_capacity, + self.occupancy, + bias_dtype=self.bias_dtype if self.enable_bias else None, + ) + + # TMEM accumulator columns: derive from the actual accumulator fragment and + # round up to a valid power-of-two allocation (handles non-power-of-two tile N), + # matching the plain BF16 grouped GEMM kernel instead of always reserving 512. + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + + # Compute A/B/C/D shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile_c, + self.num_c_stage, + ) + + self.d_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.d_dtype, + self.d_layout, + self.epi_tile, + self.num_d_stage, + ) + + # Bias SMEM layout: (tile_N, num_stages) double-buffered + if self.enable_bias: + self.bias_smem_layout_staged = cute.make_layout( + (self.mma_tiler[1], self.num_bias_stage), + stride=(1, self.mma_tiler[1]), + ) + else: + self.bias_smem_layout_staged = cute.make_layout((1, 1)) + + @staticmethod + def _compute_stages( + tiled_mma, + mma_tiler_mnk, + a_dtype, + b_dtype, + epi_tile, + epi_tile_c, + c_dtype, + c_layout, + d_dtype, + d_layout, + num_smem_capacity, + occupancy, + bias_dtype, + ): + """Compute BF16-only pipeline stages. + + Stage counts are chosen for the BF16 A/B mainloop and the GLU epilogue + shared-memory footprint. + """ + num_acc_stage = 2 + num_c_stage = 1 + num_d_stage = 1 + num_tile_stage = 2 + + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler_mnk, a_dtype, 1) + b_smem_layout_stage_one = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler_mnk, b_dtype, 1) + c_smem_layout_stage_one = sm100_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile_c, 1) + d_smem_layout_stage_one = sm100_utils.make_smem_layout_epi(d_dtype, d_layout, epi_tile, 1) + + ab_bytes_per_stage = cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + cute.size_in_bytes(b_dtype, b_smem_layout_stage_one) + mbar_helpers_bytes = 1024 + sinfo_bytes = 4 * 4 * num_tile_stage + c_bytes = cute.size_in_bytes(c_dtype, c_smem_layout_stage_one) * num_c_stage + d_bytes = cute.size_in_bytes(d_dtype, d_smem_layout_stage_one) * num_d_stage + + if bias_dtype is not None: + num_bias_stage = 2 + bias_bytes = mma_tiler_mnk[1] * num_bias_stage * (bias_dtype.width // 8) + else: + num_bias_stage = 0 + bias_bytes = 0 + + epi_bytes = c_bytes + d_bytes + bias_bytes + num_ab_stage = (num_smem_capacity // occupancy - (mbar_helpers_bytes + epi_bytes + sinfo_bytes)) // ab_bytes_per_stage + return num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, num_tile_stage, num_bias_stage + + def get_desc_workspace_bytes(self) -> int: + """Return descriptor workspace size in bytes.""" + if self.weight_mode == MoEWeightMode.DISCRETE: + return TensormapWorkspace.size_bytes(1, self.expert_cnt) + return 0 + + def get_workspace_bytes(self) -> int: + """Return descriptor workspace plus optional dynamic scheduler state.""" + desc_workspace_bytes = self.get_desc_workspace_bytes() + dynamic_sched_bytes = 4 if self.use_dynamic_sched else 0 + return desc_workspace_bytes + dynamic_sched_bytes + + @cute.jit + def _get_sched_counter_ptr(self, workspace_ptr): + counter_addr = workspace_ptr.toint() + self.get_desc_workspace_bytes() + return cute.make_ptr( + cutlass.Int32, + counter_addr, + AddressSpace.gmem, + assumed_align=4, + ) + + @cute.kernel + def helper_kernel( + self, + ptrs_b: cute.Pointer, + n: Int32, + k: Int32, + b_stride_size: cutlass.Int64, + b_major_mode: cutlass.Constexpr, + workspace_ptr, + tiled_mma_arg: cute.TiledMma, + b_smem_layout_arg, + cluster_layout_vmnk_shape_arg: cutlass.Constexpr, + ): + """Pre-main-kernel initialization. + + Launched with grid=(expert_cnt, 1, 1) for discrete mode, or + grid=(1, 1, 1) for dense+dynamic mode. + + Discrete weight: each block builds a B TMA descriptor for one expert. + Dynamic sched: block 0 resets the atomic tile counter to 0. + """ + expert_idx = cute.arch.block_idx()[0] + + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + # Discrete mode stores one B TMA descriptor per expert in the + # descriptor workspace. The main kernel later reuses those + # descriptors through the scheduler extension. + b_tma_op_arg = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma_arg.thr_id) + + b_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), cute.make_layout((self.expert_cnt,)) + ) + + c1 = cutlass.Int32(1) + c0 = cutlass.Int64(0) + c1_64 = 1 + if cutlass.const_expr(b_major_mode == OperandMajorMode.K): + stride_n = b_stride_size + stride_k = c1_64 + else: + stride_n = c1_64 + stride_k = b_stride_size + + b_ptr_val = b_ptr_tensor[expert_idx] + b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) + b_tensor_i = cute.make_tensor( + b_ptr, + cute.make_layout((n, k, c1), stride=(stride_n, stride_k, c0)), + ) + tma_atom_b, _ = cute.nvgpu.make_tiled_tma_atom_B( + b_tma_op_arg, + b_tensor_i, + b_smem_layout_arg, + self.mma_tiler, + tiled_mma_arg, + cluster_layout_vmnk_shape_arg, + ) + + workspace = TensormapWorkspace(workspace_ptr, ["b"]) + store_tma_desc(tma_atom_b, workspace.get_ptr("b", expert_idx)) + + if cutlass.const_expr(self.use_dynamic_sched): + if expert_idx == cutlass.Int32(0): + sched_counter = cute.make_tensor( + self._get_sched_counter_ptr(workspace_ptr), + cute.make_layout(1), + ) + sched_counter[0] = cutlass.Int32(0) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b, # Dense: cute.Tensor (N,K,L) | Discrete: cute.Pointer to int64[] + n: Int32, # Ignored for dense mode + k: Int32, # Ignored for dense mode + b_stride_size: cutlass.Int64, # Ignored for dense mode + b_major_mode: cutlass.Constexpr, # Ignored for dense mode + workspace_ptr, + c: cute.Tensor, + d: cute.Tensor, + padded_offsets: cute.Tensor, + alpha: cute.Tensor, + prob: cute.Tensor, + bias: Optional[cute.Tensor], + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + linear_offset: cutlass.Float32 = 0.0, + ): + """Execute the GEMM. + + Dense mode: ``b`` is a 3-D cute.Tensor (N, K, L). + Discrete mode: ``b`` is a cute.Pointer to a device int64[] array of + per-expert base addresses; ``n``, ``k``, ``b_stride_size``, and + ``b_major_mode`` describe the uniform per-expert layout. + """ + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = a.element_type + self.c_dtype: Type[cutlass.Numeric] = c.element_type + self.d_dtype: Type[cutlass.Numeric] = d.element_type + self.bias_dtype = bias.element_type if cutlass.const_expr(self.enable_bias) else cutlass.BFloat16 + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + self.d_layout = utils.LayoutEnum.from_tensor(d) + + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + else: + self.b_major_mode = b_major_mode + + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + self._setup_attributes() + + # ---- B setup (mode-dependent) ---- + # Dense mode receives a normal (N,K,L) tensor. Discrete mode receives a + # device array of per-expert B base addresses; we build a template tensor + # here only so the TMA atom has the right dtype/layout at compile time. + b_from_call_arg = b + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + c1 = cutlass.Int32(1) + c0 = cutlass.Int64(0) + c1_64 = 1 + if cutlass.const_expr(b_major_mode == OperandMajorMode.K): + b_template_stride = (b_stride_size, c1_64, c0) + else: + b_template_stride = (c1_64, b_stride_size, c0) + b_template_layout = cute.make_layout((n, k, c1), stride=b_template_stride) + b_ptr_typed = cute.make_ptr(self.b_dtype, b.toint(), AddressSpace.gmem, assumed_align=16) + b = cute.make_tensor(b_ptr_typed, b_template_layout) + + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # Setup TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + # Setup TMA store for C + c_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c, + c_smem_layout, + self.epi_tile_c, + ) + + # Setup TMA store for D + d_smem_layout = cute.slice_(self.d_smem_layout_staged, (None, None, 0)) + tma_atom_d, tma_tensor_d = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + d, + d_smem_layout, + self.epi_tile, + ) + # ---- Helper kernel: TMA desc init (discrete) + sched counter reset (dynamic) ---- + _need_helper = cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE or self.use_dynamic_sched) + if cutlass.const_expr(_need_helper): + _helper_grid_x = self.expert_cnt if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else 1 + _helper_args = ( + b_from_call_arg if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cute.make_ptr(cutlass.Int64, 0, AddressSpace.gmem), + n if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int32(0), + k if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int32(0), + b_stride_size if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int64(0), + b_major_mode if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else self.b_major_mode, + workspace_ptr, + tiled_mma, + b_smem_layout, + self.cluster_layout_vmnk.shape, + ) + self.helper_kernel(*_helper_args).launch( + grid=(_helper_grid_x, 1, 1), + block=(1, 1, 1), + stream=stream, + min_blocks_per_mp=1, + ) + + # ---- Grid computation via MoE scheduler ---- + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + b_n, b_k, b_l = cute.shape(b) + sched_expert_shape = (self.expert_cnt, b_n, b_k) + else: + sched_expert_shape = (self.expert_cnt, n, k) + + sched_params = MoESchedulerParams( + scenario="2Dx3D", + expert_shape=sched_expert_shape, + cta_tile_shape_mnk=self.cta_tile_shape_mnk, + cluster_shape_mn=self.cluster_shape_mn, + use_dynamic_sched=self.use_dynamic_sched, + ) + self.sched_params, grid = compute_grid( + sched_params, + max_active_clusters, + self.use_2cta_instrs, + ) + + self.buffer_align_bytes = 1024 + + # Define shared storage for kernel + SchedulerStorage = MoEPersistentTileScheduler.make_storage_struct(self.num_tile_stage, self.use_dynamic_sched) + + @cute.struct + class SharedStorage: + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + bias_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_bias_stage * 2 if self.enable_bias else 1] + scheduler: SchedulerStorage + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.cosize(self.c_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + sD: cute.struct.Align[ + cute.struct.MemRange[ + self.d_dtype, + cute.cosize(self.d_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # (MMA, MMA_M, MMA_K, STAGE) + sA: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + # (MMA, MMA_N, MMA_K, STAGE) + sB: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + # Bias SMEM: (tile_N, num_bias_stage) BF16 double-buffered + sBias: cute.struct.Align[ + cute.struct.MemRange[self.bias_dtype, cute.cosize(self.bias_smem_layout_staged)], + 16, + ] + + self.shared_storage = SharedStorage + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + tma_atom_d, + tma_tensor_d, + padded_offsets, + alpha, + bias, + prob, + workspace_ptr, # Contains per-expert B TMA descriptors + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.d_smem_layout_staged, + self.bias_smem_layout_staged, + self.epi_tile, + self.sched_params, + linear_offset, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + max_number_threads=[self.threads_per_cta, 1, 1], + smem=self.shared_storage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + return + + # ------------------------------------------------------------------ + # Internal: create extension based on weight_mode + # ------------------------------------------------------------------ + + @cute.jit + def _make_extension(self, workspace_ptr): + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + desc_workspace = TensormapWorkspace(workspace_ptr, ["b"]) + return DiscreteWeightGroupedGemmSchedExtension(tensormap_ctor=desc_workspace) + else: + return ContiguousGroupedGemmSchedExtension() + + @cute.jit + def store_c( + self, + tiled_copy_r2s, + tma_atom_c, + warp_idx, + tTR_rAcc, + tTR_rAcc_up, + tRS_rC, + tRS_sC, + bSG_gC, + bSG_sC, + c_pipeline, + prev_subtile_idx, + real_subtile_idx, + ) -> None: + c_buffer = prev_subtile_idx % self.num_c_stage + tRS_rC.store(tTR_rAcc.load().to(self.c_dtype)) + cute.copy( + tiled_copy_r2s, + tRS_rC[(None, None, 0)], + tRS_sC[(None, None, 0, c_buffer)], + ) + tRS_rC.store(tTR_rAcc_up.load().to(self.c_dtype)) + cute.copy( + tiled_copy_r2s, + tRS_rC[(None, None, 0)], + tRS_sC[(None, None, 1, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier.arrive_and_wait() + # + # TMA store smem to global memory + # + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, real_subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + + @cute.jit + def geglu_act(self, tCompute: cute.Tensor, acc_vec_up: cute.Tensor, acc_vec_gate: cute.Tensor, mProb: cute.Tensor, linear_offset: cutlass.Float32 = 1.0): + if cutlass.const_expr(self.vectorized_f32): + # GeGlu Packed Version + LOG2_E = cutlass.Float32(1.4426950408889634) + for i in cutlass.range_constexpr(0, cute.size(tCompute), 2): + + scaled_gate_0, scaled_gate_1 = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), + (1.702, 1.702), + rnd="rn", + ftz=False, + ) + + tCompute_log2e = cute.arch.mul_packed_f32x2( + (scaled_gate_0, scaled_gate_1), + (-LOG2_E, -LOG2_E), + rnd="rn", + ftz=False, + ) + + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.add_packed_f32x2( + ( + cute.math.exp2(tCompute_log2e[0], fastmath=True), + cute.math.exp2(tCompute_log2e[1], fastmath=True), + ), + (1.0, 1.0), + ) + + tCompute[i] = cute.arch.rcp_approx(tCompute[i]) + tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1]) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_gate[i + 0], acc_vec_gate[i + 1]), + rnd="rn", + ftz=False, + ) + ( + up_with_offset0, + up_with_offset1, + ) = cute.arch.add_packed_f32x2( + (linear_offset, linear_offset), + (acc_vec_up[i + 0], acc_vec_up[i + 1]), + rnd="rn", + ftz=False, + ) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (up_with_offset0, up_with_offset1), + rnd="rn", + ftz=False, + ) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (mProb, mProb), + rnd="rn", + ftz=False, + ) + else: + # GeGlu Unpacked Version + for i in cutlass.range_constexpr(cute.size(tCompute)): + tCompute[i] = (acc_vec_up[i] + linear_offset) * silu_f32_geglu_scaled(acc_vec_gate[i], fastmath=True) + tCompute[i] = tCompute[i] * mProb + + @cute.jit + def swiglu_act(self, tCompute: cute.Tensor, acc_vec_up: cute.Tensor, acc_vec_gate: cute.Tensor, mProb: cute.Tensor): + if cutlass.const_expr(self.vectorized_f32): + # SwiGlu Packed Version + LOG2_E = cutlass.Float32(1.4426950408889634) + for i in cutlass.range_constexpr(0, cute.size(tCompute), 2): + tCompute_log2e = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), + (-LOG2_E, -LOG2_E), + rnd="rn", + ftz=False, + ) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.add_packed_f32x2( + ( + cute.math.exp2(tCompute_log2e[0], fastmath=True), + cute.math.exp2(tCompute_log2e[1], fastmath=True), + ), + (1.0, 1.0), + ) + tCompute[i] = cute.arch.rcp_approx(tCompute[i]) + tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1]) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_gate[i + 0], acc_vec_gate[i + 1]), + rnd="rn", + ftz=False, + ) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_up[i], acc_vec_up[i + 1]), + rnd="rn", + ftz=False, + ) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (mProb, mProb), + rnd="rn", + ftz=False, + ) + else: + # SwiGlu Unpacked Version + for i in cutlass.range_constexpr(cute.size(tCompute)): + tCompute[i] = acc_vec_up[i] * silu_f32(acc_vec_gate[i], fastmath=True) + tCompute[i] = tCompute[i] * mProb + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + tma_atom_d: cute.CopyAtom, + mD_mnl: cute.Tensor, + padded_offsets: cute.Tensor, + alpha: cute.Tensor, + mBias_nl: Optional[cute.Tensor], + prob: cute.Tensor, + workspace_ptr, # Pointer to TMA descriptor workspace (from desc_init_kernel) + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + d_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + bias_smem_layout_staged: cute.Layout, + epi_tile: cute.Tile, + sched_params: MoESchedulerParams, + linear_offset: cutlass.Float32 = 0.0, + ): + """ + GPU device kernel performing the Persistent batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + cpasync.prefetch_descriptor(tma_atom_b) + if cutlass.const_expr(self.generate_c): + cpasync.prefetch_descriptor(tma_atom_c) + cpasync.prefetch_descriptor(tma_atom_d) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + total_token = padded_offsets[self.expert_cnt - 1] + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + sched_storage = storage.scheduler + + # Initialize mainloop ab_pipeline (barrier) and states + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_warp_id) * (2 if use_2cta_instrs else 1) + acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # Initialize tile info pipeline (barrier) and states + tile_info_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * 1, + ) + tile_info_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_wo_sched, + ) + tile_info_pipeline = pipeline.PipelineAsync.create( + barrier_storage=sched_storage.tile_info_mbar.data_ptr(), + num_stages=self.num_tile_stage, + producer_group=tile_info_pipeline_producer_group, + consumer_group=tile_info_pipeline_consumer_group, + ) + + scheduler = MoEPersistentTileScheduler.create( + sched_params, + padded_offsets, + cute.arch.block_idx(), + cute.arch.grid_dim(), + counter_ptr=self._get_sched_counter_ptr(workspace_ptr), + sched_storage=sched_storage, + ) + scheduler.internal_init() + + # Bias pipeline + SMEM + if cutlass.const_expr(self.enable_bias): + bias_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp, + ) + bias_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_warp_id), + ) + bias_pipeline = pipeline.PipelineCpAsync.create( + barrier_storage=storage.bias_mbar_ptr.data_ptr(), + num_stages=self.num_bias_stage, + producer_group=bias_pipeline_producer_group, + consumer_group=bias_pipeline_consumer_group, + ) + sBias = storage.sBias.get_tensor(bias_smem_layout_staged) + # (MMA_N, loopN, loopL) + gBias_nl = cute.local_tile(mBias_nl, cute.slice_(self.mma_tiler[:2], (0, None)), (None, None)) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, + ) + + # Cluster arrive after barrier init + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_arrive_relaxed() + + # + # Setup smem tensor A/B/C/D + # + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner) + sD = storage.sD.get_tensor(d_smem_layout_staged.outer, swizzle=d_smem_layout_staged.inner) + # (MMA, MMA_M, MMA_K, STAGE) + sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner) + # (MMA, MMA_N, MMA_K, STAGE) + sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner) + # (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + info_layout = cute.make_layout((4, self.num_tile_stage), stride=(1, 4)) + sInfo = sched_storage.sInfo.get_tensor(info_layout) + + # + # Compute multicast mask for A/B buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2) + b_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/D + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + # + # Cluster wait before tensor memory alloc + # + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_wait() + else: + self.cta_sync_barrier.arrive_and_wait() + + if total_token <= 0: + cute.arch.nvvm.exit() + + # + # Specialized Schedule warp (MoE Persistent Tile Scheduler) + # + if warp_idx == self.sched_warp_id: + work_tile_info = scheduler.initial_work_tile_info() + + tile_info_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_tile_stage) + + while work_tile_info.is_valid_tile: + # Write MoEWorkTileInfo directly to sInfo: + # sInfo[0] = expert_idx (>= 0 means valid) + # sInfo[1] = tile_m_idx (CTA-level M tile index) + # sInfo[2] = tile_n_idx + # sInfo[3] = k_tile_cnt + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = work_tile_info.expert_idx + sInfo[(1, tile_info_producer_state.index)] = work_tile_info.tile_m_idx + sInfo[(2, tile_info_producer_state.index)] = work_tile_info.tile_n_idx + sInfo[(3, tile_info_producer_state.index)] = work_tile_info.k_tile_cnt + cute.arch.fence_proxy("async.shared", space="cta") + + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + + work_tile_info = scheduler.advance_to_next_work() + + # Send invalid tile signal: expert_idx = -1 + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = cutlass.Int32(-1) + sInfo[(1, tile_info_producer_state.index)] = cutlass.Int32(0) + sInfo[(2, tile_info_producer_state.index)] = cutlass.Int32(0) + sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + tile_info_pipeline.producer_tail(tile_info_producer_state) + + # + # Specialized TMA load warp + # + if warp_idx == self.tma_warp_id: + ext = self._make_extension(workspace_ptr) + + ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage) + + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + # Get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + # sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + k_tile_cnt = work_tile_info.k_tile_cnt + ext.update_expert_info(padded_offsets, work_tile_info.expert_idx) + + # Get per-expert real tensors + TMA desc ptrs via extension + real_a, _ = ext.get_gmem_tensor("a", mA_mkl, padded_offsets, work_tile_info) + real_b, desc_ptr_b = ext.get_gmem_tensor("b", mB_nkl, padded_offsets, work_tile_info) + + # local_tile on per-expert tensors + gA_mkl = cute.local_tile( + real_a, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + gB_nkl = cute.local_tile( + real_b, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + + # MMA partition + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + + # TMA partition A + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + # TMA partition B + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # Convert CTA tile index to MMA tile index (matching original kernel's bidx // cta_group_size) + mma_tile_coord_m = work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape) + mma_tile_coord_n = work_tile_info.tile_n_idx + tAgA_slice = tAgA[(None, mma_tile_coord_m, None, 0)] + tBgB_slice = tBgB[(None, mma_tile_coord_n, None, 0)] + + # Peek (try_wait) AB buffer empty + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + tAgA_k = tAgA_slice[(None, ab_producer_state.count)] + tBgB_k = tBgB_slice[(None, ab_producer_state.count)] + tAsA_pipe = tAsA[(None, ab_producer_state.index)] + tBsB_pipe = tBsB[(None, ab_producer_state.index)] + + tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state) + + # Conditionally wait for AB buffer empty + ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) + ab_producer_state_next = ab_producer_state.clone() + ab_producer_state_next.advance() + if ab_producer_state_next.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state_next) + else: + peek_ab_empty_status = cutlass.Boolean(1) + + # TMA load A (contiguous, global desc via domain_offset) + cute.copy( + tma_atom_a, + tAgA_k, + tAsA_pipe, + tma_bar_ptr=tma_bar, + mcast_mask=a_full_mcast_mask, + ) + # TMA load B (discrete, per-expert desc from workspace) + cute.copy( + tma_atom_b, + tBgB_k, + tBsB_pipe, + tma_bar_ptr=tma_bar, + mcast_mask=b_full_mcast_mask, + tma_desc_ptr=desc_ptr_b, + ) + # Peek (try_wait) AB buffer empty for next k_tile + ab_producer_state.advance() + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # + # Wait A/B buffer empty + # + ab_pipeline.producer_tail(ab_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Bar sync for retrieve tensor memory ptr from shared mem + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + ab_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_ab_stage) + acd_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) + + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + # Get the first tile info from pipeline (scheduler has filtered out tiles >= num_non_exiting_tiles) + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + k_tile_cnt = tile_info[3] + + # Peek (try_wait) AB buffer full for k_tile = 0 + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + # Peek (try_wait) Acc buffer empty for k_tile = 0 + acd_producer_state.reset_count() + peek_acc_empty_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_acc_empty_status = acc_pipeline.producer_try_acquire(acd_producer_state) + + # Convert CTA tile index to MMA tile index (matching original kernel's bidx // cta_group_size) + mma_tile_coord_mnl = ( + tile_info[1] // cute.size(tiled_mma.thr_id.shape), + tile_info[2], # tile_n_idx + tile_info[0], # expert_idx + ) + + # Get accumulator stage index + tCtAcc = tCtAcc_base[(None, None, None, acd_producer_state.index)] + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acd_producer_state, peek_acc_empty_status) + # + # Mma mainloop + # + + # + # Reset the ACCUMULATE field for each tile + # + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + + if is_leader_cta: + # Conditionally wait for AB buffer full + ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + + # tCtAcc += tCrA * tCrB + num_kblocks = cute.size(tCrA, mode=[2]) + ab_consumer_state_next = ab_consumer_state.clone() + ab_consumer_state_next.advance() + if ab_consumer_state_next.count < k_tile_cnt: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state_next) + + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord = ( + None, + None, + kblock_idx, + ab_consumer_state.index, + ) + + cute.gemm( + tiled_mma, + tCtAcc, + tCrA[kblock_coord], + tCrB[kblock_coord], + tCtAcc, + ) + # Enable accumulate on tCtAcc after first kblock + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + # Async arrive AB buffer empty + ab_pipeline.consumer_release(ab_consumer_state) + ab_consumer_state = ab_consumer_state_next + + # + # Async arrive accumulator buffer full(each kblock) + # + if is_leader_cta: + acc_pipeline.producer_commit(acd_producer_state) + + # Peek (try_wait) Acc buffer empty for k_tile = k_tile + 1 + acd_producer_state.advance() + if acd_producer_state.count < k_tile_cnt: + if is_leader_cta: + peek_acc_empty_status = acc_pipeline.producer_try_acquire(acd_producer_state) + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acd_producer_state) + + # + # Specialized bias load warp — cp.async 32-bit GMEM→SMEM + # + if cutlass.const_expr(self.enable_bias): + if warp_idx == self.bias_load_warp_id and total_token > 0: + bias_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_bias_stage) + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + # 128-bit cp.async: 32 threads × (128/dtype_bits) elements = tile_N per warp + bias_elems_per_thread = 128 // self.bias_dtype.width + bias_g2s_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + self.bias_dtype, + num_bits_per_copy=128, + ) + bias_g2s_tiled = cute.make_tiled_copy_tv( + bias_g2s_atom, + cute.make_layout((self.threads_per_warp,)), + cute.make_layout((bias_elems_per_thread,)), + ) + thr_bias_g2s = bias_g2s_tiled.get_slice(cute.arch.lane_idx()) + tBs_sBias = thr_bias_g2s.partition_D(sBias) + + # Predicate tensor for bias cp.async + bias_n_total = mBias_nl.shape[0] + tBpBias = cute.make_rmem_tensor(cute.make_layout((1,)), cutlass.Boolean) + + # Get first tile info from pipeline + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + bias_producer_state.reset_count() + + # sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + mma_n_coord = tile_info[2] + expert_idx = tile_info[0] + + gBias_tile = gBias_nl[(None, mma_n_coord, expert_idx)] + tBs_gBias = thr_bias_g2s.partition_S(gBias_tile) + + # Predicate: check if this thread's chunk is within N + tBpBias[0] = mma_n_coord * self.mma_tiler[1] + cute.arch.lane_idx() * bias_elems_per_thread < bias_n_total + + bias_pipeline.producer_acquire(bias_producer_state) + cute.copy(bias_g2s_tiled, tBs_gBias[(None, 0)], tBs_sBias[(None, 0, bias_producer_state.index)], pred=tBpBias) + bias_pipeline.producer_commit(bias_producer_state) + bias_producer_state.advance() + + # Get next tile info + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + bias_pipeline.producer_tail(bias_producer_state) + + # + # Specialized epilogue warps + # + if warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # + # Partition for epilogue (shape-only: use global tensor for invariant setup) + # + epi_tidx = tidx + thr_mma_epi = tiled_mma.get_slice(mma_tile_coord_v) + # D has half as many logical N columns as the MMA accumulator because + # adjacent 32-column blocks are interpreted as gate/up pairs. + gD_mnl_shape = cute.local_tile(mD_mnl, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None)) + tCgD_shape = thr_mma_epi.partition_C(gD_mnl_shape) + + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc_gate, + tTR_rAcc_up, + ) = self.epilog_tmem_copy_and_partition(epi_tidx, tCtAcc_base, tCgD_shape, epi_tile, use_2cta_instrs) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc_gate.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition(tiled_copy_t2r, tTR_rC, epi_tidx, sC) + + tTR_rD = cute.make_rmem_tensor(tTR_rAcc_gate.shape, self.d_dtype) + tiled_copy_r2s, tRS_rD, tRS_sD = self.epilog_smem_copy_and_partition(tiled_copy_t2r, tTR_rD, epi_tidx, sD) + + epi_ext = self._make_extension(workspace_ptr) + + acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) + + c_pipeline = None + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + d_pipeline = None + # Threads/warps participating in tma store pipeline + d_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + d_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_d_stage, + producer_group=d_producer_group, + ) + + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + # Get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + if cutlass.const_expr(self.enable_bias): + bias_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_bias_stage) + bias_s2r_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.bias_dtype, num_bits_per_copy=128) + tTR_rBias_gate = cute.make_rmem_tensor(cute.make_layout(self.epi_tile[1]), self.bias_dtype) + tTR_rBias_up = cute.make_rmem_tensor(cute.make_layout(self.epi_tile[1]), self.bias_dtype) + + num_prev_subtiles = cutlass.Int32(0) + while is_valid_tile: + # sInfo format: (expert_idx, tile_m_idx, tile_n_idx, k_tile_cnt) + epi_work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + mma_tile_coord_mnl = ( + epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape), + epi_work_tile_info.tile_n_idx, + cutlass.Int32(0), + ) + + expert_idx = epi_work_tile_info.expert_idx + alpha_val = alpha[expert_idx] + epi_ext.update_expert_info(padded_offsets, epi_work_tile_info.expert_idx) + + if cutlass.const_expr(self.enable_bias): + bias_consumer_state.reset_count() + bias_pipeline.consumer_wait(bias_consumer_state) + sBias_stage = sBias[(None, bias_consumer_state.index)] + sBias_subtiles = cute.flat_divide(sBias_stage, cute.make_layout(2 * self.epi_tile[1])) + + # Get per-expert C/D tensors via extension + real_c, _ = epi_ext.get_gmem_tensor("c", mC_mnl, padded_offsets, epi_work_tile_info) + real_d, _ = epi_ext.get_gmem_tensor("d", mD_mnl, padded_offsets, epi_work_tile_info) + + # local_tile + partition on per-expert tensors + thr_mma_epi_loop = tiled_mma.get_slice(mma_tile_coord_v) + gC_mnl = cute.local_tile(real_c, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) + tCgC = thr_mma_epi_loop.partition_C(gC_mnl) + _, bSG_sC, bSG_gC_partitioned = epilog_gmem_copy_and_partition(epi_tidx, tma_atom_c, tCgC, self.epi_tile_c, sC) + + gD_mnl_loop = cute.local_tile(real_d, cute.slice_(self.mma_tiler_d, (None, None, 0)), (None, None, None)) + tCgD_loop = thr_mma_epi_loop.partition_C(gD_mnl_loop) + _, bSG_sD, bSG_gD_partitioned = epilog_gmem_copy_and_partition(epi_tidx, tma_atom_d, tCgD_loop, epi_tile, sD) + + # Slice to per-expert tile coords (L=0, domain already offset'd) + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + bSG_gD = bSG_gD_partitioned[(None, None, None, *mma_tile_coord_mnl)] + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + bSG_gD = cute.group_modes(bSG_gD, 1, cute.rank(bSG_gD)) + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] + + # + # Get PROB (per-expert via domain_offset) + # Note, it always assumes T2R_M/EPI_M is 1, otherwise it will break the result. + # + real_prob, _ = epi_ext.get_gmem_tensor("prob", prob, padded_offsets, epi_work_tile_info) + mPosition = ( + (epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape)) * self.mma_tiler[0] + + mma_tile_coord_v * (self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape)) + + tidx + ) + mProb = real_prob[mPosition, 0, 0] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + + # + # Store accumulator to global memory in subtiles + # + # Each loop consumes two adjacent accumulator subtiles: + # gate -> C and activation input + # up -> C and activation input + # C receives the pre-activation values for debugging/reference; + # D receives the final GLU result. + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in cutlass.range(0, subtile_cnt, 2, unroll=1): + real_subtile_idx = subtile_idx // 2 + + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn_gate = tTR_tAcc[(None, None, None, real_subtile_idx * 2)] + tTR_tAcc_mn_up = tTR_tAcc[(None, None, None, real_subtile_idx * 2 + 1)] + + cute.copy(tiled_copy_t2r, tTR_tAcc_mn_gate, tTR_rAcc_gate) + cute.copy(tiled_copy_t2r, tTR_tAcc_mn_up, tTR_rAcc_up) + + # + # Apply alpha (+ bias if enabled) + # + if cutlass.const_expr(self.enable_bias): + sBias_sub = sBias_subtiles[(None, real_subtile_idx)] + for i in cutlass.range_constexpr(self.epi_tile[1]): + tTR_rBias_gate[i] = sBias_sub[i] + tTR_rBias_up[i] = sBias_sub[self.epi_tile[1] + i] + bias_vec_gate = tTR_rBias_gate.load() + bias_vec_up = tTR_rBias_up.load() + + if cutlass.const_expr(self.vectorized_f32): + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc_gate), 2): + bias_gate_f32_0 = bias_vec_gate[i].to(cutlass.Float32) + bias_gate_f32_1 = bias_vec_gate[i + 1].to(cutlass.Float32) + bias_up_f32_0 = bias_vec_up[i].to(cutlass.Float32) + bias_up_f32_1 = bias_vec_up[i + 1].to(cutlass.Float32) + tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1] = cute.arch.fma_packed_f32x2( + (tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + (bias_gate_f32_0, bias_gate_f32_1), + rnd="rn", + ftz=False, + ) + tTR_rAcc_up[i], tTR_rAcc_up[i + 1] = cute.arch.fma_packed_f32x2( + (tTR_rAcc_up[i], tTR_rAcc_up[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + (bias_up_f32_0, bias_up_f32_1), + rnd="rn", + ftz=False, + ) + else: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc_gate)): + tTR_rAcc_gate[i] = tTR_rAcc_gate[i] * cutlass.Float32(alpha_val) + bias_vec_gate[i].to(cutlass.Float32) + tTR_rAcc_up[i] = tTR_rAcc_up[i] * cutlass.Float32(alpha_val) + bias_vec_up[i].to(cutlass.Float32) + + if subtile_idx == subtile_cnt - 2: + bias_pipeline.consumer_release(bias_consumer_state) + bias_consumer_state.advance() + else: + if cutlass.const_expr(self.vectorized_f32): + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc_gate), 2): + tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1] = cute.arch.mul_packed_f32x2( + (tTR_rAcc_gate[i], tTR_rAcc_gate[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + rnd="rn", + ftz=False, + ) + tTR_rAcc_up[i], tTR_rAcc_up[i + 1] = cute.arch.mul_packed_f32x2( + (tTR_rAcc_up[i], tTR_rAcc_up[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + rnd="rn", + ftz=False, + ) + else: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc_gate)): + tTR_rAcc_gate[i] = tTR_rAcc_gate[i] * cutlass.Float32(alpha_val) + tTR_rAcc_up[i] = tTR_rAcc_up[i] * cutlass.Float32(alpha_val) + + if cutlass.const_expr(self.generate_c): + self.store_c( + tiled_copy_r2s, + tma_atom_c, + warp_idx, + tTR_rAcc_gate, + tTR_rAcc_up, + tRS_rC, + tRS_sC, + bSG_gC, + bSG_sC, + c_pipeline, + num_prev_subtiles, + real_subtile_idx, + ) + + if cutlass.const_expr(self.act_func == "geglu"): + geglu_max_val = cutlass.Float32(7.0) + geglu_min_val = cutlass.Float32(-7.0) + for i in cutlass.range_constexpr(cute.size(tTR_rAcc_up)): + tTR_rAcc_gate[i] = fmin(tTR_rAcc_gate[i], geglu_max_val) + tTR_rAcc_up[i] = fmin(tTR_rAcc_up[i], geglu_max_val) + tTR_rAcc_up[i] = fmax(tTR_rAcc_up[i], geglu_min_val) + + acc_vec_gate = tTR_rAcc_gate.load() + acc_vec_up = tTR_rAcc_up.load() + + # SwiGlu or GeGLU + tCompute = cute.make_rmem_tensor(acc_vec_gate.shape, self.acc_dtype) + if cutlass.const_expr(self.act_func == "geglu"): + self.geglu_act(tCompute, acc_vec_up, acc_vec_gate, mProb, linear_offset) + elif cutlass.const_expr(self.act_func == "swiglu"): + self.swiglu_act(tCompute, acc_vec_up, acc_vec_gate, mProb) + + # + # Convert to D type + # + acc_vec = tiled_copy_r2s.retile(tCompute).load() + tRS_rD.store(acc_vec.to(self.d_dtype)) + + # + # Store D to shared memory + # + d_buffer = num_prev_subtiles % self.num_d_stage + num_prev_subtiles = num_prev_subtiles + 1 + cute.copy( + tiled_copy_r2s, + tRS_rD, + tRS_sD[(None, None, None, d_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier.arrive_and_wait() + # + # TMA store D to global memory + # + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_d, + bSG_sD[(None, d_buffer)], + bSG_gD[(None, real_subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + d_pipeline.producer_commit() + d_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + # + # Wait for C/D store complete + # + if cutlass.const_expr(self.generate_c): + c_pipeline.producer_tail() + d_pipeline.producer_tail() + + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gD_mnl: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[cutlass.Boolean, bool], + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param tAcc: The accumulator tensor to be copied and partitioned + :type tAcc: cute.Tensor + :param gD_mnl: The global tensor D + :type gD_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param use_2cta_instrs: Whether use_2cta_instrs is enabled + :type use_2cta_instrs: bool + + :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc_gate, tTR_rAcc_up) where: + - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + - tTR_tAcc: The partitioned accumulator tensor + - tTR_rAcc_gate: The partitioned accumulator tensor for acc gate + - tTR_rAcc_up: The partitioned accumulator tensor for acc up + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor] + """ + # Make tiledCopy for tensor memory load + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.d_layout, + self.d_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc[((None, None), 0, 0, None)], + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL) + gD_mnl_epi = cute.flat_divide(gD_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL) + tTR_gC = thr_copy_t2r.partition_D(gD_mnl_epi) + + # (T2R, T2R_M, T2R_N) + tTR_rAcc_gate = cute.make_rmem_tensor(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype) + # (T2R, T2R_M, T2R_N) + tTR_rAcc_up = cute.make_rmem_tensor(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc_gate, tTR_rAcc_up + + def epilog_smem_copy_and_partition( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sD: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for shared memory store, then use it to partition register array (source) and shared memory (destination). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sD: The shared memory tensor to be copied and partitioned + :type sD: cute.Tensor + + :return: A tuple containing (tiled_copy_r2s, tRS_rD, tRS_sD) where: + - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) + - tRS_rD: The partitioned tensor D (register source) + - tRS_sD: The partitioned tensor D (smem destination) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_r2s = sm100_utils.get_smem_store_op(self.d_layout, self.d_dtype, self.acc_dtype, tiled_copy_t2r) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sD = thr_copy_r2s.partition_D(sD) + # (R2S, R2S_M, R2S_N) + tRS_rD = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rD, tRS_sD diff --git a/python/cudnn/grouped_gemm/grouped_gemm_unfused/__init__.py b/python/cudnn/grouped_gemm/grouped_gemm_unfused/__init__.py new file mode 100644 index 000000000..c42539b40 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_unfused/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +from .api import GroupedGemmSm100, grouped_gemm_wrapper_sm100 + +__all__ = ["GroupedGemmSm100", "grouped_gemm_wrapper_sm100"] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_unfused/_bf16_api.py b/python/cudnn/grouped_gemm/grouped_gemm_unfused/_bf16_api.py new file mode 100644 index 000000000..75f80533d --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_unfused/_bf16_api.py @@ -0,0 +1,559 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Private APIBase API for the source-close unfused BF16 MoE kernel.""" + +import os +import weakref +from typing import Optional, Tuple + +import torch +import cutlass +import cutlass.cute as cute +from cuda.bindings import driver as cuda +from cutlass.cute.nvgpu import OperandMajorMode +from cutlass.cute.runtime import from_dlpack, make_fake_stream + +from cudnn.api_base import APIBase, TensorDesc +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor + +from ..moe_utils import MoEWeightMode +from .moe_grouped_gemm import MoEGroupedGemmBf16Kernel + +_OUTPUT_DTYPES = [torch.bfloat16, torch.float16, torch.float32] + + +class GroupedGemmBf16API(APIBase): + """Descriptor-first API for :class:`MoEGroupedGemmBf16Kernel`.""" + + def __init__( + self, + sample_a: torch.Tensor, + sample_c: torch.Tensor, + sample_d: torch.Tensor, + sample_padded_offsets: torch.Tensor, + sample_alpha: torch.Tensor, + sample_b: Optional[torch.Tensor] = None, + sample_bias: Optional[torch.Tensor] = None, + sample_prob: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + b_shape: Optional[Tuple[int, ...]] = None, + b_dtype: Optional[torch.dtype] = None, + acc_dtype: torch.dtype = torch.float32, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + vector_f32: bool = False, + m_aligned: int = 256, + generate_c: bool = False, + b_major: str = "k", + use_dynamic_sched: bool = False, + ) -> None: + super().__init__() + self._warn_experimental_api() + + if sample_b is not None and num_experts is None: + self.weight_mode = MoEWeightMode.DENSE + elif sample_b is None and num_experts is not None: + self.weight_mode = MoEWeightMode.DISCRETE + if b_shape is None or b_dtype is None: + raise ValueError("b_shape and b_dtype are required in discrete mode") + else: + raise ValueError("Provide sample_b for dense mode or (num_experts, b_shape, b_dtype) " "for discrete mode, but not both") + + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") + self.c_desc = self._make_tensor_desc(sample_c, name="sample_c") + self.d_desc = self._make_tensor_desc(sample_d, name="sample_d") + self.padded_offsets_desc = self._make_tensor_desc(sample_padded_offsets, name="sample_padded_offsets") + self.alpha_desc = self._make_tensor_desc(sample_alpha, name="sample_alpha") + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") + self.bias_desc = self._make_tensor_desc(sample_bias, name="sample_bias") + self.prob_desc = self._make_tensor_desc(sample_prob, name="sample_prob") + self._sample_offset_values = self._copy_values_to_host(sample_padded_offsets) + self._sample_offsets_ref = weakref.ref(sample_padded_offsets) + self._sample_offsets_version = int(sample_padded_offsets._version) + self._sample_data_ptrs = { + name: tensor.data_ptr() + for name, tensor in ( + ("sample_a", sample_a), + ("sample_b", sample_b), + ("sample_c", sample_c), + ("sample_d", sample_d), + ("sample_padded_offsets", sample_padded_offsets), + ("sample_alpha", sample_alpha), + ("sample_bias", sample_bias), + ("sample_prob", sample_prob), + ) + if tensor is not None + } + + self.expert_cnt = self.b_desc.shape[2] if self.weight_mode == MoEWeightMode.DENSE and self.b_desc.ndim == 3 else int(num_experts or 0) + self.b_shape = tuple(b_shape) if b_shape is not None else None + self.b_dtype = b_dtype if b_dtype is not None else self.b_desc.dtype + self.acc_dtype = acc_dtype + self.mma_tiler_mn = tuple(mma_tiler_mn) + self.use_2cta_instrs = self.mma_tiler_mn[0] == 256 + self.cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if self.use_2cta_instrs else (1, 1))) + self.vector_f32 = vector_f32 + self.m_aligned = m_aligned + self.generate_c = generate_c + self.b_major = b_major + self.use_dynamic_sched = use_dynamic_sched + self._has_bias = self.bias_desc is not None + self._kernel = MoEGroupedGemmBf16Kernel + self._workspace: Optional[torch.Tensor] = None + self._compile_b_ptrs: Optional[torch.Tensor] = None + self._validated_offsets: dict[int, tuple] = {} + self._validated_pointer_values: dict[int, tuple] = {} + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + + @staticmethod + def _expect_shape(desc: TensorDesc, expected: Tuple[int, ...], name: str) -> None: + if desc.shape != expected: + raise ValueError(f"{name} shape mismatch: expected {expected}, got {desc.shape}") + + @staticmethod + def _expect_stride(desc: TensorDesc, expected: Tuple[int, ...], name: str) -> None: + if desc.stride != expected: + raise ValueError(f"{name} must use the source-compatible layout with stride " f"{expected}, got {desc.stride}") + + @staticmethod + def _expect_device(desc: TensorDesc, device: torch.device, name: str) -> None: + if desc.device != device: + raise ValueError(f"{name} must be on {device}, got {desc.device}") + + @staticmethod + def _copy_values_to_host(tensor: torch.Tensor) -> Tuple[int, ...]: + return tuple(int(value) for value in tensor.detach().cpu().tolist()) + + @staticmethod + def _is_validation_cached(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> bool: + cached = cache.get(id(tensor)) + return bool(cached and cached[0]() is tensor and cached[1] == int(tensor._version) and cached[2] == extra) + + @staticmethod + def _remember_validation(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> None: + key = id(tensor) + + def discard(_reference, *, cache=cache, key=key): + cache.pop(key, None) + + cache[key] = ( + weakref.ref(tensor, discard), + int(tensor._version), + extra, + ) + + def _validate_offsets_once(self, offsets: torch.Tensor, *, tensor_m: int) -> None: + extra = (self.expert_cnt, tensor_m) + if self._is_validation_cached(self._validated_offsets, offsets, extra): + return + values = self._copy_values_to_host(offsets) + self._validate_offset_sequence(values, expert_cnt=self.expert_cnt, tensor_m=tensor_m) + self._remember_validation(self._validated_offsets, offsets, extra) + + def _validate_pointer_values_once(self, b_ptrs: torch.Tensor) -> None: + if self._is_validation_cached(self._validated_pointer_values, b_ptrs, self.expert_cnt): + return + pointer_values = self._copy_values_to_host(b_ptrs) + if any(value == 0 or value % 16 != 0 for value in pointer_values): + raise ValueError("b_ptrs entries must be non-null and 16-byte aligned") + self._remember_validation(self._validated_pointer_values, b_ptrs, self.expert_cnt) + + @staticmethod + def _validate_offset_sequence(values: Tuple[int, ...], *, expert_cnt: int, tensor_m: int) -> None: + if len(values) != expert_cnt: + raise ValueError(f"padded_offsets length mismatch: expected {expert_cnt}, got {len(values)}") + previous = 0 + for index, value in enumerate(values): + if value < previous: + raise ValueError("padded_offsets must be a non-decreasing cumulative sum; " f"index {index} is {value} after {previous}") + if value % MoEGroupedGemmBf16Kernel.FIX_PAD_SIZE != 0: + raise ValueError(f"padded_offsets[{index}] must be 256-aligned, got {value}") + previous = value + if not values or values[-1] <= 0 or values[-1] > tensor_m: + raise ValueError(f"padded_offsets last value must be in [1, {tensor_m}], got " f"{values[-1] if values else None}") + + @staticmethod + def _validate_data_alignment(tensor: torch.Tensor, name: str) -> None: + if tensor.data_ptr() % 16 != 0: + raise ValueError(f"{name} data pointer must be 16-byte aligned") + + @staticmethod + def _validate_pointer_array_alignment(tensor: torch.Tensor) -> None: + if tensor.data_ptr() % 8 != 0: + raise ValueError("b_ptrs data pointer must be 8-byte aligned") + + @staticmethod + def _record_pointer_stream(b_ptrs: torch.Tensor, current_stream: cuda.CUstream) -> None: + handle = int(current_stream) + torch_current = torch.cuda.current_stream(b_ptrs.device) + torch_default = torch.cuda.default_stream(b_ptrs.device) + if handle == torch_current.cuda_stream: + launch_stream = torch_current + elif handle == torch_default.cuda_stream: + launch_stream = torch_default + else: + launch_stream = torch.cuda.ExternalStream(handle, device=b_ptrs.device) + b_ptrs.record_stream(launch_stream) + + def check_support(self) -> bool: + if self.a_desc.ndim != 3: + raise ValueError(f"sample_a must be rank-3, got {self.a_desc.shape}") + tensor_m, k, one = self.a_desc.shape + if one != 1: + raise ValueError(f"sample_a trailing dimension must be 1, got {one}") + + if self.weight_mode == MoEWeightMode.DENSE: + if self.b_desc.ndim != 3: + raise ValueError(f"sample_b must be rank-3, got {self.b_desc.shape}") + n, b_k, experts = self.b_desc.shape + if b_k != k: + raise ValueError(f"sample_b K dimension ({b_k}) must match sample_a ({k})") + if experts != self.expert_cnt: + raise ValueError("sample_b expert dimension is inconsistent") + self._expect_stride(self.b_desc, (k, 1, n * k), "sample_b") + else: + if len(self.b_shape) not in (2, 3): + raise ValueError(f"b_shape must be rank-2 or rank-3, got {self.b_shape}") + n, b_k = self.b_shape[:2] + if len(self.b_shape) == 3 and self.b_shape[2] != 1: + raise ValueError(f"b_shape trailing dimension must be 1, got {self.b_shape}") + if b_k != k: + raise ValueError(f"b_shape K dimension ({b_k}) must match sample_a ({k})") + + self._expect_shape(self.c_desc, (tensor_m, n, 1), "sample_c") + self._expect_shape(self.d_desc, (tensor_m, n, 1), "sample_d") + self._expect_shape(self.padded_offsets_desc, (self.expert_cnt,), "sample_padded_offsets") + self._expect_shape(self.alpha_desc, (self.expert_cnt,), "sample_alpha") + if self.prob_desc is None: + raise ValueError("sample_prob is required") + self._expect_shape(self.prob_desc, (tensor_m, 1, 1), "sample_prob") + + self._expect_stride(self.a_desc, (k, 1, tensor_m * k), "sample_a") + self._expect_stride(self.c_desc, (n, 1, tensor_m * n), "sample_c") + self._expect_stride(self.d_desc, (n, 1, tensor_m * n), "sample_d") + self._expect_stride(self.padded_offsets_desc, (1,), "sample_padded_offsets") + self._expect_stride(self.alpha_desc, (1,), "sample_alpha") + self._expect_stride(self.prob_desc, (1, 1, 1), "sample_prob") + + self._check_dtype(self.a_desc, torch.bfloat16, "sample_a") + if self.weight_mode == MoEWeightMode.DENSE: + self._check_dtype(self.b_desc, torch.bfloat16, "sample_b") + self._check_dtype(self.b_dtype, torch.bfloat16, "b_dtype") + self._check_dtype(self.c_desc, _OUTPUT_DTYPES, "sample_c") + self._check_dtype(self.d_desc, _OUTPUT_DTYPES, "sample_d") + self._check_dtype(self.padded_offsets_desc, torch.int32, "sample_padded_offsets") + self._check_dtype(self.alpha_desc, torch.float32, "sample_alpha") + self._check_dtype(self.prob_desc, torch.float32, "sample_prob") + + device = self.a_desc.device + for desc, name in ( + (self.c_desc, "sample_c"), + (self.d_desc, "sample_d"), + (self.padded_offsets_desc, "sample_padded_offsets"), + (self.alpha_desc, "sample_alpha"), + (self.prob_desc, "sample_prob"), + ): + self._expect_device(desc, device, name) + if self.b_desc is not None: + self._expect_device(self.b_desc, device, "sample_b") + + if self.bias_desc is not None: + self._expect_shape(self.bias_desc, (n, self.expert_cnt), "sample_bias") + self._expect_stride(self.bias_desc, (1, n), "sample_bias") + self._check_dtype(self.bias_desc, _OUTPUT_DTYPES, "sample_bias") + self._expect_device(self.bias_desc, device, "sample_bias") + + for name, data_ptr in self._sample_data_ptrs.items(): + if data_ptr % 16 != 0: + raise ValueError(f"{name} data pointer must be 16-byte aligned") + + if self.acc_dtype != torch.float32: + raise ValueError(f"acc_dtype must be torch.float32, got {self.acc_dtype}") + if self.m_aligned != MoEGroupedGemmBf16Kernel.FIX_PAD_SIZE: + raise ValueError(f"m_aligned must be 256, got {self.m_aligned}") + if self.b_major != "k": + raise ValueError(f"b_major must be 'k' for the BF16 backend, got {self.b_major}") + if self.expert_cnt <= 0 or self.expert_cnt > 1024: + raise ValueError(f"expert count must be in [1, 1024], got {self.expert_cnt}") + if tensor_m % 256 != 0: + raise ValueError(f"sample_a M dimension must be 256-aligned, got {tensor_m}") + + self._validate_offset_sequence( + self._sample_offset_values, + expert_cnt=self.expert_cnt, + tensor_m=tensor_m, + ) + sample_offsets = self._sample_offsets_ref() + if sample_offsets is not None and int(sample_offsets._version) == self._sample_offsets_version: + self._remember_validation( + self._validated_offsets, + sample_offsets, + (self.expert_cnt, tensor_m), + ) + elif sample_offsets is not None: + self._validate_offsets_once(sample_offsets, tensor_m=tensor_m) + + if not self._kernel.can_implement( + _convert_to_cutlass_data_type(torch.bfloat16), + _convert_to_cutlass_data_type(self.c_desc.dtype), + _convert_to_cutlass_data_type(self.d_desc.dtype), + _convert_to_cutlass_data_type(self.acc_dtype), + self.use_2cta_instrs, + self.mma_tiler_mn, + self.cluster_shape_mn, + tensor_m, + n, + k, + self.expert_cnt, + "k", + "k", + "n", + self.m_aligned, + ): + raise ValueError("Unsupported BF16 grouped GEMM tile, cluster, alignment, or layout configuration") + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + major, minor = torch.cuda.get_device_capability(self.a_desc.device) + compute_capability = major * 10 + minor + if compute_capability < 100: + raise RuntimeError(f"GroupedGemmSm100 requires SM100+, found SM{compute_capability} " f"on {self.a_desc.device}") + + self._is_supported = True + return True + + def compile(self) -> None: + self._ensure_support_checked() + if self._compiled_kernel is not None: + return + + kernel = self._kernel( + acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), + use_2cta_instrs=self.use_2cta_instrs, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + vectorized_f32=self.vector_f32, + generate_c=self.generate_c, + enable_bias=self._has_bias, + expert_cnt=self.expert_cnt, + weight_mode=self.weight_mode, + use_dynamic_sched=self.use_dynamic_sched, + ) + + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - self.num_cluster_overlap_margin + if max_active_clusters <= 0: + raise ValueError("max_active_clusters must be > 0 after applying " "CUDNNFE_CLUSTER_OVERLAP_MARGIN") + + workspace_bytes = kernel.get_workspace_bytes() + self._workspace = torch.empty(max(workspace_bytes, 1), dtype=torch.uint8, device=self.a_desc.device) + if self._workspace.data_ptr() % 128 != 0: + raise RuntimeError("workspace allocation must be 128-byte aligned") + workspace_ptr = from_dlpack(self._workspace, assumed_align=128).iterator + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + + valid_m = cute.sym_int(divisibility=256) + a_fake = self._make_fake_cute_compact_tensor( + self.a_desc.dtype, + self.a_desc.shape, + self.a_desc.stride_order, + dynamic_mode=0, + divisibility=256, + ) + c_fake = self._make_fake_cute_compact_tensor( + self.c_desc.dtype, + self.c_desc.shape, + self.c_desc.stride_order, + dynamic_mode=0, + divisibility=256, + ) + d_fake = self._make_fake_cute_compact_tensor( + self.d_desc.dtype, + self.d_desc.shape, + self.d_desc.stride_order, + dynamic_mode=0, + divisibility=256, + ) + prob_fake = self._make_fake_cute_tensor( + self.prob_desc.dtype, + (valid_m, 1, 1), + self.prob_desc.stride, + ) + + if self.weight_mode == MoEWeightMode.DENSE: + b_fake = self._make_fake_cute_tensor_from_desc(self.b_desc) + n_value = cutlass.Int32(0) + k_value = cutlass.Int32(0) + b_stride = cutlass.Int64(0) + else: + self._compile_b_ptrs = torch.empty((self.expert_cnt,), dtype=torch.int64, device=self.a_desc.device) + self._validate_pointer_array_alignment(self._compile_b_ptrs) + b_fake = from_dlpack(self._compile_b_ptrs, assumed_align=8).iterator + n, k = self.b_shape[:2] + n_value = cutlass.Int32(n) + k_value = cutlass.Int32(k) + b_stride = cutlass.Int64(k) + + raw_compiled = cute.compile( + kernel, + a=a_fake, + b=b_fake, + n=n_value, + k=k_value, + b_stride_size=b_stride, + b_major_mode=OperandMajorMode.K, + workspace_ptr=workspace_ptr, + c=c_fake, + d=d_fake, + padded_offsets=self._make_fake_cute_tensor_from_desc(self.padded_offsets_desc), + alpha=self._make_fake_cute_tensor_from_desc(self.alpha_desc), + bias=self._make_fake_cute_tensor_from_desc(self.bias_desc), + prob=prob_fake, + max_active_clusters=max_active_clusters, + stream=fake_stream, + options="--enable-tvm-ffi", + ) + + cached_n = n_value + cached_k = k_value + cached_b_stride = b_stride + + def tensor_api( + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor], + b_ptrs: Optional[torch.Tensor], + bias_tensor: Optional[torch.Tensor], + prob_tensor: torch.Tensor, + stream: cuda.CUstream, + ) -> None: + b_arg = b_tensor if self.weight_mode == MoEWeightMode.DENSE else int(b_ptrs.data_ptr()) + raw_compiled( + a_tensor, + b_arg, + cached_n, + cached_k, + cached_b_stride, + workspace_ptr, + c_tensor, + d_tensor, + padded_offsets, + alpha_tensor, + bias_tensor, + prob_tensor, + stream, + ) + + self._compiled_kernel = tensor_api + + def _validate_live_tensor( + self, + tensor: torch.Tensor, + sample: TensorDesc, + name: str, + *, + dynamic_m: bool = False, + ) -> TensorDesc: + desc = self._make_tensor_desc(tensor, name=name) + if desc.dtype != sample.dtype: + raise ValueError(f"{name} dtype mismatch: expected {sample.dtype}, got {desc.dtype}") + if desc.device != sample.device: + raise ValueError(f"{name} device mismatch: expected {sample.device}, got {desc.device}") + if dynamic_m: + if desc.shape[1:] != sample.shape[1:]: + raise ValueError(f"{name} shape suffix mismatch: expected {sample.shape[1:]}, got {desc.shape[1:]}") + if desc.stride_order != sample.stride_order: + raise ValueError(f"{name} layout mismatch: expected stride order {sample.stride_order}, " f"got {desc.stride_order}") + elif desc.shape != sample.shape or desc.stride != sample.stride: + raise ValueError(f"{name} descriptor mismatch: expected shape/stride " f"{sample.shape}/{sample.stride}, got {desc.shape}/{desc.stride}") + return desc + + def execute( + self, + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor] = None, + b_ptrs: Optional[torch.Tensor] = None, + bias_tensor: Optional[torch.Tensor] = None, + prob_tensor: Optional[torch.Tensor] = None, + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + current_stream = self._get_default_stream(current_stream) + if self._compiled_kernel is None: + raise RuntimeError("Kernel not compiled; call compile() first") + if prob_tensor is None: + raise ValueError("prob_tensor is required") + + a_desc = self._validate_live_tensor(a_tensor, self.a_desc, "a_tensor", dynamic_m=True) + c_desc = self._validate_live_tensor(c_tensor, self.c_desc, "c_tensor", dynamic_m=True) + d_desc = self._validate_live_tensor(d_tensor, self.d_desc, "d_tensor", dynamic_m=True) + self._validate_live_tensor(padded_offsets, self.padded_offsets_desc, "padded_offsets") + self._validate_live_tensor(alpha_tensor, self.alpha_desc, "alpha_tensor") + prob_desc = self._validate_live_tensor(prob_tensor, self.prob_desc, "prob_tensor", dynamic_m=True) + + tensor_m, k, _ = a_desc.shape + n = d_desc.shape[1] + if tensor_m % 256 != 0: + raise ValueError(f"a_tensor M dimension must be 256-aligned, got {tensor_m}") + self._expect_shape(c_desc, (tensor_m, n, 1), "c_tensor") + self._expect_shape(d_desc, (tensor_m, n, 1), "d_tensor") + self._expect_shape(prob_desc, (tensor_m, 1, 1), "prob_tensor") + self._expect_stride(a_desc, (k, 1, tensor_m * k), "a_tensor") + self._expect_stride(c_desc, (n, 1, tensor_m * n), "c_tensor") + self._expect_stride(d_desc, (n, 1, tensor_m * n), "d_tensor") + self._expect_stride(prob_desc, (1, 1, 1), "prob_tensor") + self._validate_offsets_once(padded_offsets, tensor_m=tensor_m) + + for tensor, name in ( + (a_tensor, "a_tensor"), + (c_tensor, "c_tensor"), + (d_tensor, "d_tensor"), + (padded_offsets, "padded_offsets"), + (alpha_tensor, "alpha_tensor"), + (prob_tensor, "prob_tensor"), + ): + self._validate_data_alignment(tensor, name) + + if self._has_bias: + if bias_tensor is None: + raise ValueError("bias_tensor is required because the API was compiled with sample_bias") + self._validate_live_tensor(bias_tensor, self.bias_desc, "bias_tensor") + self._validate_data_alignment(bias_tensor, "bias_tensor") + elif bias_tensor is not None: + raise ValueError("bias_tensor must be omitted because the API was compiled without sample_bias") + + if self.weight_mode == MoEWeightMode.DENSE: + if b_tensor is None or b_ptrs is not None: + raise ValueError("Dense execution requires b_tensor and forbids b_ptrs") + self._validate_live_tensor(b_tensor, self.b_desc, "b_tensor") + self._validate_data_alignment(b_tensor, "b_tensor") + else: + if b_tensor is not None or b_ptrs is None: + raise ValueError("Discrete execution requires b_ptrs and forbids b_tensor") + _require_pointer_tensor(b_ptrs, "b_ptrs", self.expert_cnt) + if b_ptrs.device != self.a_desc.device: + raise ValueError(f"b_ptrs must be on the same device as a_tensor " f"({self.a_desc.device}), got {b_ptrs.device}") + self._validate_pointer_array_alignment(b_ptrs) + self._validate_pointer_values_once(b_ptrs) + self._record_pointer_stream(b_ptrs, current_stream) + + self._compiled_kernel( + a_tensor, + c_tensor, + d_tensor, + padded_offsets, + alpha_tensor, + b_tensor, + b_ptrs, + bias_tensor, + prob_tensor, + current_stream, + ) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_unfused/api.py b/python/cudnn/grouped_gemm/grouped_gemm_unfused/api.py new file mode 100644 index 000000000..972d463d7 --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_unfused/api.py @@ -0,0 +1,384 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Neutral public API for SM100 unfused BF16 grouped GEMM.""" + +from __future__ import annotations + +import os +from typing import Optional, Tuple + +import torch +from cuda.bindings import driver as cuda + +from cudnn.api_base import APIBase, TupleDict +from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor + + +from ._bf16_api import GroupedGemmBf16API +from ..moe_utils import MoEWeightMode + +__all__ = ["GroupedGemmBf16API"] + + +class GroupedGemmSm100(APIBase): + """Public lifecycle facade with deferred BF16 descriptor initialization.""" + + def __init__( + self, + sample_a: torch.Tensor, + sample_c: torch.Tensor, + sample_d: torch.Tensor, + sample_padded_offsets: torch.Tensor, + sample_alpha: torch.Tensor, + sample_b: Optional[torch.Tensor] = None, + sample_bias: Optional[torch.Tensor] = None, + sample_prob: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + b_shape: Optional[Tuple[int, ...]] = None, + b_dtype: Optional[torch.dtype] = None, + acc_dtype: torch.dtype = torch.float32, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + vector_f32: bool = False, + m_aligned: int = 256, + generate_c: bool = False, + b_major: str = "k", + use_dynamic_sched: bool = False, + ) -> None: + super().__init__() + self._pending_init_kwargs = dict(locals()) + self._pending_init_kwargs.pop("self") + self._pending_init_kwargs.pop("__class__", None) + self._implementation = None + + def check_support(self) -> bool: + if self._implementation is None: + self._implementation = GroupedGemmBf16API(**self._pending_init_kwargs) + self._kernel = self._implementation._kernel + supported = self._implementation.check_support() + self._is_supported = self._implementation._is_supported + if supported: + self._pending_init_kwargs = None + return supported + + def compile(self) -> None: + if self._implementation is None: + self.check_support() + if self._is_supported: + self._implementation._is_supported = True + self._implementation.compile() + self._is_supported = self._implementation._is_supported + self._compiled_kernel = self._implementation._compiled_kernel + + def execute( + self, + a_tensor: torch.Tensor, + c_tensor: torch.Tensor, + d_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor] = None, + b_ptrs: Optional[torch.Tensor] = None, + bias_tensor: Optional[torch.Tensor] = None, + prob_tensor: Optional[torch.Tensor] = None, + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + if self._implementation is None: + raise RuntimeError("Kernel not compiled; call compile() first") + self._implementation.execute( + a_tensor=a_tensor, + c_tensor=c_tensor, + d_tensor=d_tensor, + padded_offsets=padded_offsets, + alpha_tensor=alpha_tensor, + b_tensor=b_tensor, + b_ptrs=b_ptrs, + bias_tensor=bias_tensor, + prob_tensor=prob_tensor, + current_stream=current_stream, + ) + self._is_supported = self._implementation._is_supported + self._compiled_kernel = self._implementation._compiled_kernel + + +_cache_of_GroupedGemmSm100Objects: dict[tuple, GroupedGemmSm100] = {} + + +def _stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: + return tuple( + index + for index, _ in sorted( + enumerate(tensor.stride()), + key=lambda item: (item[1], tensor.shape[item[0]]), + ) + ) + + +def _tensor_signature(tensor: Optional[torch.Tensor], *, dynamic_m: bool = False) -> tuple: + if tensor is None: + return (None, None, None, None) + shape = (None, *tuple(tensor.shape[1:])) if dynamic_m else tuple(tensor.shape) + return ( + shape, + _stride_order(tensor), + tensor.dtype, + (tensor.device.type, tensor.device.index), + ) + + +def _validate_output( + tensor: torch.Tensor, + *, + name: str, + shape: Tuple[int, int, int], + stride: Tuple[int, int, int], + dtype: torch.dtype, + device: torch.device, +) -> None: + if tuple(tensor.shape) != shape or tuple(tensor.stride()) != stride or tensor.dtype != dtype or tensor.device != device: + raise ValueError( + f"{name} must have shape {shape}, stride {stride}, dtype {dtype}, " + f"device {device}; got shape {tuple(tensor.shape)}, stride " + f"{tuple(tensor.stride())}, dtype {tensor.dtype}, device {tensor.device}" + ) + + +def _normalize_call( + a_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + b_tensor: Optional[torch.Tensor], + b_ptrs: Optional[torch.Tensor], + n: Optional[int], + b_dtype: Optional[torch.dtype], + prob_tensor: Optional[torch.Tensor], + c_dtype: torch.dtype, + d_dtype: torch.dtype, + cd_major: str, + m_aligned: int, +) -> tuple[bool, int, int, int]: + is_dense = b_tensor is not None + is_discrete = b_ptrs is not None + if is_dense and is_discrete: + raise ValueError("Provide either b_tensor or b_ptrs, not both") + if not is_dense and not is_discrete: + raise ValueError("Must provide either b_tensor or b_ptrs") + + if a_tensor.dtype != torch.bfloat16: + raise ValueError(f"a_tensor must have dtype torch.bfloat16, got {a_tensor.dtype}") + if a_tensor.ndim != 3 or a_tensor.shape[2] != 1: + raise ValueError(f"a_tensor must have shape (m, k, 1), got {tuple(a_tensor.shape)}") + if prob_tensor is None: + raise ValueError("prob_tensor is required") + if cd_major != "n": + raise ValueError(f"cd_major must be 'n', got {cd_major}") + if c_dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise ValueError(f"c_dtype must be BF16, FP16, or FP32, got {c_dtype}") + if d_dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise ValueError(f"d_dtype must be BF16, FP16, or FP32, got {d_dtype}") + if m_aligned != 256: + raise ValueError(f"m_aligned must be 256, got {m_aligned}") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + major, minor = torch.cuda.get_device_capability(a_tensor.device) + compute_capability = major * 10 + minor + if compute_capability < 100: + raise RuntimeError(f"GroupedGemmSm100 requires SM100+, found SM{compute_capability}") + + tensor_m, k, _ = a_tensor.shape + if tensor_m % 256 != 0: + raise ValueError(f"a_tensor M dimension must be 256-aligned, got {tensor_m}") + if tuple(prob_tensor.shape) != (tensor_m, 1, 1): + raise ValueError(f"prob_tensor must have shape {(tensor_m, 1, 1)}, got " f"{tuple(prob_tensor.shape)}") + if prob_tensor.dtype != torch.float32: + raise ValueError(f"prob_tensor must have dtype torch.float32, got {prob_tensor.dtype}") + if is_dense: + if n is not None: + raise ValueError("Dense mode forbids n") + if b_dtype is not None: + raise ValueError("Dense mode forbids b_dtype") + if b_tensor.dtype != torch.bfloat16: + raise ValueError(f"b_tensor must have dtype torch.bfloat16, got {b_tensor.dtype}") + if b_tensor.ndim != 3: + raise ValueError(f"b_tensor must have shape (n, k, experts), got {tuple(b_tensor.shape)}") + n, b_k, experts = b_tensor.shape + if b_k != k: + raise ValueError(f"b_tensor K dimension ({b_k}) must match a_tensor ({k})") + else: + _require_pointer_tensor(b_ptrs, "b_ptrs") + if b_ptrs.device != a_tensor.device: + raise ValueError(f"b_ptrs must be on the same device as a_tensor " f"({a_tensor.device}), got {b_ptrs.device}") + if b_ptrs.data_ptr() % 8 != 0: + raise ValueError("b_ptrs data pointer must be 8-byte aligned") + if padded_offsets.ndim == 1 and b_ptrs.numel() != padded_offsets.numel(): + raise ValueError(f"b_ptrs length mismatch: expected {padded_offsets.numel()}, " f"got {b_ptrs.numel()}") + if n is None or b_dtype is None: + raise ValueError("Discrete mode requires n and b_dtype") + if b_dtype != torch.bfloat16: + raise ValueError(f"b_dtype must be torch.bfloat16 for the BF16 backend, got {b_dtype}") + if n <= 0: + raise ValueError(f"n must be > 0, got {n}") + experts = b_ptrs.numel() + + return is_dense, tensor_m, n, experts + + +def grouped_gemm_wrapper_sm100( + a_tensor: torch.Tensor, + padded_offsets: torch.Tensor, + alpha_tensor: torch.Tensor, + b_tensor: Optional[torch.Tensor] = None, + bias_tensor: Optional[torch.Tensor] = None, + b_ptrs: Optional[torch.Tensor] = None, + n: Optional[int] = None, + b_dtype: Optional[torch.dtype] = None, + b_major: str = "k", + prob_tensor: Optional[torch.Tensor] = None, + acc_dtype: torch.dtype = torch.float32, + c_dtype: torch.dtype = torch.bfloat16, + d_dtype: torch.dtype = torch.bfloat16, + c_tensor: Optional[torch.Tensor] = None, + d_tensor: Optional[torch.Tensor] = None, + cd_major: str = "n", + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + vector_f32: bool = False, + m_aligned: int = 256, + generate_c: bool = False, + use_dynamic_sched: bool = False, + current_stream: Optional[cuda.CUstream] = None, +) -> TupleDict: + is_dense, tensor_m, n_out, expert_cnt = _normalize_call( + a_tensor, + padded_offsets, + b_tensor, + b_ptrs, + n, + b_dtype, + prob_tensor, + c_dtype, + d_dtype, + cd_major, + m_aligned, + ) + expected_shape = (tensor_m, n_out, 1) + expected_stride = (n_out, 1, tensor_m * n_out) + if c_tensor is None: + internal_c = torch.empty_strided( + expected_shape, + expected_stride, + dtype=c_dtype, + device=a_tensor.device, + ) + else: + _validate_output( + c_tensor, + name="c_tensor", + shape=expected_shape, + stride=expected_stride, + dtype=c_dtype, + device=a_tensor.device, + ) + internal_c = c_tensor + if d_tensor is None: + d_tensor = torch.empty_strided( + expected_shape, + expected_stride, + dtype=d_dtype, + device=a_tensor.device, + ) + else: + _validate_output( + d_tensor, + name="d_tensor", + shape=expected_shape, + stride=expected_stride, + dtype=d_dtype, + device=a_tensor.device, + ) + + overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + workspace_bytes = (128 * expert_cnt if not is_dense else 0) + (4 if use_dynamic_sched else 0) + workspace_signature = ( + (max(workspace_bytes, 1),), + (a_tensor.device.type, a_tensor.device.index), + ) + cache_key = ( + "bf16", + "dense" if is_dense else "discrete", + *_tensor_signature(a_tensor, dynamic_m=True), + *_tensor_signature(b_tensor), + *_tensor_signature(b_ptrs), + *_tensor_signature(internal_c, dynamic_m=True), + *_tensor_signature(d_tensor, dynamic_m=True), + *_tensor_signature(padded_offsets), + *_tensor_signature(alpha_tensor), + *_tensor_signature(bias_tensor), + *_tensor_signature(prob_tensor, dynamic_m=True), + n_out, + expert_cnt, + b_dtype, + b_major, + acc_dtype, + c_dtype, + d_dtype, + cd_major, + tuple(mma_tiler_mn), + tuple(cluster_shape_mn) if cluster_shape_mn is not None else None, + vector_f32, + m_aligned, + generate_c, + use_dynamic_sched, + overlap_margin, + workspace_signature, + ) + + op = _cache_of_GroupedGemmSm100Objects.get(cache_key) + if op is None: + op = GroupedGemmSm100( + sample_a=a_tensor, + sample_c=internal_c, + sample_d=d_tensor, + sample_padded_offsets=padded_offsets, + sample_alpha=alpha_tensor, + sample_b=b_tensor if is_dense else None, + sample_bias=bias_tensor, + sample_prob=prob_tensor, + num_experts=None if is_dense else expert_cnt, + b_shape=None if is_dense else (n_out, a_tensor.shape[1]), + b_dtype=None if is_dense else b_dtype, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + vector_f32=vector_f32, + m_aligned=m_aligned, + generate_c=generate_c, + b_major=b_major, + use_dynamic_sched=use_dynamic_sched, + ) + assert op.check_support(), "Unsupported configuration" + op.compile() + _cache_of_GroupedGemmSm100Objects[cache_key] = op + + op.execute( + a_tensor=a_tensor, + c_tensor=internal_c, + d_tensor=d_tensor, + padded_offsets=padded_offsets, + alpha_tensor=alpha_tensor, + b_tensor=b_tensor if is_dense else None, + b_ptrs=None if is_dense else b_ptrs, + bias_tensor=bias_tensor, + prob_tensor=prob_tensor, + current_stream=current_stream, + ) + return TupleDict( + d_tensor=d_tensor, + c_tensor=internal_c if generate_c else None, + ) + + +__all__ = [ + "GroupedGemmSm100", + "grouped_gemm_wrapper_sm100", +] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_unfused/moe_grouped_gemm.py b/python/cudnn/grouped_gemm/grouped_gemm_unfused/moe_grouped_gemm.py new file mode 100644 index 000000000..988e3a0dc --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_unfused/moe_grouped_gemm.py @@ -0,0 +1,1354 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +""" +MoE BF16 Grouped GEMM Kernel. + +Supports: + - Static / Dynamic persistent tile scheduling (MoEPersistentTileScheduler) + - Dense (contiguous 3-D B) / Discrete (per-expert pointer array B) weight layout + - Optional bias and routing-probability (prob) fusion + - Optional C output (generate_c) + +This module contains only the kernel class. +MoE scheduler components live in moe_persistent_scheduler.py / moe_sched_extension.py / moe_utils.py. +""" + +from typing import Type, Tuple, Union, Optional + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import OperandMajorMode, cpasync, tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.typing import Int32, AddressSpace +from ..moe_persistent_scheduler import ( + MoEPersistentTileScheduler, + MoESchedulerParams, + MoEWorkTileInfo, +) +from ..moe_utils import ( + MoEWeightMode, + TensormapWorkspace, + store_tma_desc, +) +from ..moe_sched_extension import ( + DiscreteWeightGroupedGemmSchedExtension, + ContiguousGroupedGemmSchedExtension, +) +from ..moe_kernel_helpers import ( + can_implement_bf16_grouped_gemm, + compute_grid, + epilog_gmem_copy_and_partition, +) + + +class MoEGroupedGemmBf16Kernel: + """Plain BF16 grouped GEMM kernel with MoE tile scheduling. + + Supports both dense and discrete weight layouts, static and dynamic + scheduling, optional C output, and optional bias fusion. A/B use BF16 + storage, MMA accumulates in FP32, and C/D may be BF16, FP16, or FP32. + + :param acc_dtype: Accumulator data type (Float32). + :param use_2cta_instrs: Use 2-CTA MMA instructions. + :param mma_tiler_mn: MMA tile shape (M, N). + :param cluster_shape_mn: Cluster shape (M, N). + :param vectorized_f32: Use packed FP32 arithmetic in epilogue. + :param generate_c: Generate C output tensor. + :param enable_bias: Fuse bias addition. + :param expert_cnt: Number of experts. + :param weight_mode: ``MoEWeightMode.DENSE`` or ``MoEWeightMode.DISCRETE``. + :param use_dynamic_sched: Enable dynamic tile scheduling. + """ + + FIX_PAD_SIZE = 256 + + @staticmethod + def can_implement( + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m: int, + n: int, + k: int, + l: int, + a_major: str, + b_major: str, + cd_major: str, + m_aligned: int, + ) -> bool: + return can_implement_bf16_grouped_gemm( + ab_dtype, + c_dtype, + d_dtype, + acc_dtype, + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + m, + n, + k, + l, + a_major, + b_major, + cd_major, + m_aligned, + fix_pad_size=MoEGroupedGemmBf16Kernel.FIX_PAD_SIZE, + n_align=32, + tile_n_align=32, + ) + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + vectorized_f32: bool, + generate_c: bool, + enable_bias: bool, + expert_cnt: int, + weight_mode: MoEWeightMode = MoEWeightMode.DENSE, + use_dynamic_sched: bool = False, + ): + mma_tile_m = mma_tiler_mn[0] + if self.FIX_PAD_SIZE % mma_tile_m != 0: + raise ValueError(f"FIX_PAD_SIZE ({self.FIX_PAD_SIZE}) must be divisible by " f"mma_tiler_mn[0] ({mma_tile_m}).") + if expert_cnt > 1024: + raise ValueError("Expert count > 1024 is not supported.") + if not isinstance(weight_mode, MoEWeightMode): + raise TypeError(f"weight_mode must be a MoEWeightMode, got {type(weight_mode)}") + + self.expert_cnt = expert_cnt + self.acc_dtype = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + self.occupancy = 1 + + self.epilog_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.sched_warp_id = 6 + self.bias_load_warp_id = 7 if enable_bias else None + self.threads_per_warp = 32 + all_warps = [ + *self.epilog_warp_id, + self.mma_warp_id, + self.tma_warp_id, + self.sched_warp_id, + ] + warps_wo_sched = [*self.epilog_warp_id, self.mma_warp_id, self.tma_warp_id] + if enable_bias: + all_warps.append(self.bias_load_warp_id) + warps_wo_sched.append(self.bias_load_warp_id) + self.threads_per_cta = self.threads_per_warp * len(all_warps) + self.threads_wo_sched = self.threads_per_warp * len(warps_wo_sched) + + self.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=32 * len(self.epilog_warp_id), + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=32 * len((self.mma_warp_id, *self.epilog_warp_id)), + ) + self.sched_sync_barrier = pipeline.NamedBarrier( + barrier_id=4, + num_threads=self.threads_per_warp, + ) + + self.num_smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.vectorized_f32 = vectorized_f32 + self.generate_c = generate_c + self.enable_bias = enable_bias + self.weight_mode = weight_mode + self.use_dynamic_sched = use_dynamic_sched + self.num_epilog_warps = len(self.epilog_warp_id) + + # ------------------------------------------------------------------ + # _setup_attributes + # ------------------------------------------------------------------ + + def _setup_attributes(self): + """Configure MMA / tile / stage / SMEM layouts from GEMM inputs.""" + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + self.mma_tiler = ( + self.mma_tiler[0], + self.mma_tiler[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.epi_tile = (128, 32) + + ( + self.num_acc_stage, + self.num_ab_stage, + self.num_c_stage, + self.num_d_stage, + self.num_tile_stage, + self.num_bias_stage, + ) = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.epi_tile, + self.c_dtype, + self.c_layout, + self.d_dtype, + self.d_layout, + self.num_smem_capacity, + self.occupancy, + self.generate_c, + self.bias_dtype if self.enable_bias else None, + ) + + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + self.d_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.d_dtype, + self.d_layout, + self.epi_tile, + self.num_d_stage, + ) + if self.enable_bias: + self.bias_smem_layout_staged = cute.make_layout( + (self.mma_tiler[1], self.num_bias_stage), + stride=(1, self.mma_tiler[1]), + ) + else: + self.bias_smem_layout_staged = cute.make_layout((1, 1)) + + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + self.num_tmem_alloc_cols = utils.get_num_tmem_alloc_cols(tCtAcc_fake) + self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * self.num_acc_stage + + # ------------------------------------------------------------------ + # Stage computation + # ------------------------------------------------------------------ + + @staticmethod + def _compute_stages( + tiled_mma, + mma_tiler_mnk, + a_dtype, + b_dtype, + epi_tile, + c_dtype, + c_layout, + d_dtype, + d_layout, + num_smem_capacity, + occupancy, + generate_c, + bias_dtype, + ): + num_acc_stage = 2 + num_c_stage = 1 + num_d_stage = 1 + num_tile_stage = 2 + + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a(tiled_mma, mma_tiler_mnk, a_dtype, 1) + b_smem_layout_stage_one = sm100_utils.make_smem_layout_b(tiled_mma, mma_tiler_mnk, b_dtype, 1) + c_smem_layout_stage_one = sm100_utils.make_smem_layout_epi(c_dtype, c_layout, epi_tile, 1) + d_smem_layout_stage_one = sm100_utils.make_smem_layout_epi(d_dtype, d_layout, epi_tile, 1) + + ab_bytes_per_stage = cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + cute.size_in_bytes(b_dtype, b_smem_layout_stage_one) + mbar_helpers_bytes = 1024 + sinfo_bytes = 4 * 4 * num_tile_stage + c_bytes = cute.size_in_bytes(c_dtype, c_smem_layout_stage_one) * num_c_stage + d_bytes = cute.size_in_bytes(d_dtype, d_smem_layout_stage_one) * num_d_stage + + if bias_dtype is not None: + num_bias_stage = 2 + bias_bytes = mma_tiler_mnk[1] * num_bias_stage * (bias_dtype.width // 8) + else: + num_bias_stage = 0 + bias_bytes = 0 + + epi_bytes = c_bytes + d_bytes + bias_bytes + num_ab_stage = (num_smem_capacity // occupancy - (mbar_helpers_bytes + epi_bytes + sinfo_bytes)) // ab_bytes_per_stage + return num_acc_stage, num_ab_stage, num_c_stage, num_d_stage, num_tile_stage, num_bias_stage + + # ------------------------------------------------------------------ + # Workspace helpers + # ------------------------------------------------------------------ + + def get_desc_workspace_bytes(self) -> int: + if self.weight_mode == MoEWeightMode.DISCRETE: + return TensormapWorkspace.size_bytes(1, self.expert_cnt) + return 0 + + def get_workspace_bytes(self) -> int: + desc_workspace_bytes = self.get_desc_workspace_bytes() + dynamic_sched_bytes = 4 if self.use_dynamic_sched else 0 + return desc_workspace_bytes + dynamic_sched_bytes + + @cute.jit + def _get_sched_counter_ptr(self, workspace_ptr): + counter_addr = workspace_ptr.toint() + self.get_desc_workspace_bytes() + return cute.make_ptr( + cutlass.Int32, + counter_addr, + AddressSpace.gmem, + assumed_align=4, + ) + + # ------------------------------------------------------------------ + # helper_kernel: pre-main-kernel initialization + # - discrete weight: build per-expert B TMA descriptors + # - dynamic sched: reset the atomic tile counter + # ------------------------------------------------------------------ + + @cute.kernel + def helper_kernel( + self, + ptrs_b: cute.Pointer, + n: Int32, + k: Int32, + b_stride_size: cutlass.Int64, + b_major_mode: cutlass.Constexpr, + workspace_ptr, + tiled_mma_arg: cute.TiledMma, + b_smem_layout_arg, + cluster_layout_vmnk_shape_arg: cutlass.Constexpr, + ): + """Pre-main-kernel initialization. + + Launched with grid=(expert_cnt, 1, 1) for discrete mode, or + grid=(1, 1, 1) for dense+dynamic mode. + + Discrete weight: each block builds a B TMA descriptor for one expert. + Dynamic sched: block 0 resets the atomic tile counter to 0. + """ + expert_idx = cute.arch.block_idx()[0] + + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + b_tma_op_arg = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma_arg.thr_id) + b_ptr_tensor = cute.make_tensor( + cute.make_ptr(cutlass.Int64, ptrs_b.toint(), AddressSpace.gmem, assumed_align=8), cute.make_layout((self.expert_cnt,)) + ) + + c1 = cutlass.Int32(1) + c0 = cutlass.Int64(0) + c1_64 = 1 + if cutlass.const_expr(b_major_mode == OperandMajorMode.K): + stride_n = b_stride_size + stride_k = c1_64 + else: + stride_n = c1_64 + stride_k = b_stride_size + + b_ptr_val = b_ptr_tensor[expert_idx] + b_ptr = cute.make_ptr(self.b_dtype, b_ptr_val, AddressSpace.gmem) + b_tensor_i = cute.make_tensor( + b_ptr, + cute.make_layout((n, k, c1), stride=(stride_n, stride_k, c0)), + ) + tma_atom_b, _ = cute.nvgpu.make_tiled_tma_atom_B( + b_tma_op_arg, + b_tensor_i, + b_smem_layout_arg, + self.mma_tiler, + tiled_mma_arg, + cluster_layout_vmnk_shape_arg, + ) + workspace = TensormapWorkspace(workspace_ptr, ["b"]) + store_tma_desc(tma_atom_b, workspace.get_ptr("b", expert_idx)) + + if cutlass.const_expr(self.use_dynamic_sched): + if expert_idx == cutlass.Int32(0): + sched_counter = cute.make_tensor( + self._get_sched_counter_ptr(workspace_ptr), + cute.make_layout(1), + ) + sched_counter[0] = cutlass.Int32(0) + + # ------------------------------------------------------------------ + # __call__ + # ------------------------------------------------------------------ + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b, # Dense: cute.Tensor (N,K,L) | Discrete: cute.Pointer to int64[] + n: Int32, # Ignored for dense mode + k: Int32, # Ignored for dense mode + b_stride_size: cutlass.Int64, # Ignored for dense mode + b_major_mode: cutlass.Constexpr, # Ignored for dense mode + workspace_ptr, + c: cute.Tensor, + d: cute.Tensor, + padded_offsets: cute.Tensor, + alpha: cute.Tensor, + bias: Optional[cute.Tensor], + prob: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + """Execute the GEMM. + + Dense mode: ``b`` is a 3-D cute.Tensor (N, K, L). + Discrete mode: ``b`` is a cute.Pointer to a device int64[] array of + per-expert base addresses; ``n``, ``k``, ``b_stride_size``, and + ``b_major_mode`` describe the uniform per-expert layout. + """ + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = a.element_type + self.c_dtype: Type[cutlass.Numeric] = c.element_type + self.d_dtype: Type[cutlass.Numeric] = d.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + self.d_layout = utils.LayoutEnum.from_tensor(d) + self.bias_dtype = bias.element_type if cutlass.const_expr(self.enable_bias) else cutlass.BFloat16 + + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + else: + self.b_major_mode = b_major_mode + + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"A/B dtype must match: {self.a_dtype} != {self.b_dtype}") + + self._setup_attributes() + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # ---- B setup (mode-dependent) ---- + # Discrete mode receives a pointer array, then builds a template B tensor. + # helper_kernel still needs the original pointer-array argument. + b_from_call_arg = b + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + c1 = cutlass.Int32(1) + c0 = cutlass.Int64(0) + c1_64 = 1 + if cutlass.const_expr(b_major_mode == OperandMajorMode.K): + b_template_stride = (b_stride_size, c1_64, c0) + else: + b_template_stride = (c1_64, b_stride_size, c0) + b_template_layout = cute.make_layout((n, k, c1), stride=b_template_stride) + b_ptr_typed = cute.make_ptr(self.b_dtype, b.toint(), AddressSpace.gmem, assumed_align=16) + b = cute.make_tensor(b_ptr_typed, b_template_layout) + + # ---- TMA atoms ---- + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + c_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c, + c_smem_layout, + self.epi_tile, + ) + d_smem_layout = cute.slice_(self.d_smem_layout_staged, (None, None, 0)) + tma_atom_d, tma_tensor_d = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + d, + d_smem_layout, + self.epi_tile, + ) + + # ---- Helper kernel: TMA desc init (discrete) + sched counter reset (dynamic) ---- + _need_helper = cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE or self.use_dynamic_sched) + if cutlass.const_expr(_need_helper): + _helper_grid_x = self.expert_cnt if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else 1 + _helper_args = ( + b_from_call_arg if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cute.make_ptr(cutlass.Int64, 0, AddressSpace.gmem), + n if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int32(0), + k if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int32(0), + b_stride_size if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else cutlass.Int64(0), + b_major_mode if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else self.b_major_mode, + workspace_ptr, + tiled_mma, + b_smem_layout, + self.cluster_layout_vmnk.shape, + ) + self.helper_kernel(*_helper_args).launch( + grid=(_helper_grid_x, 1, 1), + block=(1, 1, 1), + stream=stream, + min_blocks_per_mp=1, + ) + + # ---- Grid computation via MoE scheduler ---- + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + b_n, b_k, b_l = cute.shape(b) + sched_expert_shape = (self.expert_cnt, b_n, b_k) + else: + sched_expert_shape = (self.expert_cnt, n, k) + + sched_params = MoESchedulerParams( + scenario="2Dx3D", + expert_shape=sched_expert_shape, + cta_tile_shape_mnk=self.cta_tile_shape_mnk, + cluster_shape_mn=self.cluster_shape_mn, + use_dynamic_sched=self.use_dynamic_sched, + ) + self.sched_params, grid = compute_grid(sched_params, max_active_clusters, self.use_2cta_instrs) + self.buffer_align_bytes = 1024 + + SchedulerStorage = MoEPersistentTileScheduler.make_storage_struct(self.num_tile_stage, self.use_dynamic_sched) + + # ---- Shared storage ---- + @cute.struct + class SharedStorage: + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + scheduler: SchedulerStorage + if cutlass.const_expr(self.enable_bias): + bias_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_bias_stage * 2] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + sC: cute.struct.Align[ + cute.struct.MemRange[self.c_dtype, cute.cosize(self.c_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sD: cute.struct.Align[ + cute.struct.MemRange[self.d_dtype, cute.cosize(self.d_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sA: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sB: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + if cutlass.const_expr(self.enable_bias): + sBias: cute.struct.Align[ + cute.struct.MemRange[self.bias_dtype, cute.cosize(self.bias_smem_layout_staged)], + 16, + ] + + self.shared_storage = SharedStorage + + # ---- Launch ---- + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + tma_atom_d, + tma_tensor_d, + padded_offsets, + alpha, + bias, + prob, + workspace_ptr, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.d_smem_layout_staged, + self.bias_smem_layout_staged, + self.epi_tile, + self.sched_params, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + max_number_threads=[self.threads_per_cta, 1, 1], + smem=self.shared_storage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + return + + # ------------------------------------------------------------------ + # Helper methods + # ------------------------------------------------------------------ + + @cute.jit + def store_c( + self, + tiled_copy_r2s, + tma_atom_c, + warp_idx, + tTR_rAcc, + tRS_rC, + tRS_sC, + bSG_gC, + bSG_sC, + c_pipeline, + prev_subtile_idx, + real_subtile_idx, + ): + c_buffer = prev_subtile_idx % self.num_c_stage + tRS_rC.store(tTR_rAcc.load().to(self.c_dtype)) + cute.copy(tiled_copy_r2s, tRS_rC[(None, None, 0)], tRS_sC[(None, None, 0, c_buffer)]) + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier.arrive_and_wait() + if warp_idx == self.epilog_warp_id[0]: + cute.copy(tma_atom_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, real_subtile_idx)]) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + + def epilog_tmem_copy_and_partition(self, tidx, tAcc, gD_mnl, epi_tile, use_2cta_instrs): + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.d_layout, + self.d_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + gD_mnl_epi = cute.flat_divide(gD_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + tTR_gC = thr_copy_t2r.partition_D(gD_mnl_epi) + tTR_rAcc = cute.make_rmem_tensor(tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc + + def epilog_smem_copy_and_partition(self, tiled_copy_t2r, tTR_rD, tidx, sD): + copy_atom_r2s = sm100_utils.get_smem_store_op(self.d_layout, self.d_dtype, self.acc_dtype, tiled_copy_t2r) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sD = thr_copy_r2s.partition_D(sD) + tRS_rD = tiled_copy_r2s.retile(tTR_rD) + return tiled_copy_r2s, tRS_rD, tRS_sD + + # ------------------------------------------------------------------ + # GPU device kernel + # ------------------------------------------------------------------ + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + tma_atom_d: cute.CopyAtom, + mD_mnl: cute.Tensor, + padded_offsets: cute.Tensor, + alpha: cute.Tensor, + mBias_nl: Optional[cute.Tensor], + prob: cute.Tensor, + workspace_ptr, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + d_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + bias_smem_layout_staged: Optional[cute.Layout], + epi_tile: cute.Tile, + sched_params: MoESchedulerParams, + epilogue_op: cutlass.Constexpr, + ): + del epilogue_op + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_d) + if cutlass.const_expr(self.generate_c): + cpasync.prefetch_descriptor(tma_atom_c) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + total_token = padded_offsets[self.expert_cnt - 1] + + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + tidx, _, _ = cute.arch.thread_idx() + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + sched_storage = storage.scheduler + + # ---- Pipeline setup ---- + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_warp_id) * (2 if use_2cta_instrs else 1) + acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + tile_info_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.threads_per_warp) + tile_info_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.threads_wo_sched) + tile_info_pipeline = pipeline.PipelineAsync.create( + barrier_storage=sched_storage.tile_info_mbar.data_ptr(), + num_stages=self.num_tile_stage, + producer_group=tile_info_pipeline_producer_group, + consumer_group=tile_info_pipeline_consumer_group, + ) + + if cutlass.const_expr(self.enable_bias): + bias_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.threads_per_warp) + bias_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.epilog_warp_id), + ) + bias_pipeline = pipeline.PipelineCpAsync.create( + barrier_storage=storage.bias_mbar_ptr.data_ptr(), + num_stages=self.num_bias_stage, + producer_group=bias_pipeline_producer_group, + consumer_group=bias_pipeline_consumer_group, + ) + sBias = storage.sBias.get_tensor(bias_smem_layout_staged) + + # ---- Scheduler and TMEM allocator ---- + scheduler = MoEPersistentTileScheduler.create( + sched_params, + padded_offsets, + cute.arch.block_idx(), + cute.arch.grid_dim(), + counter_ptr=self._get_sched_counter_ptr(workspace_ptr), + sched_storage=sched_storage, + ) + scheduler.internal_init() + + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, + ) + + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_arrive_relaxed() + + sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner) + sD = storage.sD.get_tensor(d_smem_layout_staged.outer, swizzle=d_smem_layout_staged.inner) + sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner) + sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner) + info_layout = cute.make_layout((4, self.num_tile_stage), stride=(1, 4)) + sInfo = sched_storage.sInfo.get_tensor(info_layout) + + a_full_mcast_mask = None + b_full_mcast_mask = None + # Multicast masks must be created together when any mcast or 2CTA is active. + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2) + b_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1) + + # SMEM fragments for MMA and the TMEM accumulator shape shared by MMA/epilogue. + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + if cute.size(self.cluster_shape_mn) > 1: + cute.arch.cluster_wait() + else: + self.cta_sync_barrier.arrive_and_wait() + + if total_token <= 0: + cute.arch.nvvm.exit() + + # ============================================================== + # Scheduler warp (MoE Persistent Tile Scheduler) + # ============================================================== + if warp_idx == self.sched_warp_id: + work_tile_info = scheduler.initial_work_tile_info() + tile_info_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_tile_stage) + while work_tile_info.is_valid_tile: + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = work_tile_info.expert_idx + sInfo[(1, tile_info_producer_state.index)] = work_tile_info.tile_m_idx + sInfo[(2, tile_info_producer_state.index)] = work_tile_info.tile_n_idx + sInfo[(3, tile_info_producer_state.index)] = work_tile_info.k_tile_cnt + cute.arch.fence_proxy("async.shared", space="cta") + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + work_tile_info = scheduler.advance_to_next_work() + + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = cutlass.Int32(-1) + sInfo[(1, tile_info_producer_state.index)] = cutlass.Int32(0) + sInfo[(2, tile_info_producer_state.index)] = cutlass.Int32(0) + sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + tile_info_pipeline.producer_tail(tile_info_producer_state) + + # ============================================================== + # Bias load warp + # ============================================================== + if cutlass.const_expr(self.enable_bias): + if warp_idx == self.bias_load_warp_id: + bias_ext = self._make_extension(workspace_ptr) + bias_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_bias_stage) + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + bias_g2s_atom = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(cache_mode=cute.nvgpu.cpasync.LoadCacheMode.GLOBAL), + self.bias_dtype, + num_bits_per_copy=128, + ) + bias_g2s_tiled = cute.make_tiled_copy_tv( + bias_g2s_atom, + cute.make_layout((32,)), + cute.make_layout((8,)), + ) + thr_bias_g2s = bias_g2s_tiled.get_slice(cute.arch.lane_idx()) + tBs_sBias = thr_bias_g2s.partition_D(sBias) + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + bias_producer_state.reset_count() + work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + bias_ext.update_expert_info(padded_offsets, work_tile_info.expert_idx) + real_bias, _ = bias_ext.get_gmem_tensor("bias", mBias_nl, padded_offsets, work_tile_info) + gBias_expert = cute.local_tile(real_bias, cute.slice_(self.mma_tiler[:2], (0, None)), (None, None)) + bias_tile = gBias_expert[(None, work_tile_info.tile_n_idx, 0)] + bias_identity_tensor = cute.make_identity_tensor(bias_tile.shape) + bias_partitioned_by_g2s = thr_bias_g2s.partition_S(bias_tile) + bias_coord_partitioned_by_g2s = thr_bias_g2s.partition_S(bias_identity_tensor) + residue_n = sched_params.intermediate - work_tile_info.tile_n_idx * self.cta_tile_shape_mnk[1] + bias_pred_tensor = cute.make_rmem_tensor(bias_coord_partitioned_by_g2s[(None, 0)].shape, cutlass.Boolean) + for vi in cutlass.range_constexpr(cute.size(bias_pred_tensor)): + bias_pred_tensor[vi] = cute.elem_less(bias_coord_partitioned_by_g2s[(vi, 0)], (residue_n,)) + bias_pred_tensor = bias_pred_tensor[((0, None),)] + + bias_pipeline.producer_acquire(bias_producer_state) + cute.copy( + bias_g2s_tiled, + bias_partitioned_by_g2s[(None, 0)], + tBs_sBias[(None, 0, bias_producer_state.index)], + pred=bias_pred_tensor, + ) + bias_pipeline.producer_commit(bias_producer_state) + bias_producer_state.advance() + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + bias_pipeline.producer_tail(bias_producer_state) + + # ============================================================== + # DMA / TMA load warp + # ============================================================== + if warp_idx == self.tma_warp_id: + ext = self._make_extension(workspace_ptr) + ab_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_ab_stage) + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + k_tile_cnt = work_tile_info.k_tile_cnt + ext.update_expert_info(padded_offsets, work_tile_info.expert_idx) + + real_a, _ = ext.get_gmem_tensor("a", mA_mkl, padded_offsets, work_tile_info) + real_b, desc_ptr_b = ext.get_gmem_tensor("b", mB_nkl, padded_offsets, work_tile_info) + gA_mkl = cute.local_tile(real_a, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None)) + gB_nkl = cute.local_tile(real_b, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None)) + + thr_mma_dma = tiled_mma.get_slice(mma_tile_coord_v) + tCgA = thr_mma_dma.partition_A(gA_mkl) + tCgB = thr_mma_dma.partition_B(gB_nkl) + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + mma_tile_coord_m = work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape) + mma_tile_coord_n = work_tile_info.tile_n_idx + tAgA_slice = tAgA[(None, mma_tile_coord_m, None, 0)] + tBgB_slice = tBgB[(None, mma_tile_coord_n, None, 0)] + + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + tAgA_k = tAgA_slice[(None, ab_producer_state.count)] + tBgB_k = tBgB_slice[(None, ab_producer_state.count)] + tAsA_pipe = tAsA[(None, ab_producer_state.index)] + tBsB_pipe = tBsB[(None, ab_producer_state.index)] + tma_bar = ab_pipeline.producer_get_barrier(ab_producer_state) + ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) + cute.copy(tma_atom_a, tAgA_k, tAsA_pipe, tma_bar_ptr=tma_bar, mcast_mask=a_full_mcast_mask) + cute.copy(tma_atom_b, tBgB_k, tBsB_pipe, tma_bar_ptr=tma_bar, mcast_mask=b_full_mcast_mask, tma_desc_ptr=desc_ptr_b) + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + ab_pipeline.producer_tail(ab_producer_state) + + # ============================================================== + # MMA warp + # ============================================================== + if warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + ab_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_ab_stage) + acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_stage) + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + k_tile_cnt = tile_info[3] + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + acc_producer_state.reset_count() + peek_acc_empty_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_acc_empty_status = acc_pipeline.producer_try_acquire(acc_producer_state) + + acc_stage_index = acc_producer_state.index + tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] + + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state, peek_acc_empty_status) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + if is_leader_cta: + ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + num_kblocks = cute.size(tCrA, mode=[2]) + ab_consumer_state_next = ab_consumer_state.clone() + ab_consumer_state_next.advance() + if ab_consumer_state_next.count < k_tile_cnt: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state_next) + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + kblock_coord = (None, None, kblock_idx, ab_consumer_state.index) + cute.gemm(tiled_mma, tCtAcc, tCrA[kblock_coord], tCrB[kblock_coord], tCtAcc) + tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + ab_pipeline.consumer_release(ab_consumer_state) + ab_consumer_state = ab_consumer_state_next + + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + + acc_producer_state.advance() + if acc_producer_state.count < k_tile_cnt: + if is_leader_cta: + peek_acc_empty_status = acc_pipeline.producer_try_acquire(acc_producer_state) + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + acc_pipeline.producer_tail(acc_producer_state) + + # ============================================================== + # Epilogue warps + # ============================================================== + if warp_idx < self.mma_warp_id and total_token > 0: + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + epi_tidx = tidx + thr_mma_epi = tiled_mma.get_slice(mma_tile_coord_v) + gD_mnl_shape = cute.local_tile(mD_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) + tCgD_shape = thr_mma_epi.partition_C(gD_mnl_shape) + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = self.epilog_tmem_copy_and_partition( + epi_tidx, + tCtAcc_base, + tCgD_shape, + epi_tile, + use_2cta_instrs, + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = self.epilog_smem_copy_and_partition( + tiled_copy_t2r, + tTR_rC, + epi_tidx, + sC, + ) + tTR_rD = cute.make_rmem_tensor(tTR_rAcc.shape, self.d_dtype) + tiled_copy_r2s, tRS_rD, tRS_sD = self.epilog_smem_copy_and_partition( + tiled_copy_t2r, + tTR_rD, + epi_tidx, + sD, + ) + + epi_ext = self._make_extension(workspace_ptr) + acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_stage) + c_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilog_warp_id)) + c_pipeline = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_producer_group) + d_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilog_warp_id)) + d_pipeline = pipeline.PipelineTmaStore.create(num_stages=self.num_d_stage, producer_group=d_producer_group) + + tile_info_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_tile_stage) + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + if cutlass.const_expr(self.enable_bias): + bias_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_bias_stage) + bias_s2r_tom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.bias_dtype, num_bits_per_copy=128) + tTR_rBias = cute.make_rmem_tensor(cute.make_layout(self.epi_tile[1]), self.bias_dtype) + + num_prev_subtiles = cutlass.Int32(0) + while is_valid_tile: + epi_work_tile_info = MoEWorkTileInfo( + expert_idx=tile_info[0], + tile_m_idx=tile_info[1], + tile_n_idx=tile_info[2], + k_tile_cnt=tile_info[3], + ) + expert_idx = epi_work_tile_info.expert_idx + epi_ext.update_expert_info(padded_offsets, expert_idx) + alpha_val = alpha[expert_idx] + + if cutlass.const_expr(self.enable_bias): + bias_consumer_state.reset_count() + bias_pipeline.consumer_wait(bias_consumer_state) + sBias_stage = sBias[(None, bias_consumer_state.index)] + sBias_subtiles = cute.flat_divide(sBias_stage, cute.make_layout(self.epi_tile[1])) + + real_d, _ = epi_ext.get_gmem_tensor("d", mD_mnl, padded_offsets, epi_work_tile_info) + real_c, _ = epi_ext.get_gmem_tensor("c", mC_mnl, padded_offsets, epi_work_tile_info) + thr_mma_epi_loop = tiled_mma.get_slice(mma_tile_coord_v) + + gD_mnl_loop = cute.local_tile(real_d, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) + tCgD_loop = thr_mma_epi_loop.partition_C(gD_mnl_loop) + _, bSG_sD, bSG_gD_partitioned = epilog_gmem_copy_and_partition( + epi_tidx, + tma_atom_d, + tCgD_loop, + epi_tile, + sD, + ) + + gC_mnl_loop = cute.local_tile(real_c, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)) + tCgC_loop = thr_mma_epi_loop.partition_C(gC_mnl_loop) + _, bSG_sC, bSG_gC_partitioned = epilog_gmem_copy_and_partition( + epi_tidx, + tma_atom_c, + tCgC_loop, + epi_tile, + sC, + ) + + epi_mma_tile_coord = ( + epi_work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape), + epi_work_tile_info.tile_n_idx, + 0, + ) + bSG_gC = bSG_gC_partitioned[(None, None, None, *epi_mma_tile_coord)] + bSG_gD = bSG_gD_partitioned[(None, None, None, *epi_mma_tile_coord)] + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + bSG_gD = cute.group_modes(bSG_gD, 1, cute.rank(bSG_gD)) + + mPosition = epi_work_tile_info.tile_m_idx * self.cta_tile_shape_mnk[0] + tidx + real_prob, _ = epi_ext.get_gmem_tensor("prob", prob, padded_offsets, epi_work_tile_info) + mProb = real_prob[mPosition, 0, 0] + + acc_stage_index = acc_consumer_state.index + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)] + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + acc_pipeline.consumer_wait(acc_consumer_state) + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + + for subtile_idx in cutlass.range(0, subtile_cnt, 1, unroll=1): + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + if cutlass.const_expr(self.enable_bias): + sBias_sub = sBias_subtiles[(None, subtile_idx)] + cute.copy(bias_s2r_tom, sBias_sub, tTR_rBias) + bias_vec = tTR_rBias.load() + if cutlass.const_expr(self.vectorized_f32): + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc), 2): + bias_f32_0 = bias_vec[i].to(cutlass.Float32) + bias_f32_1 = bias_vec[i + 1].to(cutlass.Float32) + bias_f32_0, bias_f32_1 = cute.arch.mul_packed_f32x2( + (mProb, mProb), + (bias_f32_0, bias_f32_1), + rnd="rn", + ftz=False, + ) + tTR_rAcc[i], tTR_rAcc[i + 1] = cute.arch.fma_packed_f32x2( + (tTR_rAcc[i], tTR_rAcc[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + (bias_f32_0, bias_f32_1), + rnd="rn", + ftz=False, + ) + else: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + tTR_rAcc[i] = tTR_rAcc[i] * cutlass.Float32(alpha_val) + bias_vec[i].to(cutlass.Float32) * mProb + else: + if cutlass.const_expr(self.vectorized_f32): + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc), 2): + tTR_rAcc[i], tTR_rAcc[i + 1] = cute.arch.mul_packed_f32x2( + (tTR_rAcc[i], tTR_rAcc[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + rnd="rn", + ftz=False, + ) + else: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + tTR_rAcc[i] = tTR_rAcc[i] * cutlass.Float32(alpha_val) + + if cutlass.const_expr(self.generate_c): + self.store_c( + tiled_copy_r2s, + tma_atom_c, + warp_idx, + tTR_rAcc, + tRS_rC, + tRS_sC, + bSG_gC, + bSG_sC, + c_pipeline, + num_prev_subtiles, + subtile_idx, + ) + + acc_vec = tTR_rAcc.load() + if cutlass.const_expr(not self.enable_bias): + tCompute = cute.make_rmem_tensor(acc_vec.shape, self.acc_dtype) + if cutlass.const_expr(self.vectorized_f32): + for i in cutlass.range_constexpr(0, cute.size(tTR_rAcc), 2): + tCompute[i], tCompute[i + 1] = cute.arch.mul_packed_f32x2( + (acc_vec[i], acc_vec[i + 1]), + (mProb, mProb), + rnd="rn", + ftz=False, + ) + else: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + tCompute[i] = acc_vec[i] * mProb + else: + tCompute = tTR_rAcc + + acc_vec = tiled_copy_r2s.retile(tCompute).load() + tRS_rD.store(acc_vec.to(self.d_dtype)) + d_buffer = num_prev_subtiles % self.num_d_stage + num_prev_subtiles = num_prev_subtiles + 1 + cute.copy(tiled_copy_r2s, tRS_rD, tRS_sD[(None, None, None, d_buffer)]) + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier.arrive_and_wait() + if warp_idx == self.epilog_warp_id[0]: + cute.copy(tma_atom_d, bSG_sD[(None, d_buffer)], bSG_gD[(None, subtile_idx)]) + d_pipeline.producer_commit() + d_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + if cutlass.const_expr(self.enable_bias): + bias_pipeline.consumer_release(bias_consumer_state) + bias_consumer_state.advance() + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[0] >= cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + if cutlass.const_expr(self.generate_c): + c_pipeline.producer_tail() + d_pipeline.producer_tail() + + # ------------------------------------------------------------------ + # Internal: create extension based on weight_mode + # ------------------------------------------------------------------ + + @cute.jit + def _make_extension(self, workspace_ptr): + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + desc_workspace = TensormapWorkspace(workspace_ptr, ["b"]) + return DiscreteWeightGroupedGemmSchedExtension(tensormap_ctor=desc_workspace) + else: + return ContiguousGroupedGemmSchedExtension() diff --git a/python/cudnn/grouped_gemm/grouped_gemm_utils.py b/python/cudnn/grouped_gemm/grouped_gemm_utils.py new file mode 100644 index 000000000..d0d1d8e9f --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_utils.py @@ -0,0 +1,58 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +from contextlib import contextmanager +from enum import Enum +from typing import Iterator, Optional + +import torch +from cuda.bindings import driver as cuda + + +class GroupedGemmBackend(str, Enum): + BF16 = "bf16" + BLOCK_SCALED = "block_scaled" + + +@contextmanager +def _torch_stream_context(current_stream: Optional[cuda.CUstream], device: torch.device) -> Iterator[None]: + """Run PyTorch work on the CUDA stream used for the kernel launch.""" + if current_stream is None: + yield + return + handle = int(current_stream) + torch_current = torch.cuda.current_stream(device) + torch_default = torch.cuda.default_stream(device) + if handle == torch_current.cuda_stream: + launch_stream = torch_current + elif handle == torch_default.cuda_stream: + launch_stream = torch_default + else: + launch_stream = torch.cuda.ExternalStream(handle, device=device) + with torch.cuda.stream(launch_stream): + yield + + +def select_grouped_gemm_backend( + *, + operation, + a_dtype, + b_dtype, + scale_controls, + block_scaled_dtype_pairs, +): + bf16_operands = (a_dtype == torch.bfloat16, b_dtype == torch.bfloat16) + if any(bf16_operands): + if not all(bf16_operands): + raise ValueError(f"{operation}: mixed dtype families: a_dtype={a_dtype}, " f"b_dtype={b_dtype}") + forbidden = [name for name, value in scale_controls if value is not None] + if forbidden: + raise ValueError(f"{operation}: BF16 forbids scale control {forbidden[0]}") + return GroupedGemmBackend.BF16 + if (a_dtype, b_dtype) in block_scaled_dtype_pairs: + return GroupedGemmBackend.BLOCK_SCALED + raise ValueError(f"{operation}: unsupported dtype pair a_dtype={a_dtype}, " f"b_dtype={b_dtype}") + + +def backend_cache_key(backend, *components): + return (backend.value, *components) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/_bf16_api.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/_bf16_api.py new file mode 100644 index 000000000..2a29d894b --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/_bf16_api.py @@ -0,0 +1,514 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Private descriptor-first BF16 API for SM100 grouped GEMM wgrad.""" + +import os +import weakref +from typing import Optional, Tuple, Union + +import cutlass +import cutlass.cute as cute +from cuda.bindings import driver as cuda +from cutlass.cute.runtime import from_dlpack, make_fake_stream +import torch + +from cudnn.api_base import APIBase, TensorDesc +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor + +from ..grouped_gemm_utils import _torch_stream_context +from ..moe_utils import MoEWeightMode, WGradInputOrder +from .moe_grouped_gemm_wgrad import MoEGroupedGemmWgradBF16Kernel + +_OUTPUT_DTYPES = [torch.bfloat16, torch.float16, torch.float32] + + +class GroupedGemmWgradBf16API(APIBase): + """Descriptor-first lifecycle API for the source BF16 wgrad kernel.""" + + def __init__( + self, + sample_a: torch.Tensor, + sample_b: torch.Tensor, + sample_sfa: Optional[torch.Tensor], + sample_sfb: Optional[torch.Tensor], + sample_offsets: torch.Tensor, + sample_wgrad: Optional[torch.Tensor] = None, + sample_wgrad_expert: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + wgrad_shape: Optional[Tuple[int, int]] = None, + wgrad_dtype: Optional[torch.dtype] = None, + sample_global_scale_a: Optional[torch.Tensor] = None, + sample_global_scale_b: Optional[torch.Tensor] = None, + acc_dtype: torch.dtype = torch.float32, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + sf_vec_size: int = 16, + accumulate_on_output: bool = False, + input_order: Union[WGradInputOrder, str] = WGradInputOrder.Tensor2D, + ) -> None: + super().__init__() + self._warn_experimental_api() + self.input_order = WGradInputOrder(input_order) + if sample_wgrad is not None and num_experts is None: + self.weight_mode = MoEWeightMode.DENSE + elif sample_wgrad is None and num_experts is not None: + self.weight_mode = MoEWeightMode.DISCRETE + if wgrad_shape is None or wgrad_dtype is None: + raise ValueError("wgrad_shape and wgrad_dtype are required in discrete mode") + else: + raise ValueError("Provide either sample_wgrad for dense mode or " "(num_experts, wgrad_shape, wgrad_dtype) for discrete mode, but not both") + + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") + self.offsets_desc = self._make_tensor_desc(sample_offsets, name="sample_offsets") + self.wgrad_desc = self._make_tensor_desc(sample_wgrad, name="sample_wgrad") + self.single_expert_wgrad_desc = self._make_tensor_desc(sample_wgrad_expert, name="sample_wgrad_expert") + self.expert_cnt = self.wgrad_desc.shape[0] if self.weight_mode == MoEWeightMode.DENSE and self.wgrad_desc.ndim == 3 else int(num_experts or 0) + self.wgrad_shape = self.wgrad_desc.shape[1:] if self.weight_mode == MoEWeightMode.DENSE and self.wgrad_desc.ndim == 3 else tuple(wgrad_shape or ()) + self.wgrad_dtype = self.wgrad_desc.dtype if self.weight_mode == MoEWeightMode.DENSE else wgrad_dtype + if self.weight_mode == MoEWeightMode.DISCRETE and self.single_expert_wgrad_desc is None: + self.single_expert_wgrad_desc = TensorDesc( + dtype=self.wgrad_dtype, + shape=self.wgrad_shape, + stride=(self.wgrad_shape[1], 1), + stride_order=(1, 0), + device=self.a_desc.device, + name="single_expert_wgrad", + ) + + self.acc_dtype = acc_dtype + self.mma_tiler_mn = tuple(mma_tiler_mn) + self.use_2cta_instrs = self.mma_tiler_mn[0] == 256 + self.cluster_shape_mn = tuple(cluster_shape_mn or ((2, 1) if self.use_2cta_instrs else (1, 1))) + self.accumulate_on_output = accumulate_on_output + self.sf_vec_size = sf_vec_size + self._scale_controls = ( + sample_sfa, + sample_sfb, + sample_global_scale_a, + sample_global_scale_b, + ) + self._kernel = MoEGroupedGemmWgradBF16Kernel + self._workspace: Optional[torch.Tensor] = None + self._compile_wgrad_ptrs: Optional[torch.Tensor] = None + self._single_expert_placeholder: Optional[torch.Tensor] = None + self._validated_offsets: dict[int, tuple] = {} + self._validated_pointer_values: dict[int, tuple] = {} + self._sample_offset_values = self._copy_values_to_host(sample_offsets) + self._sample_offsets_ref = weakref.ref(sample_offsets) + self._sample_offsets_version = int(sample_offsets._version) + self._sample_data_ptrs = { + name: tensor.data_ptr() + for name, tensor in ( + ("sample_a", sample_a), + ("sample_b", sample_b), + ("sample_offsets", sample_offsets), + ("sample_wgrad", sample_wgrad), + ("sample_wgrad_expert", sample_wgrad_expert), + ) + if tensor is not None + } + self.num_cluster_overlap_margin = int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")) + self.a_major: Optional[str] = None + self.b_major: Optional[str] = None + + @staticmethod + def _copy_values_to_host(tensor: torch.Tensor) -> Tuple[int, ...]: + return tuple(int(value) for value in tensor.detach().cpu().tolist()) + + @staticmethod + def _is_validation_cached(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> bool: + cached = cache.get(id(tensor)) + return bool(cached and cached[0]() is tensor and cached[1] == int(tensor._version) and cached[2] == extra) + + @staticmethod + def _remember_validation(cache: dict[int, tuple], tensor: torch.Tensor, extra) -> None: + key = id(tensor) + + def discard(_reference, *, cache=cache, key=key): + cache.pop(key, None) + + cache[key] = (weakref.ref(tensor, discard), int(tensor._version), extra) + + @staticmethod + def _validate_offset_sequence(values: Tuple[int, ...], *, expert_cnt: int, tokens_sum: int) -> Tuple[int, ...]: + if len(values) != expert_cnt: + raise ValueError(f"sample_offsets length mismatch: expected {expert_cnt}, got {len(values)}") + groups = [] + previous = 0 + for index, value in enumerate(values): + if value < previous: + raise ValueError("sample_offsets must be a non-decreasing cumulative sum; " f"index {index} is {value} after {previous}") + group_k = value - previous + if group_k % MoEGroupedGemmWgradBF16Kernel.FIX_PAD_SIZE != 0: + raise ValueError(f"sample_offsets group {index} must be 256-aligned, got {group_k}") + groups.append(group_k) + previous = value + if not values or values[-1] != tokens_sum: + raise ValueError(f"sample_offsets last value must equal total tokens {tokens_sum}, " f"got {values[-1] if values else None}") + return tuple(groups) + + def _validate_offsets_once(self, offsets: torch.Tensor, *, tokens_sum: int) -> None: + extra = (self.expert_cnt, tokens_sum) + if self._is_validation_cached(self._validated_offsets, offsets, extra): + return + values = self._copy_values_to_host(offsets) + self._validate_offset_sequence(values, expert_cnt=self.expert_cnt, tokens_sum=tokens_sum) + self._remember_validation(self._validated_offsets, offsets, extra) + + def _validate_pointer_values_once(self, pointers: torch.Tensor) -> None: + if self._is_validation_cached(self._validated_pointer_values, pointers, self.expert_cnt): + return + values = self._copy_values_to_host(pointers) + if any(value == 0 or value % 16 != 0 for value in values): + raise ValueError("wgrad_ptrs entries must be non-null and 16-byte aligned") + self._remember_validation(self._validated_pointer_values, pointers, self.expert_cnt) + + @staticmethod + def _validate_pointer_array_alignment(tensor: torch.Tensor) -> None: + if tensor.data_ptr() % 8 != 0: + raise ValueError("wgrad_ptrs data pointer must be 8-byte aligned") + + @staticmethod + def _validate_data_alignment(tensor: torch.Tensor, name: str, alignment: int = 16) -> None: + if tensor.data_ptr() % alignment != 0: + raise ValueError(f"{name} data pointer must be {alignment}-byte aligned") + + @staticmethod + def _record_pointer_stream(pointers: torch.Tensor, current_stream: cuda.CUstream) -> None: + handle = int(current_stream) + torch_current = torch.cuda.current_stream(pointers.device) + torch_default = torch.cuda.default_stream(pointers.device) + if handle == torch_current.cuda_stream: + launch_stream = torch_current + elif handle == torch_default.cuda_stream: + launch_stream = torch_default + else: + launch_stream = torch.cuda.ExternalStream(handle, device=pointers.device) + pointers.record_stream(launch_stream) + + @staticmethod + def _infer_a_major(desc: TensorDesc) -> str: + m, tokens = desc.shape + if desc.stride == (tokens, 1): + return "k" + if desc.stride == (1, m): + return "m" + raise ValueError(f"A tensor must use a supported K-major or M-major layout, got stride {desc.stride}") + + @staticmethod + def _infer_b_major(desc: TensorDesc) -> str: + tokens, n = desc.shape + if desc.stride == (1, tokens): + return "k" + if desc.stride == (n, 1): + return "n" + raise ValueError(f"B tensor must use a supported K-major or N-major layout, got stride {desc.stride}") + + @staticmethod + def _expect_device(desc: TensorDesc, device: torch.device, name: str) -> None: + if desc.device != device: + raise ValueError(f"{name} must be on {device}, got {desc.device}") + + def check_support(self) -> bool: + if self.a_desc.ndim != 2: + raise ValueError(f"sample_a must be rank-2, got {self.a_desc.shape}") + if self.b_desc.ndim != 2: + raise ValueError(f"sample_b must be rank-2, got {self.b_desc.shape}") + m, tokens_sum = self.a_desc.shape + tokens_b, n = self.b_desc.shape + if tokens_b != tokens_sum: + raise ValueError(f"sample_a and sample_b token dimensions must match, got {tokens_sum} and {tokens_b}") + self.a_major = self._infer_a_major(self.a_desc) + self.b_major = self._infer_b_major(self.b_desc) + self._check_dtype(self.a_desc, torch.bfloat16, "sample_a") + self._check_dtype(self.b_desc, torch.bfloat16, "sample_b") + self._check_dtype(self.offsets_desc, torch.int32, "sample_offsets") + self._check_dtype(self.wgrad_dtype, _OUTPUT_DTYPES, "wgrad_dtype") + if self.acc_dtype != torch.float32: + raise ValueError(f"acc_dtype must be torch.float32, got {self.acc_dtype}") + if any(control is not None for control in self._scale_controls): + raise ValueError("BF16 wgrad forbids scale and global-scale tensors") + if self.sf_vec_size != 16: + raise ValueError(f"BF16 wgrad requires sf_vec_size=16, got {self.sf_vec_size}") + if self.offsets_desc.shape != (self.expert_cnt,): + raise ValueError(f"sample_offsets must have shape {(self.expert_cnt,)}, got {self.offsets_desc.shape}") + if self.offsets_desc.stride != (1,): + raise ValueError("sample_offsets must be contiguous") + + if self.weight_mode == MoEWeightMode.DENSE: + if self.wgrad_desc.ndim != 3: + raise ValueError(f"sample_wgrad must be rank-3, got {self.wgrad_desc.shape}") + if self.wgrad_desc.shape != (self.expert_cnt, m, n): + raise ValueError(f"sample_wgrad shape mismatch: expected {(self.expert_cnt, m, n)}, got {self.wgrad_desc.shape}") + if self.wgrad_desc.stride != (m * n, n, 1): + raise ValueError("sample_wgrad must be contiguous in expert/M/N order") + output_desc = self.wgrad_desc + else: + if self.wgrad_shape != (m, n): + raise ValueError(f"wgrad_shape mismatch: expected {(m, n)}, got {self.wgrad_shape}") + output_desc = self.single_expert_wgrad_desc + if output_desc.shape not in ((m, n), (m, n, 1)): + raise ValueError(f"sample_wgrad_expert shape mismatch: expected {(m, n)}, got {output_desc.shape}") + expected_stride = (n, 1) if output_desc.ndim == 2 else (n, 1, 1) + if output_desc.stride != expected_stride: + raise ValueError("sample_wgrad_expert must be contiguous in M/N order") + self._check_dtype(output_desc, self.wgrad_dtype, "sample_wgrad_expert") + + device = self.a_desc.device + for desc, name in ( + (self.b_desc, "sample_b"), + (self.offsets_desc, "sample_offsets"), + (output_desc, "sample_wgrad"), + ): + self._expect_device(desc, device, name) + for name, pointer in self._sample_data_ptrs.items(): + alignment = 4 if name == "sample_offsets" else 16 + if pointer % alignment: + raise ValueError(f"{name} data pointer must be {alignment}-byte aligned") + + groups = self._validate_offset_sequence( + self._sample_offset_values, + expert_cnt=self.expert_cnt, + tokens_sum=tokens_sum, + ) + sample_offsets = self._sample_offsets_ref() + if sample_offsets is not None and int(sample_offsets._version) == self._sample_offsets_version: + self._remember_validation( + self._validated_offsets, + sample_offsets, + (self.expert_cnt, tokens_sum), + ) + elif sample_offsets is not None: + self._validate_offsets_once(sample_offsets, tokens_sum=tokens_sum) + + if not self._kernel.can_implement( + _convert_to_cutlass_data_type(torch.bfloat16), + _convert_to_cutlass_data_type(self.wgrad_dtype), + _convert_to_cutlass_data_type(self.acc_dtype), + self.use_2cta_instrs, + self.mma_tiler_mn, + self.cluster_shape_mn, + m, + n, + list(groups), + self.expert_cnt, + self.a_major, + self.b_major, + self.weight_mode, + self.input_order, + ): + raise ValueError("Unsupported BF16 grouped GEMM wgrad configuration: check mma_tiler, cluster, and alignment") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + major, minor = torch.cuda.get_device_capability(self.a_desc.device) + capability = major * 10 + minor + if capability < 100: + raise RuntimeError(f"GroupedGemmWgradSm100 requires SM100+, found SM{capability} on {self.a_desc.device}") + self._is_supported = True + return True + + def compile(self) -> None: + self._ensure_support_checked() + if self._compiled_kernel is not None: + return + kernel = self._kernel( + acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), + use_2cta_instrs=self.use_2cta_instrs, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + accumulate_on_output=self.accumulate_on_output, + expert_cnt=self.expert_cnt, + weight_mode=self.weight_mode, + input_order=self.input_order, + ) + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - self.num_cluster_overlap_margin + if max_active_clusters <= 0: + raise ValueError("max_active_clusters must be > 0 after applying CUDNNFE_CLUSTER_OVERLAP_MARGIN") + self._workspace = torch.empty( + max(kernel.get_workspace_bytes(), 1), + dtype=torch.uint8, + device=self.a_desc.device, + ) + self._validate_data_alignment(self._workspace, "workspace", 128) + workspace_fake = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + a_fake = self._make_fake_cute_compact_tensor( + dtype=self.a_desc.dtype, + shape=self.a_desc.shape, + stride_order=self.a_desc.stride_order, + assumed_align=16, + dynamic_mode=1, + divisibility=256, + ) + b_fake = self._make_fake_cute_compact_tensor( + dtype=self.b_desc.dtype, + shape=self.b_desc.shape, + stride_order=self.b_desc.stride_order, + assumed_align=16, + dynamic_mode=0, + divisibility=256, + ) + offsets_fake = self._make_fake_cute_tensor_from_desc(self.offsets_desc, assumed_align=4) + if self.weight_mode == MoEWeightMode.DENSE: + out_fake = self._make_fake_cute_tensor_from_desc(self.wgrad_desc, assumed_align=16) + single_expert_fake = None + else: + self._compile_wgrad_ptrs = torch.empty((self.expert_cnt,), dtype=torch.int64, device=self.a_desc.device) + self._validate_pointer_array_alignment(self._compile_wgrad_ptrs) + out_fake = from_dlpack(self._compile_wgrad_ptrs, assumed_align=8).iterator + single_expert_fake = self._make_fake_cute_tensor_from_desc(self.single_expert_wgrad_desc, assumed_align=16) + raw_compiled = cute.compile( + kernel, + a_fake, + b_fake, + out_fake, + offsets_fake, + workspace_fake, + max_active_clusters, + fake_stream, + single_expert_fake, + options="--enable-tvm-ffi", + ) + cached_workspace = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) + if self.weight_mode == MoEWeightMode.DISCRETE: + self._single_expert_placeholder = torch.empty_strided( + self.single_expert_wgrad_desc.shape, + self.single_expert_wgrad_desc.stride, + dtype=self.single_expert_wgrad_desc.dtype, + device=self.single_expert_wgrad_desc.device, + ) + self._validate_data_alignment(self._single_expert_placeholder, "single expert placeholder") + cached_single_expert = from_dlpack( + self._single_expert_placeholder, + assumed_align=16, + enable_tvm_ffi=True, + ) + else: + cached_single_expert = None + + def tensor_api(a_tensor, b_tensor, output, offsets, stream) -> None: + out_arg = output if self.weight_mode == MoEWeightMode.DENSE else int(output.data_ptr()) + raw_compiled( + a_tensor, + b_tensor, + out_arg, + offsets, + cached_workspace, + stream, + cached_single_expert, + ) + + self._compiled_kernel = tensor_api + + def _validate_live_input( + self, + tensor: torch.Tensor, + sample: TensorDesc, + name: str, + *, + token_axis: int, + ) -> TensorDesc: + desc = self._make_tensor_desc(tensor, name=name) + if desc.dtype != sample.dtype: + raise ValueError(f"{name} dtype mismatch: expected {sample.dtype}, got {desc.dtype}") + if desc.device != sample.device: + raise ValueError(f"{name} device mismatch: expected {sample.device}, got {desc.device}") + for axis, (actual, expected) in enumerate(zip(desc.shape, sample.shape)): + if axis != token_axis and actual != expected: + raise ValueError(f"{name} shape mismatch: expected static dimension {expected} at axis {axis}, got {actual}") + expected_major = self.a_major if name == "a_tensor" else self.b_major + actual_major = self._infer_a_major(desc) if name == "a_tensor" else self._infer_b_major(desc) + if actual_major != expected_major: + raise ValueError(f"{name} layout mismatch: expected {expected_major}-major, got {actual_major}-major") + return desc + + def _validate_live_output(self, tensor: torch.Tensor) -> None: + desc = self._make_tensor_desc(tensor, name="wgrad_tensor") + expected = (self.expert_cnt, *self.wgrad_shape) + if desc.shape != expected: + raise ValueError(f"wgrad_tensor shape mismatch: expected {expected}, got {desc.shape}") + if desc.stride != ( + self.wgrad_shape[0] * self.wgrad_shape[1], + self.wgrad_shape[1], + 1, + ): + raise ValueError("wgrad_tensor must be contiguous in expert/M/N order") + if desc.dtype != self.wgrad_dtype: + raise ValueError(f"wgrad_tensor dtype mismatch: expected {self.wgrad_dtype}, got {desc.dtype}") + if desc.device != self.a_desc.device: + raise ValueError(f"wgrad_tensor device mismatch: expected {self.a_desc.device}, got {desc.device}") + + def execute( + self, + a_tensor: torch.Tensor, + b_tensor: torch.Tensor, + sfa_tensor: Optional[torch.Tensor], + sfb_tensor: Optional[torch.Tensor], + offsets_tensor: torch.Tensor, + wgrad_tensor: Optional[torch.Tensor] = None, + wgrad_ptrs: Optional[torch.Tensor] = None, + global_scale_a: Optional[torch.Tensor] = None, + global_scale_b: Optional[torch.Tensor] = None, + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + current_stream = self._get_default_stream(current_stream) + if self._compiled_kernel is None: + raise RuntimeError("Kernel not compiled; call compile() first") + forbidden = ( + ("sfa_tensor", sfa_tensor), + ("sfb_tensor", sfb_tensor), + ("global_scale_a", global_scale_a), + ("global_scale_b", global_scale_b), + ) + for name, value in forbidden: + if value is not None: + raise ValueError(f"BF16 forbids scale control {name}") + a_desc = self._validate_live_input(a_tensor, self.a_desc, "a_tensor", token_axis=1) + b_desc = self._validate_live_input(b_tensor, self.b_desc, "b_tensor", token_axis=0) + tokens_sum = a_desc.shape[1] + if b_desc.shape[0] != tokens_sum: + raise ValueError("a_tensor and b_tensor token dimensions must match") + offsets_desc = self._make_tensor_desc(offsets_tensor, name="offsets_tensor") + if offsets_desc.shape != (self.expert_cnt,) or offsets_desc.stride != (1,) or offsets_desc.dtype != torch.int32: + raise ValueError("offsets_tensor must be a contiguous rank-1 int32 tensor with one entry per expert") + if offsets_desc.device != self.a_desc.device: + raise ValueError(f"offsets_tensor device mismatch: expected {self.a_desc.device}, got {offsets_desc.device}") + self._validate_offsets_once(offsets_tensor, tokens_sum=tokens_sum) + self._validate_data_alignment(a_tensor, "a_tensor") + self._validate_data_alignment(b_tensor, "b_tensor") + self._validate_data_alignment(offsets_tensor, "offsets_tensor", 4) + + if self.weight_mode == MoEWeightMode.DENSE: + if wgrad_tensor is None or wgrad_ptrs is not None: + raise ValueError("Dense execution requires wgrad_tensor and forbids wgrad_ptrs") + self._validate_live_output(wgrad_tensor) + self._validate_data_alignment(wgrad_tensor, "wgrad_tensor") + output = wgrad_tensor + else: + if wgrad_tensor is not None: + self._validate_live_output(wgrad_tensor) + self._validate_data_alignment(wgrad_tensor, "wgrad_tensor") + generated_wgrad_ptrs = wgrad_ptrs is None + if wgrad_ptrs is None: + if wgrad_tensor is None: + raise ValueError("Discrete execution requires wgrad_tensor or wgrad_ptrs") + stride_bytes = wgrad_tensor.stride(0) * wgrad_tensor.element_size() + with _torch_stream_context(current_stream, wgrad_tensor.device): + wgrad_ptrs = torch.tensor( + [wgrad_tensor.data_ptr() + index * stride_bytes for index in range(self.expert_cnt)], + dtype=torch.int64, + device=wgrad_tensor.device, + ) + _require_pointer_tensor(wgrad_ptrs, "wgrad_ptrs", self.expert_cnt) + if wgrad_ptrs.device != self.a_desc.device: + raise ValueError(f"wgrad_ptrs must be on {self.a_desc.device}, got {wgrad_ptrs.device}") + self._validate_pointer_array_alignment(wgrad_ptrs) + if not generated_wgrad_ptrs: + self._validate_pointer_values_once(wgrad_ptrs) + self._record_pointer_stream(wgrad_ptrs, current_stream) + output = wgrad_ptrs + self._compiled_kernel(a_tensor, b_tensor, output, offsets_tensor, current_stream) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/_blockscaled_api.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/_blockscaled_api.py new file mode 100644 index 000000000..7d73d6eed --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/_blockscaled_api.py @@ -0,0 +1,539 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Unified FE API for grouped GEMM wgrad on SM100+.""" + +from typing import Optional, Tuple, Union + +import torch +import cutlass +import cutlass.cute as cute +from cuda.bindings import driver as cuda +from cutlass.cute.runtime import from_dlpack, make_fake_stream + +from cudnn.api_base import APIBase, TensorDesc, ceil_div, is_power_of_2 +from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor + +from .moe_blockscaled_grouped_gemm_wgrad import BlockScaledMoEGroupedGemmWgradKernel +from ..moe_utils import MoEWeightMode, WGradInputOrder + + +def _round_up(a: int, b: int) -> int: + return ceil_div(a, b) * b + + +class GroupedGemmWgradBlockScaledAPI(APIBase): + """Unified grouped GEMM wgrad FE API for SM100+ GPUs.""" + + def __init__( + self, + sample_a: torch.Tensor, + sample_b: torch.Tensor, + sample_sfa: torch.Tensor, + sample_sfb: torch.Tensor, + sample_offsets: torch.Tensor, + sample_wgrad: Optional[torch.Tensor] = None, + sample_wgrad_expert: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + wgrad_shape: Optional[Tuple[int, int]] = None, + wgrad_dtype: Optional[torch.dtype] = None, + sample_global_scale_a: Optional[torch.Tensor] = None, + sample_global_scale_b: Optional[torch.Tensor] = None, + acc_dtype: torch.dtype = torch.float32, + mma_tiler_mn: Tuple[int, int] = (256, 256), + cluster_shape_mn: Optional[Tuple[int, int]] = None, + sf_vec_size: int = 16, + accumulate_on_output: bool = False, + input_order: Union[WGradInputOrder, str] = WGradInputOrder.Tensor2D, + ): + super().__init__() + self._warn_experimental_api() + self.input_order = WGradInputOrder(input_order) + + if sample_wgrad is not None and num_experts is None: + self.weight_mode = MoEWeightMode.DENSE + elif sample_wgrad is None and num_experts is not None: + self.weight_mode = MoEWeightMode.DISCRETE + if wgrad_shape is None or wgrad_dtype is None: + raise ValueError("wgrad_shape and wgrad_dtype are required in discrete mode") + else: + raise ValueError("Provide either sample_wgrad for dense mode or " "(num_experts, wgrad_shape, wgrad_dtype) for discrete mode, but not both") + + self._interpret_uint8_as_fp4x2 = True + self.sample_a_tensor = sample_a if self._is_fp4x2(sample_a) else None + self.sample_b_tensor = sample_b if self._is_fp4x2(sample_b) else None + self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") + self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") + self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") + self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb") + self.offsets_desc = self._make_tensor_desc(sample_offsets, name="sample_offsets") + self.global_scale_a_desc = self._make_tensor_desc(sample_global_scale_a, name="sample_global_scale_a") + self.global_scale_b_desc = self._make_tensor_desc(sample_global_scale_b, name="sample_global_scale_b") + self.sf_vec_size = sf_vec_size + tokens_sum_a = self.a_desc.shape[1] + tokens_sum_b = self.b_desc.shape[0] + self._value_error_if( + tokens_sum_a != tokens_sum_b, + f"sample_a and sample_b token dimensions must match, got {tokens_sum_a} and {tokens_sum_b}", + ) + self._offset_values = self._validate_offsets(sample_offsets, tokens_sum_a, name="sample_offsets") + self._scale_cols = _round_up(ceil_div(tokens_sum_a, self.sf_vec_size), 4) + + if self.weight_mode == MoEWeightMode.DENSE: + self.wgrad_desc = self._make_tensor_desc(sample_wgrad, name="sample_wgrad") + self.expert_cnt = self.wgrad_desc.shape[0] + self.wgrad_shape = self.wgrad_desc.shape[1:] + self.wgrad_dtype = self.wgrad_desc.dtype + self.single_expert_wgrad_desc = TensorDesc( + dtype=self.wgrad_desc.dtype, + shape=self.wgrad_desc.shape[1:], + stride=self.wgrad_desc.stride[1:], + stride_order=tuple(i for i, s in sorted(enumerate(self.wgrad_desc.stride[1:]), key=lambda x: x[1])), + device=self.wgrad_desc.device, + name="single_expert_wgrad", + ) + else: # MoEWeightMode.DISCRETE + self.expert_cnt = num_experts + self.wgrad_shape = tuple(wgrad_shape) + self.wgrad_dtype = wgrad_dtype + self.wgrad_desc = None + if sample_wgrad_expert is not None: + self.single_expert_wgrad_desc = self._make_tensor_desc( + sample_wgrad_expert, + name="sample_wgrad_expert", + ) + else: + self.single_expert_wgrad_desc = TensorDesc( + dtype=wgrad_dtype, + shape=self.wgrad_shape, + stride=(self.wgrad_shape[1], 1), + stride_order=(1, 0), + device=self.a_desc.device, + name="single_expert_wgrad", + ) + + self.acc_dtype = acc_dtype + self.mma_tiler_mn = mma_tiler_mn + self.use_2cta_instrs = mma_tiler_mn[0] == 256 + self.cluster_shape_mn = cluster_shape_mn or ((2, 1) if self.use_2cta_instrs else (1, 1)) + self.accumulate_on_output = accumulate_on_output + self._kernel = BlockScaledMoEGroupedGemmWgradKernel + self._workspace = None + + def _validate_offsets(self, offsets_tensor: torch.Tensor, tokens_sum: int, name: str) -> Tuple[int, ...]: + self._value_error_if(offsets_tensor.ndim != 1, f"{name} must be rank-1, got shape {tuple(offsets_tensor.shape)}") + + offset_values = tuple(int(offset) for offset in offsets_tensor.detach().cpu().tolist()) + prev_offset = 0 + for idx, offset in enumerate(offset_values): + self._value_error_if( + offset < prev_offset, + f"{name} must be a non-decreasing cumulative sum, but index {idx} has {offset} after {prev_offset}", + ) + prev_offset = offset + + if offset_values: + self._value_error_if( + offset_values[-1] > tokens_sum, + f"{name} last value must not exceed total tokens {tokens_sum}, got {offset_values[-1]}", + ) + else: + self._value_error_if(tokens_sum != 0, f"{name} cannot be empty when total tokens is {tokens_sum}") + + return offset_values + + def check_support(self) -> bool: + m, tokens_sum = self._tensor_shape(self.a_desc, name="sample_a") + _, n = self._tensor_shape(self.b_desc, name="sample_b") + + _ = self._check_tensor_shape(self.a_desc, (m, tokens_sum), "sample_a") + _ = self._check_tensor_shape(self.b_desc, (tokens_sum, n), "sample_b") + _ = self._check_tensor_shape(self.sfa_desc, (_round_up(m, 128), self._scale_cols), "sample_sfa") + _ = self._check_tensor_shape(self.sfb_desc, (_round_up(n, 128), self._scale_cols), "sample_sfb") + _ = self._check_tensor_shape(self.offsets_desc, (self.expert_cnt,), "sample_offsets") + + dtype = self._check_dtype(self.a_desc, [torch.float4_e2m1fn_x2, torch.uint8, torch.float8_e5m2, torch.float8_e4m3fn], "sample_a") + self._check_dtype(self.b_desc, dtype, "sample_b", extra_error_msg="sample_b must have the same dtype as sample_a") + self._check_dtype( + self.sfa_desc, + [torch.float8_e8m0fnu, torch.float8_e4m3fn], + "sample_sfa", + extra_error_msg="sample_sfa must have dtype float8_e8m0fnu or float8_e4m3fn", + ) + self._check_dtype( + self.sfb_desc, + [torch.float8_e8m0fnu, torch.float8_e4m3fn], + "sample_sfb", + extra_error_msg="sample_sfb must have dtype float8_e8m0fnu or float8_e4m3fn", + ) + self._check_dtype(self.offsets_desc, torch.int32, "sample_offsets", extra_error_msg="sample_offsets must be int32") + self._check_dtype( + self.wgrad_dtype, [torch.bfloat16, torch.float16, torch.float32], "wgrad_dtype", extra_error_msg="wgrad_dtype must be bfloat16, float16, or float32" + ) + + if self.weight_mode == MoEWeightMode.DENSE: + self._check_tensor_shape(self.wgrad_desc, (self.expert_cnt, m, n), "sample_wgrad") + else: + self._check_tensor_shape(self.wgrad_shape, (m, n), "wgrad_shape") + self._check_tensor_shape(self.single_expert_wgrad_desc, (m, n), "single_expert_wgrad") + self._check_dtype( + self.single_expert_wgrad_desc, + self.wgrad_dtype, + "sample_wgrad_expert", + extra_error_msg="sample_wgrad_expert must have the same dtype as wgrad_dtype", + ) + + self._value_error_if(self.mma_tiler_mn[0] not in (128, 256), f"mma_tiler_mn[0] must be 128 or 256, got {self.mma_tiler_mn[0]}") + self._value_error_if(self.mma_tiler_mn[1] not in (128, 256), f"mma_tiler_mn[1] must be 128 or 256, got {self.mma_tiler_mn[1]}") + self._value_error_if( + self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0, + f"cluster_shape_mn[0] must be divisible by 2 when use_2cta_instrs=True, got {self.cluster_shape_mn[0]}", + ) + self._value_error_if(self.cluster_shape_mn[0] * self.cluster_shape_mn[1] > 16, f"cluster shape product must be <= 16, got {self.cluster_shape_mn}") + self._value_error_if( + not (is_power_of_2(self.cluster_shape_mn[0]) and is_power_of_2(self.cluster_shape_mn[1])), + f"cluster shape values must be powers of 2, got {self.cluster_shape_mn}", + ) + + has_global_scale = self.global_scale_a_desc is not None or self.global_scale_b_desc is not None + if has_global_scale: + self._value_error_if( + self.global_scale_a_desc is None or self.global_scale_b_desc is None, + "sample_global_scale_a and sample_global_scale_b must be provided together", + ) + self._value_error_if( + self.global_scale_a_desc.shape != (self.expert_cnt,), + f"sample_global_scale_a must have shape {(self.expert_cnt,)}, got {self.global_scale_a_desc.shape}", + ) + self._value_error_if( + self.global_scale_b_desc.shape != (self.expert_cnt,), + f"sample_global_scale_b must have shape {(self.expert_cnt,)}, got {self.global_scale_b_desc.shape}", + ) + self._check_dtype(self.global_scale_a_desc, torch.float32, "sample_global_scale_a") + self._check_dtype(self.global_scale_b_desc, torch.float32, "sample_global_scale_b") + + requires_global_scale = ( + self._is_fp4x2(self.a_desc) and self.sf_vec_size == 16 and self.sfa_desc.dtype == torch.float8_e4m3fn and self.sfb_desc.dtype == torch.float8_e4m3fn + ) + self._value_error_if(requires_global_scale and not has_global_scale, "NVFP4 wgrad requires sample_global_scale_a and sample_global_scale_b") + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available") + device = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(device) + compute_capability = major * 10 + minor + if compute_capability < 100: + raise RuntimeError(f"GroupedGemmWgrad requires SM100+ compute capability, but found SM{compute_capability} on device {device}") + + self._is_supported = True + return True + + def compile(self) -> None: + self._ensure_support_checked() + if self._compiled_kernel is not None: + return + + kernel = self._kernel( + sf_vec_size=self.sf_vec_size, + acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), + use_2cta_instrs=self.use_2cta_instrs, + mma_tiler_mn=self.mma_tiler_mn, + cluster_shape_mn=self.cluster_shape_mn, + accumulate_on_output=self.accumulate_on_output, + expert_cnt=self.expert_cnt, + weight_mode=self.weight_mode, + input_order=self.input_order, + ) + + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) + self._workspace = torch.empty(max(kernel.get_workspace_bytes(), 1), dtype=torch.uint8, device=self.a_desc.device) + fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) + + if self.weight_mode == MoEWeightMode.DENSE: + self._compile_dense(kernel, max_active_clusters, fake_stream) + else: + self._compile_discrete(kernel, max_active_clusters, fake_stream) + + if self.sample_a_tensor is not None: + del self.sample_a_tensor + if self.sample_b_tensor is not None: + del self.sample_b_tensor + + def _compile_dense(self, kernel, max_active_clusters, fake_stream) -> None: + a_fake = ( + from_dlpack(self.sample_a_tensor, assumed_align=16, enable_tvm_ffi=True).mark_compact_shape_dynamic( + mode=1, + stride_order=self.sample_a_tensor.dim_order(), + divisibility=16, + ) + if self.sample_a_tensor is not None + else self._make_fake_cute_compact_tensor( + dtype=self.a_desc.dtype, + shape=self.a_desc.shape, + stride_order=self.a_desc.stride_order, + assumed_align=16, + dynamic_mode=1, + divisibility=16, + ) + ) + b_fake = ( + from_dlpack(self.sample_b_tensor, assumed_align=16, enable_tvm_ffi=True).mark_compact_shape_dynamic( + mode=0, + stride_order=self.sample_b_tensor.dim_order(), + divisibility=16, + ) + if self.sample_b_tensor is not None + else self._make_fake_cute_compact_tensor( + dtype=self.b_desc.dtype, + shape=self.b_desc.shape, + stride_order=self.b_desc.stride_order, + assumed_align=16, + dynamic_mode=0, + divisibility=16, + ) + ) + sfa_fake = self._make_fake_cute_compact_tensor( + dtype=self.sfa_desc.dtype, + shape=self.sfa_desc.shape, + stride_order=self.sfa_desc.stride_order, + assumed_align=16, + dynamic_mode=1, + divisibility=4, + ) + sfb_fake = self._make_fake_cute_compact_tensor( + dtype=self.sfb_desc.dtype, + shape=self.sfb_desc.shape, + stride_order=self.sfb_desc.stride_order, + assumed_align=16, + dynamic_mode=1, + divisibility=4, + ) + wgrad_fake = self._make_fake_cute_tensor_from_desc(self.wgrad_desc, assumed_align=16) + offsets_fake = self._make_fake_cute_tensor_from_desc(self.offsets_desc, assumed_align=4) + workspace_fake = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) + gs_a_fake = self._make_fake_cute_tensor_from_desc(self.global_scale_a_desc, assumed_align=4) + gs_b_fake = self._make_fake_cute_tensor_from_desc(self.global_scale_b_desc, assumed_align=4) + + compiled = cute.compile( + kernel, + a_fake, + b_fake, + sfa_fake, + sfb_fake, + wgrad_fake, + offsets_fake, + workspace_fake, + max_active_clusters, + fake_stream, + gs_a_fake, + gs_b_fake, + None, + options="--enable-tvm-ffi", + ) + + cached_workspace = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) + + def tensor_api( + a_tensor: torch.Tensor, + b_tensor: torch.Tensor, + sfa_tensor: torch.Tensor, + sfb_tensor: torch.Tensor, + wgrad_tensor: torch.Tensor, + offsets_tensor: torch.Tensor, + stream: cuda.CUstream, + global_scale_a: Optional[torch.Tensor], + global_scale_b: Optional[torch.Tensor], + ) -> None: + compiled( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + wgrad_tensor, + offsets_tensor, + cached_workspace, + stream, + global_scale_a, + global_scale_b, + None, + ) + + self._compiled_kernel = tensor_api + + def _compile_discrete(self, kernel, max_active_clusters, fake_stream) -> None: + a_fake = ( + from_dlpack(self.sample_a_tensor, assumed_align=16, enable_tvm_ffi=True).mark_compact_shape_dynamic( + mode=1, + stride_order=self.sample_a_tensor.dim_order(), + divisibility=16, + ) + if self.sample_a_tensor is not None + else self._make_fake_cute_compact_tensor( + dtype=self.a_desc.dtype, + shape=self.a_desc.shape, + stride_order=self.a_desc.stride_order, + assumed_align=16, + dynamic_mode=1, + divisibility=16, + ) + ) + b_fake = ( + from_dlpack(self.sample_b_tensor, assumed_align=16, enable_tvm_ffi=True).mark_compact_shape_dynamic( + mode=0, + stride_order=self.sample_b_tensor.dim_order(), + divisibility=16, + ) + if self.sample_b_tensor is not None + else self._make_fake_cute_compact_tensor( + dtype=self.b_desc.dtype, + shape=self.b_desc.shape, + stride_order=self.b_desc.stride_order, + assumed_align=16, + dynamic_mode=0, + divisibility=16, + ) + ) + sfa_fake = self._make_fake_cute_compact_tensor( + dtype=self.sfa_desc.dtype, + shape=self.sfa_desc.shape, + stride_order=self.sfa_desc.stride_order, + assumed_align=16, + dynamic_mode=1, + divisibility=4, + ) + sfb_fake = self._make_fake_cute_compact_tensor( + dtype=self.sfb_desc.dtype, + shape=self.sfb_desc.shape, + stride_order=self.sfb_desc.stride_order, + assumed_align=16, + dynamic_mode=1, + divisibility=4, + ) + offsets_fake = self._make_fake_cute_tensor_from_desc(self.offsets_desc, assumed_align=4) + workspace_fake = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) + gs_a_fake = self._make_fake_cute_tensor_from_desc(self.global_scale_a_desc, assumed_align=4) + gs_b_fake = self._make_fake_cute_tensor_from_desc(self.global_scale_b_desc, assumed_align=4) + wgrad_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device=self.a_desc.device) + wgrad_ptrs_fake = from_dlpack(wgrad_ptrs_placeholder, assumed_align=8, enable_tvm_ffi=True).iterator + single_expert_fake = self._make_fake_cute_tensor( + dtype=self.single_expert_wgrad_desc.dtype, + shape=self.single_expert_wgrad_desc.shape, + stride=self.single_expert_wgrad_desc.stride, + assumed_align=16, + ) + + compiled = cute.compile( + kernel, + a_fake, + b_fake, + sfa_fake, + sfb_fake, + wgrad_ptrs_fake, + offsets_fake, + workspace_fake, + max_active_clusters, + fake_stream, + gs_a_fake, + gs_b_fake, + single_expert_fake, + options="--enable-tvm-ffi", + ) + + cached_workspace = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) + single_expert_placeholder = torch.empty_strided( + self.single_expert_wgrad_desc.shape, + self.single_expert_wgrad_desc.stride, + dtype=self.single_expert_wgrad_desc.dtype, + device=self.single_expert_wgrad_desc.device, + ) + cached_single_expert = from_dlpack( + single_expert_placeholder, + assumed_align=16, + enable_tvm_ffi=True, + ) + + def tensor_api( + a_tensor: torch.Tensor, + b_tensor: torch.Tensor, + sfa_tensor: torch.Tensor, + sfb_tensor: torch.Tensor, + wgrad_ptrs: torch.Tensor, + offsets_tensor: torch.Tensor, + stream: cuda.CUstream, + global_scale_a: Optional[torch.Tensor], + global_scale_b: Optional[torch.Tensor], + ) -> None: + compiled( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + wgrad_ptrs.data_ptr(), + offsets_tensor, + cached_workspace, + stream, + global_scale_a, + global_scale_b, + cached_single_expert, + ) + + self._compiled_kernel = tensor_api + + def execute( + self, + a_tensor: torch.Tensor, + b_tensor: torch.Tensor, + sfa_tensor: torch.Tensor, + sfb_tensor: torch.Tensor, + offsets_tensor: torch.Tensor, + wgrad_tensor: Optional[torch.Tensor] = None, + wgrad_ptrs: Optional[torch.Tensor] = None, + global_scale_a: Optional[torch.Tensor] = None, + global_scale_b: Optional[torch.Tensor] = None, + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + current_stream = self._get_default_stream(current_stream) + self._runtime_error_if(self._compiled_kernel is None, "Kernel not compiled; call compile() first") + + if self.weight_mode == MoEWeightMode.DENSE: + self._value_error_if(wgrad_tensor is None, "wgrad_tensor is required in dense mode") + self._compiled_kernel( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + wgrad_tensor, + offsets_tensor, + current_stream, + global_scale_a, + global_scale_b, + ) + return + + if wgrad_ptrs is None: + self._value_error_if(wgrad_tensor is None, "Provide wgrad_tensor or wgrad_ptrs in discrete mode") + self._value_error_if(wgrad_tensor.ndim != 3, f"wgrad_tensor must be rank-3, got {tuple(wgrad_tensor.shape)}") + self._value_error_if(not wgrad_tensor.is_cuda, f"wgrad_tensor must be a CUDA tensor, got {wgrad_tensor.device}") + if wgrad_tensor.shape[0] == 0: + wgrad_ptrs = torch.empty((0,), dtype=torch.int64, device=wgrad_tensor.device) + else: + expert_stride_bytes = wgrad_tensor.stride(0) * wgrad_tensor.element_size() + ptrs = [wgrad_tensor.data_ptr() + i * expert_stride_bytes for i in range(wgrad_tensor.shape[0])] + wgrad_ptrs = torch.tensor(ptrs, dtype=torch.int64, device=wgrad_tensor.device) + _require_pointer_tensor(wgrad_ptrs, "wgrad_ptrs", self.expert_cnt) + self._compiled_kernel( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + wgrad_ptrs, + offsets_tensor, + current_stream, + global_scale_a, + global_scale_b, + ) + + +__all__ = ["GroupedGemmWgradBlockScaledAPI"] diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py index 29b3d7328..79989bb72 100644 --- a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/api.py @@ -5,36 +5,58 @@ from __future__ import annotations -from typing import Optional, Tuple -import logging +from typing import Any, Optional, Tuple, overload +import os import torch -import cutlass -import cutlass.cute as cute from cuda.bindings import driver as cuda -from cutlass.cute.runtime import from_dlpack, make_fake_stream -from cudnn.api_base import APIBase, TensorDesc, TupleDict, ceil_div, is_power_of_2 -from cudnn.datatypes import _convert_to_cutlass_data_type +from cudnn.api_base import APIBase, TupleDict from cudnn.discrete_grouped_gemm.discrete_kernel_utils import _require_pointer_tensor -from .moe_blockscaled_grouped_gemm_wgrad import BlockScaledMoEGroupedGemmWgradKernel -from ..moe_utils import MoEWeightMode, WGradInputOrder - +from ..grouped_gemm_utils import ( + GroupedGemmBackend, + _torch_stream_context, + backend_cache_key, + select_grouped_gemm_backend, +) +from ..moe_utils import WGradInputOrder + +_BLOCK_SCALED_DTYPE_PAIRS = { + (dtype, dtype) + for dtype in ( + torch.float4_e2m1fn_x2, + torch.uint8, + torch.float8_e5m2, + torch.float8_e4m3fn, + ) +} -def _round_up(a: int, b: int) -> int: - return ceil_div(a, b) * b +_cache_of_GroupedGemmWgradSm100Objects = {} -def _normalize_input_order(input_order: WGradInputOrder | str) -> WGradInputOrder: - if isinstance(input_order, WGradInputOrder): - return input_order - return WGradInputOrder(input_order) +from ._bf16_api import GroupedGemmWgradBf16API +from ._blockscaled_api import GroupedGemmWgradBlockScaledAPI class GroupedGemmWgradSm100(APIBase): - """Unified grouped GEMM wgrad FE API for SM100+ GPUs.""" + """Stable public facade that selects the WGrad backend during support checking.""" + + # BF16 implementation + @overload + def __init__( + self, + sample_a: torch.Tensor, + sample_b: torch.Tensor, + sample_sfa: None, + sample_sfb: None, + sample_offsets: torch.Tensor, + *args: Any, + **kwargs: Any, + ) -> None: ... + # Block-scaled implementation + @overload def __init__( self, sample_a: torch.Tensor, @@ -42,6 +64,17 @@ def __init__( sample_sfa: torch.Tensor, sample_sfb: torch.Tensor, sample_offsets: torch.Tensor, + *args: Any, + **kwargs: Any, + ) -> None: ... + + def __init__( + self, + sample_a: torch.Tensor, + sample_b: torch.Tensor, + sample_sfa: Optional[torch.Tensor], + sample_sfb: Optional[torch.Tensor], + sample_offsets: torch.Tensor, sample_wgrad: Optional[torch.Tensor] = None, sample_wgrad_expert: Optional[torch.Tensor] = None, num_experts: Optional[int] = None, @@ -55,442 +88,82 @@ def __init__( sf_vec_size: int = 16, accumulate_on_output: bool = False, input_order: WGradInputOrder | str = WGradInputOrder.Tensor2D, - ): + ) -> None: super().__init__() - self._warn_experimental_api() - self.input_order = _normalize_input_order(input_order) - - if sample_wgrad is not None and num_experts is None: - self.weight_mode = MoEWeightMode.DENSE - elif sample_wgrad is None and num_experts is not None: - self.weight_mode = MoEWeightMode.DISCRETE - if wgrad_shape is None or wgrad_dtype is None: - raise ValueError("wgrad_shape and wgrad_dtype are required in discrete mode") - else: - raise ValueError("Provide either sample_wgrad for dense mode or " "(num_experts, wgrad_shape, wgrad_dtype) for discrete mode, but not both") - - self._interpret_uint8_as_fp4x2 = True - self.sample_a_tensor = sample_a if self._is_fp4x2(sample_a) else None - self.sample_b_tensor = sample_b if self._is_fp4x2(sample_b) else None - self.a_desc = self._make_tensor_desc(sample_a, name="sample_a") - self.b_desc = self._make_tensor_desc(sample_b, name="sample_b") - self.sfa_desc = self._make_tensor_desc(sample_sfa, name="sample_sfa") - self.sfb_desc = self._make_tensor_desc(sample_sfb, name="sample_sfb") - self.offsets_desc = self._make_tensor_desc(sample_offsets, name="sample_offsets") - self.global_scale_a_desc = self._make_tensor_desc(sample_global_scale_a, name="sample_global_scale_a") - self.global_scale_b_desc = self._make_tensor_desc(sample_global_scale_b, name="sample_global_scale_b") - self.sf_vec_size = sf_vec_size - tokens_sum_a = self.a_desc.shape[1] - tokens_sum_b = self.b_desc.shape[0] - self._value_error_if( - tokens_sum_a != tokens_sum_b, - f"sample_a and sample_b token dimensions must match, got {tokens_sum_a} and {tokens_sum_b}", - ) - self._offset_values = self._validate_offsets(sample_offsets, tokens_sum_a, name="sample_offsets") - self._scale_cols = _round_up(ceil_div(tokens_sum_a, self.sf_vec_size), 4) - - if self.weight_mode == MoEWeightMode.DENSE: - self.wgrad_desc = self._make_tensor_desc(sample_wgrad, name="sample_wgrad") - self.expert_cnt = self.wgrad_desc.shape[0] - self.wgrad_shape = self.wgrad_desc.shape[1:] - self.wgrad_dtype = self.wgrad_desc.dtype - self.single_expert_wgrad_desc = TensorDesc( - dtype=self.wgrad_desc.dtype, - shape=self.wgrad_desc.shape[1:], - stride=self.wgrad_desc.stride[1:], - stride_order=tuple(i for i, s in sorted(enumerate(self.wgrad_desc.stride[1:]), key=lambda x: x[1])), - device=self.wgrad_desc.device, - name="single_expert_wgrad", - ) - else: # MoEWeightMode.DISCRETE - self.expert_cnt = num_experts - self.wgrad_shape = tuple(wgrad_shape) - self.wgrad_dtype = wgrad_dtype - self.wgrad_desc = None - if sample_wgrad_expert is not None: - self.single_expert_wgrad_desc = self._make_tensor_desc( - sample_wgrad_expert, - name="sample_wgrad_expert", - ) - else: - self.single_expert_wgrad_desc = TensorDesc( - dtype=wgrad_dtype, - shape=self.wgrad_shape, - stride=(self.wgrad_shape[1], 1), - stride_order=(1, 0), - device=self.a_desc.device, - name="single_expert_wgrad", - ) - - self.acc_dtype = acc_dtype - self.mma_tiler_mn = mma_tiler_mn - self.use_2cta_instrs = mma_tiler_mn[0] == 256 - self.cluster_shape_mn = cluster_shape_mn or ((2, 1) if self.use_2cta_instrs else (1, 1)) - self.accumulate_on_output = accumulate_on_output - self._kernel = BlockScaledMoEGroupedGemmWgradKernel - self._workspace = None - - def _validate_offsets(self, offsets_tensor: torch.Tensor, tokens_sum: int, name: str) -> Tuple[int, ...]: - self._value_error_if(offsets_tensor.ndim != 1, f"{name} must be rank-1, got shape {tuple(offsets_tensor.shape)}") - - offset_values = tuple(int(offset) for offset in offsets_tensor.detach().cpu().tolist()) - prev_offset = 0 - for idx, offset in enumerate(offset_values): - self._value_error_if( - offset < prev_offset, - f"{name} must be a non-decreasing cumulative sum, but index {idx} has {offset} after {prev_offset}", - ) - prev_offset = offset - - if offset_values: - self._value_error_if( - offset_values[-1] > tokens_sum, - f"{name} last value must not exceed total tokens {tokens_sum}, got {offset_values[-1]}", - ) - else: - self._value_error_if(tokens_sum != 0, f"{name} cannot be empty when total tokens is {tokens_sum}") - - return offset_values + self._pending_init_kwargs = dict(locals()) + self._pending_init_kwargs.pop("self") + self._pending_init_kwargs.pop("__class__", None) + self._implementation = None def check_support(self) -> bool: - m, tokens_sum = self._tensor_shape(self.a_desc, name="sample_a") - _, n = self._tensor_shape(self.b_desc, name="sample_b") - - _ = self._check_tensor_shape(self.a_desc, (m, tokens_sum), "sample_a") - _ = self._check_tensor_shape(self.b_desc, (tokens_sum, n), "sample_b") - _ = self._check_tensor_shape(self.sfa_desc, (_round_up(m, 128), self._scale_cols), "sample_sfa") - _ = self._check_tensor_shape(self.sfb_desc, (_round_up(n, 128), self._scale_cols), "sample_sfb") - _ = self._check_tensor_shape(self.offsets_desc, (self.expert_cnt,), "sample_offsets") - - dtype = self._check_dtype( - self.a_desc, [torch.float4_e2m1fn_x2, torch.uint8, torch.float8_e5m2, torch.float8_e4m3fn, torch.bfloat16], "sample_a" - ) # TODO @mingyangw: check if bfloat16 is supported - self._check_dtype(self.b_desc, dtype, "sample_b", extra_error_msg="sample_b must have the same dtype as sample_a") - self._check_dtype( - self.sfa_desc, - [torch.float8_e8m0fnu, torch.float8_e4m3fn], - "sample_sfa", - extra_error_msg="sample_sfa must have dtype float8_e8m0fnu or float8_e4m3fn", - ) - self._check_dtype( - self.sfb_desc, - [torch.float8_e8m0fnu, torch.float8_e4m3fn], - "sample_sfb", - extra_error_msg="sample_sfb must have dtype float8_e8m0fnu or float8_e4m3fn", - ) - self._check_dtype(self.offsets_desc, torch.int32, "sample_offsets", extra_error_msg="sample_offsets must be int32") - self._check_dtype( - self.wgrad_dtype, [torch.bfloat16, torch.float16, torch.float32], "wgrad_dtype", extra_error_msg="wgrad_dtype must be bfloat16, float16, or float32" - ) - - if self.weight_mode == MoEWeightMode.DENSE: - self._check_tensor_shape(self.wgrad_desc, (self.expert_cnt, m, n), "sample_wgrad") - else: - self._check_tensor_shape(self.wgrad_shape, (m, n), "wgrad_shape") - self._check_tensor_shape(self.single_expert_wgrad_desc, (m, n), "single_expert_wgrad") - self._check_dtype( - self.single_expert_wgrad_desc, - self.wgrad_dtype, - "sample_wgrad_expert", - extra_error_msg="sample_wgrad_expert must have the same dtype as wgrad_dtype", - ) - - self._value_error_if(self.mma_tiler_mn[0] not in (128, 256), f"mma_tiler_mn[0] must be 128 or 256, got {self.mma_tiler_mn[0]}") - self._value_error_if(self.mma_tiler_mn[1] not in (128, 256), f"mma_tiler_mn[1] must be 128 or 256, got {self.mma_tiler_mn[1]}") - self._value_error_if( - self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0, - f"cluster_shape_mn[0] must be divisible by 2 when use_2cta_instrs=True, got {self.cluster_shape_mn[0]}", - ) - self._value_error_if(self.cluster_shape_mn[0] * self.cluster_shape_mn[1] > 16, f"cluster shape product must be <= 16, got {self.cluster_shape_mn}") - self._value_error_if( - not (is_power_of_2(self.cluster_shape_mn[0]) and is_power_of_2(self.cluster_shape_mn[1])), - f"cluster shape values must be powers of 2, got {self.cluster_shape_mn}", - ) - - has_global_scale = self.global_scale_a_desc is not None or self.global_scale_b_desc is not None - if has_global_scale: - self._value_error_if( - self.global_scale_a_desc is None or self.global_scale_b_desc is None, - "sample_global_scale_a and sample_global_scale_b must be provided together", - ) - self._value_error_if( - self.global_scale_a_desc.shape != (self.expert_cnt,), - f"sample_global_scale_a must have shape {(self.expert_cnt,)}, got {self.global_scale_a_desc.shape}", + if self._implementation is None: + kwargs = self._pending_init_kwargs + backend = select_grouped_gemm_backend( + operation="grouped_gemm_wgrad_sm100", + a_dtype=kwargs["sample_a"].dtype, + b_dtype=kwargs["sample_b"].dtype, + scale_controls=( + ("sample_sfa", kwargs["sample_sfa"]), + ("sample_sfb", kwargs["sample_sfb"]), + ("sample_global_scale_a", kwargs["sample_global_scale_a"]), + ("sample_global_scale_b", kwargs["sample_global_scale_b"]), + ("sf_vec_size", kwargs["sf_vec_size"] if kwargs["sf_vec_size"] != 16 else None), + ), + block_scaled_dtype_pairs=_BLOCK_SCALED_DTYPE_PAIRS, ) - self._value_error_if( - self.global_scale_b_desc.shape != (self.expert_cnt,), - f"sample_global_scale_b must have shape {(self.expert_cnt,)}, got {self.global_scale_b_desc.shape}", - ) - self._check_dtype(self.global_scale_a_desc, torch.float32, "sample_global_scale_a") - self._check_dtype(self.global_scale_b_desc, torch.float32, "sample_global_scale_b") - - requires_global_scale = ( - self._is_fp4x2(self.a_desc) and self.sf_vec_size == 16 and self.sfa_desc.dtype == torch.float8_e4m3fn and self.sfb_desc.dtype == torch.float8_e4m3fn - ) - self._value_error_if(requires_global_scale and not has_global_scale, "NVFP4 wgrad requires sample_global_scale_a and sample_global_scale_b") - - if not torch.cuda.is_available(): - raise RuntimeError("CUDA is not available") - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) - compute_capability = major * 10 + minor - if compute_capability < 100: - raise RuntimeError(f"GroupedGemmWgrad requires SM100+ compute capability, but found SM{compute_capability} on device {device}") - - self._is_supported = True - return True + self.backend = backend + if backend is GroupedGemmBackend.BF16: + self._implementation = GroupedGemmWgradBf16API(**kwargs) + else: + self._implementation = GroupedGemmWgradBlockScaledAPI(**kwargs) + self._kernel = self._implementation._kernel + self.weight_mode = self._implementation.weight_mode + supported = self._implementation.check_support() + self._is_supported = self._implementation._is_supported + if supported: + self._pending_init_kwargs = None + return supported def compile(self) -> None: - self._ensure_support_checked() - if self._compiled_kernel is not None: - return - - kernel = self._kernel( - sf_vec_size=self.sf_vec_size, - acc_dtype=_convert_to_cutlass_data_type(self.acc_dtype), - use_2cta_instrs=self.use_2cta_instrs, - mma_tiler_mn=self.mma_tiler_mn, - cluster_shape_mn=self.cluster_shape_mn, - accumulate_on_output=self.accumulate_on_output, - expert_cnt=self.expert_cnt, - weight_mode=self.weight_mode, - input_order=self.input_order, - ) - - hardware_info = cutlass.utils.HardwareInfo() - max_active_clusters = hardware_info.get_max_active_clusters(self.cluster_shape_mn[0] * self.cluster_shape_mn[1]) - self._workspace = torch.empty(max(kernel.get_workspace_bytes(), 1), dtype=torch.uint8, device=self.a_desc.device) - fake_stream = make_fake_stream(use_tvm_ffi_env_stream=False) - - if self.weight_mode == MoEWeightMode.DENSE: - self._compile_dense(kernel, max_active_clusters, fake_stream) - else: - self._compile_discrete(kernel, max_active_clusters, fake_stream) - - if self.sample_a_tensor is not None: - del self.sample_a_tensor - if self.sample_b_tensor is not None: - del self.sample_b_tensor - - def _compile_dense(self, kernel, max_active_clusters, fake_stream) -> None: - a_fake = ( - from_dlpack(self.sample_a_tensor, assumed_align=16, enable_tvm_ffi=True).mark_compact_shape_dynamic( - mode=1, - stride_order=self.sample_a_tensor.dim_order(), - divisibility=16, - ) - if self.sample_a_tensor is not None - else self._make_fake_cute_compact_tensor( - dtype=self.a_desc.dtype, - shape=self.a_desc.shape, - stride_order=self.a_desc.stride_order, - assumed_align=16, - dynamic_mode=1, - divisibility=16, - ) - ) - b_fake = ( - from_dlpack(self.sample_b_tensor, assumed_align=16, enable_tvm_ffi=True).mark_compact_shape_dynamic( - mode=0, - stride_order=self.sample_b_tensor.dim_order(), - divisibility=16, - ) - if self.sample_b_tensor is not None - else self._make_fake_cute_compact_tensor( - dtype=self.b_desc.dtype, - shape=self.b_desc.shape, - stride_order=self.b_desc.stride_order, - assumed_align=16, - dynamic_mode=0, - divisibility=16, - ) - ) - sfa_fake = self._make_fake_cute_compact_tensor( - dtype=self.sfa_desc.dtype, - shape=self.sfa_desc.shape, - stride_order=self.sfa_desc.stride_order, - assumed_align=16, - dynamic_mode=1, - divisibility=4, - ) - sfb_fake = self._make_fake_cute_compact_tensor( - dtype=self.sfb_desc.dtype, - shape=self.sfb_desc.shape, - stride_order=self.sfb_desc.stride_order, - assumed_align=16, - dynamic_mode=1, - divisibility=4, - ) - wgrad_fake = self._make_fake_cute_tensor_from_desc(self.wgrad_desc, assumed_align=16) - offsets_fake = self._make_fake_cute_tensor_from_desc(self.offsets_desc, assumed_align=4) - workspace_fake = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) - gs_a_fake = self._make_fake_cute_tensor_from_desc(self.global_scale_a_desc, assumed_align=4) - gs_b_fake = self._make_fake_cute_tensor_from_desc(self.global_scale_b_desc, assumed_align=4) - - compiled = cute.compile( - kernel, - a_fake, - b_fake, - sfa_fake, - sfb_fake, - wgrad_fake, - offsets_fake, - workspace_fake, - max_active_clusters, - fake_stream, - gs_a_fake, - gs_b_fake, - None, - options="--enable-tvm-ffi", - ) - - cached_workspace = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) - - def tensor_api( - a_tensor: torch.Tensor, - b_tensor: torch.Tensor, - sfa_tensor: torch.Tensor, - sfb_tensor: torch.Tensor, - wgrad_tensor: torch.Tensor, - offsets_tensor: torch.Tensor, - stream: cuda.CUstream, - global_scale_a: Optional[torch.Tensor], - global_scale_b: Optional[torch.Tensor], - ) -> None: - compiled( - a_tensor, - b_tensor, - sfa_tensor, - sfb_tensor, - wgrad_tensor, - offsets_tensor, - cached_workspace, - stream, - global_scale_a, - global_scale_b, - None, - ) - - self._compiled_kernel = tensor_api - - def _compile_discrete(self, kernel, max_active_clusters, fake_stream) -> None: - a_fake = ( - from_dlpack(self.sample_a_tensor, assumed_align=16, enable_tvm_ffi=True).mark_compact_shape_dynamic( - mode=1, - stride_order=self.sample_a_tensor.dim_order(), - divisibility=16, - ) - if self.sample_a_tensor is not None - else self._make_fake_cute_compact_tensor( - dtype=self.a_desc.dtype, - shape=self.a_desc.shape, - stride_order=self.a_desc.stride_order, - assumed_align=16, - dynamic_mode=1, - divisibility=16, - ) - ) - b_fake = ( - from_dlpack(self.sample_b_tensor, assumed_align=16, enable_tvm_ffi=True).mark_compact_shape_dynamic( - mode=0, - stride_order=self.sample_b_tensor.dim_order(), - divisibility=16, - ) - if self.sample_b_tensor is not None - else self._make_fake_cute_compact_tensor( - dtype=self.b_desc.dtype, - shape=self.b_desc.shape, - stride_order=self.b_desc.stride_order, - assumed_align=16, - dynamic_mode=0, - divisibility=16, - ) - ) - sfa_fake = self._make_fake_cute_compact_tensor( - dtype=self.sfa_desc.dtype, - shape=self.sfa_desc.shape, - stride_order=self.sfa_desc.stride_order, - assumed_align=16, - dynamic_mode=1, - divisibility=4, - ) - sfb_fake = self._make_fake_cute_compact_tensor( - dtype=self.sfb_desc.dtype, - shape=self.sfb_desc.shape, - stride_order=self.sfb_desc.stride_order, - assumed_align=16, - dynamic_mode=1, - divisibility=4, - ) - offsets_fake = self._make_fake_cute_tensor_from_desc(self.offsets_desc, assumed_align=4) - workspace_fake = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) - gs_a_fake = self._make_fake_cute_tensor_from_desc(self.global_scale_a_desc, assumed_align=4) - gs_b_fake = self._make_fake_cute_tensor_from_desc(self.global_scale_b_desc, assumed_align=4) - wgrad_ptrs_placeholder = torch.empty((self.expert_cnt,), dtype=torch.int64, device=self.a_desc.device) - wgrad_ptrs_fake = from_dlpack(wgrad_ptrs_placeholder, assumed_align=8, enable_tvm_ffi=True).iterator - single_expert_fake = self._make_fake_cute_tensor( - dtype=self.single_expert_wgrad_desc.dtype, - shape=self.single_expert_wgrad_desc.shape, - stride=self.single_expert_wgrad_desc.stride, - assumed_align=16, - ) - - compiled = cute.compile( - kernel, - a_fake, - b_fake, - sfa_fake, - sfb_fake, - wgrad_ptrs_fake, - offsets_fake, - workspace_fake, - max_active_clusters, - fake_stream, - gs_a_fake, - gs_b_fake, - single_expert_fake, - options="--enable-tvm-ffi", - ) - - cached_workspace = from_dlpack(self._workspace, assumed_align=128, enable_tvm_ffi=True) - single_expert_placeholder = torch.empty_strided( - self.single_expert_wgrad_desc.shape, - self.single_expert_wgrad_desc.stride, - dtype=self.single_expert_wgrad_desc.dtype, - device=self.single_expert_wgrad_desc.device, - ) - cached_single_expert = from_dlpack( - single_expert_placeholder, - assumed_align=16, - enable_tvm_ffi=True, - ) - - def tensor_api( - a_tensor: torch.Tensor, - b_tensor: torch.Tensor, - sfa_tensor: torch.Tensor, - sfb_tensor: torch.Tensor, - wgrad_ptrs: torch.Tensor, - offsets_tensor: torch.Tensor, - stream: cuda.CUstream, - global_scale_a: Optional[torch.Tensor], - global_scale_b: Optional[torch.Tensor], - ) -> None: - compiled( - a_tensor, - b_tensor, - sfa_tensor, - sfb_tensor, - wgrad_ptrs.data_ptr(), - offsets_tensor, - cached_workspace, - stream, - global_scale_a, - global_scale_b, - cached_single_expert, - ) + if self._implementation is None: + self.check_support() + if self._is_supported: + self._implementation._is_supported = True + self._implementation.compile() + self._is_supported = self._implementation._is_supported + self._compiled_kernel = self._implementation._compiled_kernel + + # BF16 implementation + @overload + def execute( + self, + a_tensor: torch.Tensor, + b_tensor: torch.Tensor, + sfa_tensor: None, + sfb_tensor: None, + offsets_tensor: torch.Tensor, + wgrad_tensor: Optional[torch.Tensor] = None, + wgrad_ptrs: Optional[torch.Tensor] = None, + *, + global_scale_a: None = None, + global_scale_b: None = None, + ) -> None: ... - self._compiled_kernel = tensor_api + # Block-scaled implementation + @overload + def execute( + self, + a_tensor: torch.Tensor, + b_tensor: torch.Tensor, + sfa_tensor: torch.Tensor, + sfb_tensor: torch.Tensor, + offsets_tensor: torch.Tensor, + wgrad_tensor: Optional[torch.Tensor] = None, + wgrad_ptrs: Optional[torch.Tensor] = None, + *, + global_scale_a: Optional[torch.Tensor] = None, + global_scale_b: Optional[torch.Tensor] = None, + ) -> None: ... def execute( self, @@ -505,64 +178,29 @@ def execute( global_scale_b: Optional[torch.Tensor] = None, current_stream: Optional[cuda.CUstream] = None, ) -> None: - current_stream = self._get_default_stream(current_stream) - self._runtime_error_if(self._compiled_kernel is None, "Kernel not compiled; call compile() first") - - if self.weight_mode == MoEWeightMode.DENSE: - self._value_error_if(wgrad_tensor is None, "wgrad_tensor is required in dense mode") - self._compiled_kernel( - a_tensor, - b_tensor, - sfa_tensor, - sfb_tensor, - wgrad_tensor, - offsets_tensor, - current_stream, - global_scale_a, - global_scale_b, - ) - return - - if wgrad_ptrs is None: - self._value_error_if(wgrad_tensor is None, "Provide wgrad_tensor or wgrad_ptrs in discrete mode") - self._value_error_if(wgrad_tensor.ndim != 3, f"wgrad_tensor must be rank-3, got {tuple(wgrad_tensor.shape)}") - self._value_error_if(not wgrad_tensor.is_cuda, f"wgrad_tensor must be a CUDA tensor, got {wgrad_tensor.device}") - if wgrad_tensor.shape[0] == 0: - wgrad_ptrs = torch.empty((0,), dtype=torch.int64, device=wgrad_tensor.device) - else: - expert_stride_bytes = wgrad_tensor.stride(0) * wgrad_tensor.element_size() - ptrs = [wgrad_tensor.data_ptr() + i * expert_stride_bytes for i in range(wgrad_tensor.shape[0])] - wgrad_ptrs = torch.tensor(ptrs, dtype=torch.int64, device=wgrad_tensor.device) - _require_pointer_tensor(wgrad_ptrs, "wgrad_ptrs", self.expert_cnt) - self._compiled_kernel( - a_tensor, - b_tensor, - sfa_tensor, - sfb_tensor, - wgrad_ptrs, - offsets_tensor, - current_stream, - global_scale_a, - global_scale_b, + if self._implementation is None: + raise RuntimeError("Kernel not compiled; call compile() first") + self._implementation.execute( + a_tensor=a_tensor, + b_tensor=b_tensor, + sfa_tensor=sfa_tensor, + sfb_tensor=sfb_tensor, + offsets_tensor=offsets_tensor, + wgrad_tensor=wgrad_tensor, + wgrad_ptrs=wgrad_ptrs, + global_scale_a=global_scale_a, + global_scale_b=global_scale_b, + current_stream=current_stream, ) -_logger = logging.getLogger(__name__) -_cache_of_GroupedGemmWgradSm100Objects = {} - - -def _stride_order(tensor: torch.Tensor) -> Tuple[int, ...]: - return tuple(i for i, s in sorted(enumerate(tensor.stride()), key=lambda x: (x[1], tensor.shape[x[0]]))) - - -def _dynamic_dim_tensor_signature( - tensor: Optional[torch.Tensor], - dynamic_dims: Tuple[int, ...], -) -> Tuple[Optional[Tuple[Optional[int], ...]], Optional[Tuple[int, ...]], Optional[torch.dtype]]: +def _wgrad_tensor_signature(tensor: Optional[torch.Tensor], *, dynamic_dims: tuple[int, ...] = (), exact_stride: bool): if tensor is None: - return None, None, None - static_shape = tuple(None if i in dynamic_dims else int(dim) for i, dim in enumerate(tensor.shape)) - return static_shape, _stride_order(tensor), tensor.dtype + return None + shape = tuple(None if index in dynamic_dims else int(value) for index, value in enumerate(tensor.shape)) + stride = tuple(int(value) for value in tensor.stride()) + layout = stride if exact_stride else tuple(index for index, _ in sorted(enumerate(stride), key=lambda item: (item[1], tensor.shape[item[0]]))) + return (shape, layout, tensor.dtype, tensor.device) def grouped_gemm_wgrad_wrapper_sm100( @@ -585,91 +223,94 @@ def grouped_gemm_wgrad_wrapper_sm100( input_order: WGradInputOrder | str = WGradInputOrder.Tensor2D, current_stream: Optional[cuda.CUstream] = None, ) -> TupleDict: - """Compile and execute grouped GEMM wgrad in one call.""" - input_order = _normalize_input_order(input_order) - hidden, _ = a_tensor.shape - _, intermediate = b_tensor.shape - wgrad_shape = (hidden, intermediate) - expert_cnt = offsets_tensor.shape[0] - - if output_mode not in {"dense", "discrete"}: - raise ValueError(f"output_mode must be 'dense' or 'discrete', got {output_mode}") - + """Compile and execute grouped GEMM wgrad through the selected backend API.""" + if output_mode not in ("dense", "discrete"): + raise ValueError(f'output_mode must be "dense" or "discrete", got {output_mode}') + if a_tensor.ndim != 2 or b_tensor.ndim != 2: + raise ValueError("a_tensor and b_tensor must both be rank-2") + hidden, tokens_sum = a_tensor.shape + tokens_b, intermediate = b_tensor.shape + if tokens_sum != tokens_b: + raise ValueError(f"a_tensor and b_tensor token dimensions must match, got {tokens_sum} and {tokens_b}") + if offsets_tensor.ndim != 1: + raise ValueError(f"offsets_tensor must be rank-1, got shape {tuple(offsets_tensor.shape)}") + input_order = WGradInputOrder(input_order) + expert_cnt = offsets_tensor.numel() + if output_mode == "dense" and wgrad_ptrs is not None: + raise ValueError("dense output_mode forbids wgrad_ptrs") + if wgrad_ptrs is not None: + _require_pointer_tensor(wgrad_ptrs, "wgrad_ptrs", expert_cnt) + backend = select_grouped_gemm_backend( + operation="grouped_gemm_wgrad_sm100", + a_dtype=a_tensor.dtype, + b_dtype=b_tensor.dtype, + scale_controls=( + ("sfa_tensor", sfa_tensor), + ("sfb_tensor", sfb_tensor), + ("global_scale_a", global_scale_a), + ("global_scale_b", global_scale_b), + ("sf_vec_size", sf_vec_size if sf_vec_size != 16 else None), + ), + block_scaled_dtype_pairs=_BLOCK_SCALED_DTYPE_PAIRS, + ) if wgrad_tensor is None and wgrad_ptrs is None: - # Backward compatibility: Dense mode. - if accumulate_on_output: - wgrad_tensor = torch.zeros((expert_cnt, *wgrad_shape), dtype=wgrad_dtype, device=a_tensor.device) - else: - wgrad_tensor = torch.empty((expert_cnt, *wgrad_shape), dtype=wgrad_dtype, device=a_tensor.device) - - cache_key = ( + allocator = torch.zeros if accumulate_on_output else torch.empty + with _torch_stream_context(current_stream, a_tensor.device): + wgrad_tensor = allocator((expert_cnt, hidden, intermediate), dtype=wgrad_dtype, device=a_tensor.device) + cache_key = backend_cache_key( + backend, output_mode, - *_dynamic_dim_tensor_signature(a_tensor, dynamic_dims=(1,)), - *_dynamic_dim_tensor_signature(b_tensor, dynamic_dims=(0,)), - *_dynamic_dim_tensor_signature(sfa_tensor, dynamic_dims=(1,)), - *_dynamic_dim_tensor_signature(sfb_tensor, dynamic_dims=(1,)), - tuple(offsets_tensor.shape), - tuple(offsets_tensor.stride()), - offsets_tensor.dtype, - *_dynamic_dim_tensor_signature(wgrad_tensor, dynamic_dims=()), - tuple(global_scale_a.shape) if global_scale_a is not None else None, - global_scale_a.dtype if global_scale_a is not None else None, - tuple(global_scale_b.shape) if global_scale_b is not None else None, - global_scale_b.dtype if global_scale_b is not None else None, + _wgrad_tensor_signature(a_tensor, dynamic_dims=(1,), exact_stride=False), + _wgrad_tensor_signature(b_tensor, dynamic_dims=(0,), exact_stride=False), + _wgrad_tensor_signature(sfa_tensor, dynamic_dims=(1,), exact_stride=False), + _wgrad_tensor_signature(sfb_tensor, dynamic_dims=(1,), exact_stride=False), + _wgrad_tensor_signature(offsets_tensor, exact_stride=True), + _wgrad_tensor_signature(wgrad_tensor, exact_stride=True), + _wgrad_tensor_signature(wgrad_ptrs, exact_stride=True), + _wgrad_tensor_signature(global_scale_a, exact_stride=True), + _wgrad_tensor_signature(global_scale_b, exact_stride=True), acc_dtype, wgrad_dtype, - mma_tiler_mn, - cluster_shape_mn, + tuple(mma_tiler_mn), + tuple(cluster_shape_mn) if cluster_shape_mn is not None else None, sf_vec_size, accumulate_on_output, input_order, + int(os.getenv("CUDNNFE_CLUSTER_OVERLAP_MARGIN", "0")), ) - - if cache_key in _cache_of_GroupedGemmWgradSm100Objects: - op = _cache_of_GroupedGemmWgradSm100Objects[cache_key] - else: + op = _cache_of_GroupedGemmWgradSm100Objects.get(cache_key) + if op is None: + common = dict( + sample_a=a_tensor, + sample_b=b_tensor, + sample_sfa=sfa_tensor, + sample_sfb=sfb_tensor, + sample_offsets=offsets_tensor, + sample_global_scale_a=global_scale_a, + sample_global_scale_b=global_scale_b, + acc_dtype=acc_dtype, + mma_tiler_mn=mma_tiler_mn, + cluster_shape_mn=cluster_shape_mn, + sf_vec_size=sf_vec_size, + accumulate_on_output=accumulate_on_output, + input_order=input_order, + ) if output_mode == "dense": - op = GroupedGemmWgradSm100( - sample_a=a_tensor, - sample_b=b_tensor, - sample_sfa=sfa_tensor, - sample_sfb=sfb_tensor, - sample_offsets=offsets_tensor, - sample_wgrad=wgrad_tensor, - sample_global_scale_a=global_scale_a, - sample_global_scale_b=global_scale_b, - acc_dtype=acc_dtype, - mma_tiler_mn=mma_tiler_mn, - cluster_shape_mn=cluster_shape_mn, - sf_vec_size=sf_vec_size, - accumulate_on_output=accumulate_on_output, - input_order=input_order, - ) + common["sample_wgrad"] = wgrad_tensor else: - sample_expert = torch.empty(wgrad_shape, dtype=wgrad_dtype, device=a_tensor.device) - op = GroupedGemmWgradSm100( - sample_a=a_tensor, - sample_b=b_tensor, - sample_sfa=sfa_tensor, - sample_sfb=sfb_tensor, - sample_offsets=offsets_tensor, - sample_wgrad_expert=sample_expert, + common.update( + sample_wgrad_expert=( + wgrad_tensor[0] if wgrad_tensor is not None else torch.empty((hidden, intermediate), dtype=wgrad_dtype, device=a_tensor.device) + ), num_experts=expert_cnt, wgrad_shape=(hidden, intermediate), wgrad_dtype=wgrad_dtype, - sample_global_scale_a=global_scale_a, - sample_global_scale_b=global_scale_b, - acc_dtype=acc_dtype, - mma_tiler_mn=mma_tiler_mn, - cluster_shape_mn=cluster_shape_mn, - sf_vec_size=sf_vec_size, - accumulate_on_output=accumulate_on_output, - input_order=input_order, ) - assert op.check_support(), "Unsupported configuration" + op = GroupedGemmWgradSm100(**common) + if not op.check_support(): + raise RuntimeError("Unsupported configuration") op.compile() _cache_of_GroupedGemmWgradSm100Objects[cache_key] = op - op.execute( a_tensor=a_tensor, b_tensor=b_tensor, @@ -682,5 +323,4 @@ def grouped_gemm_wgrad_wrapper_sm100( global_scale_b=global_scale_b, current_stream=current_stream, ) - return TupleDict(wgrad_tensor=wgrad_tensor) diff --git a/python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_grouped_gemm_wgrad.py b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_grouped_gemm_wgrad.py new file mode 100644 index 000000000..14c4c471e --- /dev/null +++ b/python/cudnn/grouped_gemm/grouped_gemm_wgrad/moe_grouped_gemm_wgrad.py @@ -0,0 +1,1154 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +""" +MoE BF16 Grouped GEMM Kernel — Weight Gradient (2Dx2D). + +Computes: A(hidden, tokens_sum) x B(tokens_sum, intermediate) + -> C(experts, hidden, intermediate) + +where C is the weight gradient. K (tokens) varies per expert; +M (hidden) and N (intermediate) are fixed across all experts. + +Supports: + - CLC-based dynamic persistent tile scheduling + - Dense (contiguous 3-D C) / Discrete (per-expert pointer array C) output + - accumulate_on_output (TMA reduce for atomic accumulation) + - k_tile_cnt == 0 handling (zero output for empty experts) + +This module contains only the kernel class. +Scheduler: moe_persistent_scheduler.py (CLC mode, scenario="2Dx2D") +Extension: moe_sched_extension.py (WgradDense / WgradDiscrete) +""" + +from typing import Type, Tuple, Optional + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.utils.gemm.sm100 import ( + transform_partitioned_tensor_layout, + epilogue_tmem_copy_and_partition, + epilogue_smem_copy_and_partition, +) +from ..moe_persistent_scheduler import ( + MoEPersistentTileScheduler, + MoESchedulerParams, + MoEWorkTileInfo, +) +from ..moe_utils import ( + MoEWeightMode, + WGradInputOrder, + WgradTensormapConstructor, +) +from ..moe_sched_extension import ( + WgradGemmSchedExtension, +) +from ..moe_kernel_helpers import ( + compute_stages_wgrad_bf16, +) + + +class MoEGroupedGemmWgradBF16Kernel: + """BF16 grouped GEMM kernel for MoE weight gradient (2Dx2D). + + :param acc_dtype: Accumulator data type (Float32). + :param use_2cta_instrs: Use 2-CTA MMA instructions. + :param mma_tiler_mn: MMA tile shape (M, N). + :param cluster_shape_mn: Cluster shape (M, N). + :param accumulate_on_output: Use TMA reduce for atomic accumulation. + :param expert_cnt: Number of experts. + :param weight_mode: ``MoEWeightMode.DENSE`` or ``MoEWeightMode.DISCRETE`` for output. + :param input_order: ``WGradInputOrder.Tensor2D`` (default, single global K + stride across experts) or ``WGradInputOrder.TensorRagged`` (per-expert + K-contiguous blocks concatenated in memory). Only the unit-stride axis + of host-side ``mat_a/mat_b`` is meaningful in ``TensorRagged`` mode; + the non-unit stride values are ignored. + """ + + FIX_PAD_SIZE = 256 + + @staticmethod + def can_implement( + ab_dtype: Type[cutlass.Numeric], + out_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m: int, + n: int, + group_k_list: list, + expert_cnt: int, + a_major: str, + b_major: str, + weight_mode: MoEWeightMode, + input_order: WGradInputOrder, + ) -> bool: + """Check whether a BF16 wgrad testcase is supported.""" + result = True + + if ab_dtype != cutlass.BFloat16: + result = False + if out_dtype not in (cutlass.BFloat16, cutlass.Float16, cutlass.Float32): + result = False + if acc_dtype != cutlass.Float32: + result = False + + if a_major not in ("k", "m") or b_major not in ("k", "n"): + result = False + if weight_mode not in (MoEWeightMode.DENSE, MoEWeightMode.DISCRETE): + result = False + if input_order not in (WGradInputOrder.Tensor2D, WGradInputOrder.TensorRagged): + result = False + + if m <= 0 or n <= 0 or expert_cnt <= 0: + result = False + if group_k_list is None: + return False + if len(group_k_list) != expert_cnt: + result = False + + tokens_sum = 0 + for k_val in group_k_list: + if k_val < 0 or k_val % MoEGroupedGemmWgradBF16Kernel.FIX_PAD_SIZE != 0: + result = False + tokens_sum += k_val + + if len(mma_tiler_mn) != 2 or len(cluster_shape_mn) != 2: + return False + + tile_m, tile_n = mma_tiler_mn + cluster_m, cluster_n = cluster_shape_mn + if not ((not use_2cta_instrs and tile_m == 128) or (use_2cta_instrs and tile_m == 256)): + result = False + if tile_n not in range(64, 257, 64): + result = False + if cluster_m % (2 if use_2cta_instrs else 1) != 0: + result = False + + def is_power_of_2(x: int) -> bool: + return x > 0 and (x & (x - 1)) == 0 + + if cluster_m * cluster_n > 16 or not is_power_of_2(cluster_m) or not is_power_of_2(cluster_n): + result = False + + def is_16b_aligned(dtype: Type[cutlass.Numeric], contiguous_elems: int) -> bool: + elems_per_16b = 16 * 8 // dtype.width + return contiguous_elems % elems_per_16b == 0 + + a_contiguous_elems = tokens_sum if a_major == "k" else m + b_contiguous_elems = tokens_sum if b_major == "k" else n + c_contiguous_elems = n + if not ( + is_16b_aligned(ab_dtype, a_contiguous_elems) and is_16b_aligned(ab_dtype, b_contiguous_elems) and is_16b_aligned(out_dtype, c_contiguous_elems) + ): + result = False + + return result + + def __init__( + self, + acc_dtype: Type[cutlass.Numeric] = cutlass.Float32, + use_2cta_instrs: bool = False, + mma_tiler_mn: Tuple[int, int] = (128, 128), + cluster_shape_mn: Tuple[int, int] = (1, 1), + accumulate_on_output: bool = False, + expert_cnt: int = 1, + weight_mode: MoEWeightMode = MoEWeightMode.DENSE, + input_order: WGradInputOrder = WGradInputOrder.Tensor2D, + ): + self.expert_cnt = expert_cnt + self.acc_dtype = acc_dtype + self.use_2cta_instrs = use_2cta_instrs + self.cluster_shape_mn = cluster_shape_mn + self.mma_tiler = (*mma_tiler_mn, 1) + self.accumulate_on_output = accumulate_on_output + self.weight_mode = weight_mode + self.input_order = input_order + + self.cta_group = tcgen05.CtaGroup.TWO if use_2cta_instrs else tcgen05.CtaGroup.ONE + + self.occupancy = 1 + self.epilogue_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.sched_warp_id = 6 + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_warp_id, + self.sched_warp_id, + *self.epilogue_warp_id, + ) + ) + + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + + self.smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + + # ------------------------------------------------------------------ + # Workspace + # ------------------------------------------------------------------ + + def get_workspace_bytes(self) -> int: + return WgradTensormapConstructor.get_workspace_size(self.input_order, self.weight_mode, self.expert_cnt) + + # ------------------------------------------------------------------ + # _setup_attributes + # ------------------------------------------------------------------ + + def _setup_attributes(self) -> None: + self.mma_inst_shape_mn = (self.mma_tiler[0], self.mma_tiler[1]) + + tiled_mma = self._create_tiled_mma() + + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 + mma_tiler_k = mma_inst_shape_k * mma_inst_tile_k + self.mma_tiler = ( + self.mma_inst_shape_mn[0], + self.mma_inst_shape_mn[1], + mma_tiler_k, + ) + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + self.epi_tile = sm100_utils.compute_epilogue_tile_shape( + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + self.epi_tile_n = cute.size(self.epi_tile[1]) + + self.num_acc_stage, self.num_ab_stage, self.num_c_stage = compute_stages_wgrad_bf16( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.epi_tile, + self.c_dtype, + self.c_layout, + self.smem_capacity, + self.occupancy, + ) + self.num_sched_stages = 2 + + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + + self.num_acc_pipeline_stages = self.num_acc_stage + self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * self.num_acc_stage + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + self.num_tma_load_bytes = (a_copy_size + b_copy_size) * atom_thr_size + + # ------------------------------------------------------------------ + # MMA helpers + # ------------------------------------------------------------------ + + def _create_tiled_mma(self): + return sm100_utils.make_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_inst_shape_mn, + ) + + @cute.jit + def __call__( + self, + mat_a: cute.Tensor, # (hidden, tokens_sum) — activation^T + mat_b: cute.Tensor, # (tokens_sum, intermediate) — activation + out, # Dense: cute.Tensor (experts, hidden, intermediate) + # Discrete: cute.Pointer to int64[] + offs, # Union[cute.Tensor, cute.Pointer] (experts,) cumsum end offsets, int32 + workspace: cute.Tensor, # expert-wise TMA desc (discrete only) + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + # Discrete-only: template tensor for a single expert's output (M, N) or (M, N, 1) + out_single_expert: Optional[cute.Tensor] = None, + ) -> None: + + # ================================================================= + # Step 1: Transform to GEMM domain (2Dx2D) + # ================================================================= + c1 = cutlass.Int32(1) + c0 = cutlass.Int32(0) + + # mat_a: (hidden, tokens_sum) -> A: (M=hidden, K=tokens_sum, L=1) + hidden, tokens_sum = mat_a.shape + a_gemm = cute.make_tensor( + mat_a.iterator, + cute.make_layout( + (hidden, tokens_sum, c1), + stride=(mat_a.stride[0], mat_a.stride[1], c0), + ), + ) + # mat_b: (tokens_sum, intermediate) -> B: (N=intermediate, K=tokens_sum, L=1) + tokens_sum_b, intermediate = mat_b.shape + b_gemm = cute.make_tensor( + mat_b.iterator, + cute.make_layout( + (intermediate, tokens_sum_b, c1), + stride=(mat_b.stride[1], mat_b.stride[0], c0), + ), + ) + + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + # out: (experts, hidden, intermediate) -> C: (M=hidden, N=intermediate, L=experts) + experts, hidden_c, intermediate_c = out.shape + c_gemm = cute.make_tensor( + out.iterator, + cute.make_layout( + (hidden_c, intermediate_c, experts), + stride=(out.stride[1], out.stride[2], out.stride[0]), + ), + ) + expert_cnt = experts + else: + # Discrete: out is a Pointer to int64[] of per-expert base addresses + expert_cnt = self.expert_cnt + # Normalize out_single_expert to rank-3 (M, N, 1) if rank-2 + if cutlass.const_expr(cute.rank(out_single_expert.layout) == 2): + out_single_expert = cute.make_tensor( + out_single_expert.iterator, + cute.make_layout( + (*out_single_expert.shape, c1), + stride=(*out_single_expert.stride, c0), + ), + ) + c_gemm = out_single_expert + + intermediate_dim = intermediate + hidden_dim = hidden + + # ================================================================= + # Step 2: Infer dtypes and major modes + # ================================================================= + self.a_dtype = a_gemm.element_type + self.b_dtype = b_gemm.element_type + self.c_dtype = c_gemm.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a_gemm).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b_gemm).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c_gemm) + + # ================================================================= + # Step 3: Setup kernel attributes + # ================================================================= + self._setup_attributes() + tiled_mma = self._create_tiled_mma() + + # ================================================================= + # Step 4: Create TMA ops (atoms built after helper kernel) + # ================================================================= + # All atoms are built AFTER the helper kernel launch. In ragged + # input_order, the helper kernel calls make_tiled_tma_atom_A/B with + # a_op/b_op, which mutates the op's smem_layout with kernel-region + # block arguments. Re-building the host-region atoms after the launch + # restores valid host-region values. C follows the same idiom. + + # TMA load A + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + + # TMA load B + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + + # TMA store/reduce C + if cutlass.const_expr(self.accumulate_on_output): + c_tma_op = cpasync.CopyReduceBulkTensorTileS2GOp() + else: + c_tma_op = cpasync.CopyBulkTensorTileS2GOp() + + # ================================================================= + # Step 5: Scheduler params and grid + # ================================================================= + sched_params = MoESchedulerParams( + scenario="2Dx2D", + expert_shape=(expert_cnt, intermediate_dim, hidden_dim), + cta_tile_shape_mnk=self.cta_tile_shape_mnk, + cluster_shape_mn=self.cluster_shape_mn, + ) + grid = MoESchedulerParams.get_grid_shape(sched_params, max_active_clusters) + + # ================================================================= + # Step 6: Launch helper kernel when descriptors are needed + # ================================================================= + # Builds A/B (TensorRagged only) + C (Discrete weight_mode only). + # The WgradTensormapConstructor + # is created inside the kernel body from the raw params — it has too + # many Constexpr fields for MLIR serialization as a kernel argument. + epi_smem_layout_helper = cute.select(self.c_smem_layout_staged, mode=[0, 1]) if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else None + # A/B helper kernel inputs (used only when input_order==TensorRagged) + if cutlass.const_expr(self.input_order == WGradInputOrder.TensorRagged): + a_smem_layout_helper = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + b_smem_layout_helper = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + a_gemm_helper = a_gemm + b_gemm_helper = b_gemm + a_op_helper = a_op + b_op_helper = b_op + else: + a_smem_layout_helper = None + b_smem_layout_helper = None + a_gemm_helper = None + b_gemm_helper = None + a_op_helper = None + b_op_helper = None + + if cutlass.const_expr(self.input_order == WGradInputOrder.TensorRagged or self.weight_mode == MoEWeightMode.DISCRETE): + self.helper_kernel( + offs, + workspace.iterator, + tiled_mma, + self.cluster_layout_vmnk.shape, + out if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else None, + c_gemm if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else None, + c_tma_op if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else None, + epi_smem_layout_helper, + self.epi_tile if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE) else None, + a_gemm_helper, + b_gemm_helper, + a_op_helper, + b_op_helper, + a_smem_layout_helper, + b_smem_layout_helper, + ).launch( + grid=(expert_cnt, 1, 1), + block=(1, 1, 1), + stream=stream, + min_blocks_per_mp=1, + ) + + # Build A, B, C TMA atoms AFTER the helper kernel launch. + # make_tiled_tma_atom_*() stores smem_layout on the TMA op object. + # Because the ops are passed as Constexpr to the helper kernel, the + # helper's own make_tiled_tma_atom_*() calls contaminate them with + # kernel-region block arguments. Creating fresh atoms here re-sets + # the ops' smem_layout to valid host-region values. + # A/B atoms are unconditionally rebuilt here: in TensorRagged mode + # the helper has already mutated a_op/b_op; in Tensor2D mode this is + # the first build. In TensorRagged mode the host-region atom's + # internal desc carries Tensor2D-flavored stride that is meaningless + # at runtime — the actual desc used by TMA load is read from the + # workspace pointer (see desc_ptr_a/b in the main kernel). + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a_gemm, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b_gemm, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + epi_smem_layout = cute.select(self.c_smem_layout_staged, mode=[0, 1]) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + c_tma_op, + c_gemm, + epi_smem_layout, + self.epi_tile, + ) + + # ================================================================= + # Step 7: Launch main kernel + # ================================================================= + self.kernel( + tiled_mma, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_c, + tma_tensor_c, + a_gemm, + b_gemm, + c_gemm, + offs, + sched_params, + workspace.iterator, + self.cluster_layout_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.c_smem_layout_staged, + self.epi_tile, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + min_blocks_per_mp=self.occupancy, + ) + + # ------------------------------------------------------------------ + # helper_kernel (expert-wise TMA desc init via construct_and_write) + # ------------------------------------------------------------------ + + @cute.kernel + def helper_kernel( + self, + offs: cute.Tensor, + workspace_ptr, + tiled_mma: cute.TiledMma, + cluster_layout_vmnk_shape: cutlass.Constexpr, + c_ptrs=None, + c_single_expert=None, + c_tma_op: cutlass.Constexpr = None, + epi_smem_layout=None, + epi_tile=None, + a_tensor=None, + b_tensor=None, + a_tma_op: cutlass.Constexpr = None, + b_tma_op: cutlass.Constexpr = None, + a_smem_layout=None, + b_smem_layout=None, + ): + """Build per-expert TMA descriptors. + + Builds: + - A, B when ``input_order == WGradInputOrder.TensorRagged`` + - C when ``weight_mode == MoEWeightMode.DISCRETE`` + + Launched with grid=(expert_cnt, 1, 1). Each block handles one expert. + """ + from ..moe_utils import WgradTensormapConstructor + + ctor = WgradTensormapConstructor( + weight_mode=self.weight_mode, + tiled_mma=tiled_mma, + mma_tiler=self.mma_tiler, + cluster_layout_vmnk_shape=cluster_layout_vmnk_shape, + offs=offs, + workspace_ptr=workspace_ptr, + c_tma_op=c_tma_op, + epi_smem_layout=epi_smem_layout, + epi_tile=epi_tile, + c_ptrs=c_ptrs, + c_single_expert=c_single_expert, + expert_cnt=self.expert_cnt, + input_order=self.input_order, + a_tma_op=a_tma_op, + b_tma_op=b_tma_op, + a_smem_layout=a_smem_layout, + b_smem_layout=b_smem_layout, + a_major_mode=self.a_major_mode, + b_major_mode=self.b_major_mode, + a_tensor=a_tensor, + b_tensor=b_tensor, + ) + expert_idx = cute.arch.block_idx()[0] + ctor.construct_and_write(expert_idx) + + # ------------------------------------------------------------------ + # kernel (GPU device kernel) + # ------------------------------------------------------------------ + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + a_gemm: cute.Tensor, + b_gemm: cute.Tensor, + c_gemm: cute.Tensor, + offs: cute.Tensor, + sched_params: MoESchedulerParams, + workspace_ptr, + cluster_layout_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + c_smem_layout_staged: cute.ComposedLayout, + epi_tile: cute.Tile, + ): + """GPU device kernel for MoE BF16 wgrad with CLC scheduling.""" + + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + tidx, _, _ = cute.arch.thread_idx() + + # ================================================================= + # SharedStorage + # ================================================================= + SchedulerStorage = MoEPersistentTileScheduler.make_storage_struct(self.num_sched_stages, use_dynamic_sched=True) + + @cute.struct + class SharedStorage: + ab_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_pipeline_stages * 2] + scheduler: SchedulerStorage + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + sched_storage = storage.scheduler + + # ================================================================= + # Pipelines + # ================================================================= + + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_tma_producer) + ab_producer, ab_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_full_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_warp_id) * 32 * (2 if use_2cta_instrs else 1) + acc_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_acc_consumer_threads) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_pipeline_stages, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Scheduler pipeline (sched warp -> tma/mma/epi warps) + sched_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32) + num_sched_consumer_threads = 32 * len((self.tma_warp_id, self.mma_warp_id, *self.epilogue_warp_id)) + sched_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread, num_sched_consumer_threads) + sched_pipeline = pipeline.PipelineAsync.create( + num_stages=self.num_sched_stages, + producer_group=sched_producer_group, + consumer_group=sched_consumer_group, + barrier_storage=sched_storage.tile_info_mbar.data_ptr(), + defer_sync=True, + ) + + # TMEM allocator + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, + ) + + # Scheduler (CLC-based for 2Dx2D) + scheduler = MoEPersistentTileScheduler.create( + sched_params, + offs, + cute.arch.block_idx(), + cute.arch.grid_dim(), + counter_ptr=None, + sched_storage=sched_storage, + ) + scheduler.internal_init() + + # Cluster barrier sync after init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # ================================================================= + # SMEM tensors + # ================================================================= + sA = smem.allocate_tensor( + element_type=self.a_dtype, + layout=a_smem_layout_staged.outer, + byte_alignment=128, + swizzle=a_smem_layout_staged.inner, + ) + sB = smem.allocate_tensor( + element_type=self.b_dtype, + layout=b_smem_layout_staged.outer, + byte_alignment=128, + swizzle=b_smem_layout_staged.inner, + ) + + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + # Scheduler buf tensor for sched_pipeline broadcast + sched_buf_ptr = sched_storage.sInfo.data_ptr() + sched_copy_atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Int32, num_bits_per_copy=128) + sched_buf_tensor = cute.make_tensor( + sched_buf_ptr, + cute.make_layout((4, self.num_sched_stages), stride=(1, 4)), + ) + + # Build extension + from ..moe_utils import TensormapWorkspace, WgradTensormapConstructor + + slot_names = WgradTensormapConstructor.slot_names(self.input_order, self.weight_mode) + desc_workspace = TensormapWorkspace(workspace_ptr, slot_names) + ext = WgradGemmSchedExtension( + tensormap_ctor=desc_workspace, + weight_mode=self.weight_mode, + input_order=self.input_order, + ) + + # Cluster wait + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # ================================================================= + # Scheduler warp (warp 6) + # ================================================================= + if warp_idx == self.sched_warp_id: + sched_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_sched_stages) + + work_tile_info = scheduler.initial_work_tile_info() + + sched_pipeline.producer_acquire(sched_producer_state) + rmem = work_tile_info.to_rmem_tensor() + cute.copy( + sched_copy_atom, + rmem, + sched_buf_tensor[(None, sched_producer_state.index)], + ) + cute.arch.fence_proxy("async.shared", space="cta") + sched_pipeline.producer_commit(sched_producer_state) + sched_producer_state.advance() + + work_tile_info = scheduler.advance_to_next_work() + while work_tile_info.is_valid_tile: + sched_pipeline.producer_acquire(sched_producer_state) + rmem = work_tile_info.to_rmem_tensor() + cute.copy( + sched_copy_atom, + rmem, + sched_buf_tensor[(None, sched_producer_state.index)], + ) + cute.arch.fence_proxy("async.shared", space="cta") + sched_pipeline.producer_commit(sched_producer_state) + sched_producer_state.advance() + + work_tile_info = scheduler.advance_to_next_work() + + sched_pipeline.producer_acquire(sched_producer_state) + sentinel = MoEWorkTileInfo( + cutlass.Int32(-1), + cutlass.Int32(0), + cutlass.Int32(0), + cutlass.Int32(0), + ) + rmem = sentinel.to_rmem_tensor() + cute.copy( + sched_copy_atom, + rmem, + sched_buf_tensor[(None, sched_producer_state.index)], + ) + cute.arch.fence_proxy("async.shared", space="cta") + sched_pipeline.producer_commit(sched_producer_state) + sched_pipeline.producer_tail(sched_producer_state) + + # ================================================================= + # TMA load warp (warp 5) + # ================================================================= + if warp_idx == self.tma_warp_id: + a_full_mcast_mask = None + b_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2) + b_full_mcast_mask = cpasync.create_tma_multicast_mask(cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1) + + sched_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_sched_stages) + + sched_pipeline.consumer_wait(sched_consumer_state) + rmem = cute.make_rmem_tensor((4,), cutlass.Int32) + cute.copy( + sched_copy_atom, + sched_buf_tensor[(None, sched_consumer_state.index)], + rmem, + ) + work_tile_info = MoEWorkTileInfo.from_rmem_tensor(rmem) + cute.arch.fence_acq_rel_cta() + sched_pipeline.consumer_release(sched_consumer_state) + sched_consumer_state.advance() + + while work_tile_info.is_valid_tile: + k_tile_cnt = work_tile_info.k_tile_cnt + ext.update_expert_info(offs, work_tile_info.expert_idx) + + real_a, desc_ptr_a = ext.get_gmem_tensor( + "a", + mA_mkl, + offs, + work_tile_info, + ) + real_b, desc_ptr_b = ext.get_gmem_tensor( + "b", + mB_nkl, + offs, + work_tile_info, + ) + + gA_mkl = cute.local_tile( + real_a, + cute.slice_(self.mma_tiler, (None, 0, None)), + (None, None, None), + ) + gB_nkl = cute.local_tile( + real_b, + cute.slice_(self.mma_tiler, (0, None, None)), + (None, None, None), + ) + + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + mma_tile_m = work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape) + tAgA_slice = tAgA[(None, mma_tile_m, None, 0)] + tBgB_slice = tBgB[(None, work_tile_info.tile_n_idx, None, 0)] + + ab_producer.reset() + peek_ab_empty_status = ab_producer.try_acquire() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_producer.acquire_and_advance(peek_ab_empty_status) + peek_ab_empty_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_empty_status = ab_producer.try_acquire() + cute.copy( + tma_atom_a, + tAgA_slice[(None, handle.count)], + tAsA[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_a, + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, handle.count)], + tBsB[(None, handle.index)], + tma_bar_ptr=handle.barrier, + tma_desc_ptr=desc_ptr_b, + mcast_mask=b_full_mcast_mask, + ) + + sched_pipeline.consumer_wait(sched_consumer_state) + rmem = cute.make_rmem_tensor((4,), cutlass.Int32) + cute.copy( + sched_copy_atom, + sched_buf_tensor[(None, sched_consumer_state.index)], + rmem, + ) + work_tile_info = MoEWorkTileInfo.from_rmem_tensor(rmem) + cute.arch.fence_acq_rel_cta() + sched_pipeline.consumer_release(sched_consumer_state) + sched_consumer_state.advance() + ab_producer.tail() + + # ================================================================= + # MMA warp (warp 4) + # ================================================================= + if warp_idx == self.mma_warp_id: + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + acc_producer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Producer, self.num_acc_pipeline_stages) + sched_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_sched_stages) + + sched_pipeline.consumer_wait(sched_consumer_state) + rmem = cute.make_rmem_tensor((4,), cutlass.Int32) + cute.copy( + sched_copy_atom, + sched_buf_tensor[(None, sched_consumer_state.index)], + rmem, + ) + work_tile_info = MoEWorkTileInfo.from_rmem_tensor(rmem) + cute.arch.fence_acq_rel_cta() + sched_pipeline.consumer_release(sched_consumer_state) + sched_consumer_state.advance() + + while work_tile_info.is_valid_tile: + k_tile_cnt = work_tile_info.k_tile_cnt + acc_stage_index = acc_producer_state.index + + if is_leader_cta: + tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] + + ab_consumer.reset() + peek_ab_full_status = cutlass.Boolean(1) + if k_tile_cnt > 0: + peek_ab_full_status = ab_consumer.try_wait() + acc_pipeline.producer_acquire(acc_producer_state) + + tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + handle = ab_consumer.wait_and_advance(peek_ab_full_status) + peek_ab_full_status = cutlass.Boolean(1) + if handle.count + 1 < k_tile_cnt: + peek_ab_full_status = ab_consumer.try_wait() + + tiled_mma.set(tcgen05.Field.ACCUMULATE, k_tile != 0) + tile_crd = (None, None, None, handle.index) + cute.gemm(tiled_mma, tCtAcc, tCrA[tile_crd], tCrB[tile_crd], tCtAcc) + handle.release() + + if k_tile_cnt > 0: + acc_pipeline.producer_commit(acc_producer_state) + if k_tile_cnt > 0: + acc_producer_state.advance() + + sched_pipeline.consumer_wait(sched_consumer_state) + rmem = cute.make_rmem_tensor((4,), cutlass.Int32) + cute.copy( + sched_copy_atom, + sched_buf_tensor[(None, sched_consumer_state.index)], + rmem, + ) + work_tile_info = MoEWorkTileInfo.from_rmem_tensor(rmem) + cute.arch.fence_acq_rel_cta() + sched_pipeline.consumer_release(sched_consumer_state) + sched_consumer_state.advance() + + acc_pipeline.producer_tail(acc_producer_state) + + # ================================================================= + # SMEM tensor C (allocated after MMA section) + # ================================================================= + sC = smem.allocate_tensor( + element_type=self.c_dtype, + layout=c_smem_layout_staged.outer, + byte_alignment=128, + swizzle=c_smem_layout_staged.inner, + ) + + # ================================================================= + # Epilogue warps (warps 0-3) + # ================================================================= + if warp_idx < self.mma_warp_id: + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + acc_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_acc_pipeline_stages) + sched_consumer_state = pipeline.make_pipeline_state(pipeline.PipelineUserType.Consumer, self.num_sched_stages) + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_producer_group) + + epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.epilog_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) + + tCtAcc_transformed = transform_partitioned_tensor_layout(tCtAcc_base) + + num_tiles_executed = cutlass.Int32(0) + + sched_pipeline.consumer_wait(sched_consumer_state) + rmem = cute.make_rmem_tensor((4,), cutlass.Int32) + cute.copy( + sched_copy_atom, + sched_buf_tensor[(None, sched_consumer_state.index)], + rmem, + ) + work_tile_info = MoEWorkTileInfo.from_rmem_tensor(rmem) + cute.arch.fence_acq_rel_cta() + sched_pipeline.consumer_release(sched_consumer_state) + sched_consumer_state.advance() + + while work_tile_info.is_valid_tile: + k_tile_cnt = work_tile_info.k_tile_cnt + ext.update_expert_info(offs, work_tile_info.expert_idx) + + real_c, desc_ptr_c = ext.get_gmem_tensor( + "c", + mC_mnl, + offs, + work_tile_info, + ) + + gC_mnl = cute.local_tile( + real_c, + cute.slice_(self.mma_tiler, (None, None, 0)), + (None, None, None), + ) + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + tCgC = thr_mma.partition_C(gC_mnl) + tCgC_transformed = transform_partitioned_tensor_layout(tCgC) + + mma_tile_coord_mnl = ( + work_tile_info.tile_m_idx // cute.size(tiled_mma.thr_id.shape), + work_tile_info.tile_n_idx, + cutlass.Int32(0), + ) + + tiled_copy_t2r, tTR_tAcc_base_epi, tTR_rAcc = epilogue_tmem_copy_and_partition( + self, + tidx, + tCtAcc_transformed, + tCgC_transformed, + epi_tile, + use_2cta_instrs, + ) + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = epilogue_smem_copy_and_partition(self, tiled_copy_t2r, tTR_rC, tidx, sC) + + tCgC_epi = cute.flat_divide(tCgC_transformed, epi_tile) + bSG_sC, bSG_gC_partitioned = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + cute.group_modes(sC, 0, 2), + cute.group_modes(tCgC_epi, 0, 2), + ) + bSG_gC = bSG_gC_partitioned[(None, None, None, *mma_tile_coord_mnl)] + + acc_stage_index = acc_consumer_state.index + tTR_tAcc = tTR_tAcc_base_epi[(None, None, None, None, None, acc_stage_index)] + + if k_tile_cnt > 0: + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) + + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + num_prev_subtiles = num_tiles_executed * subtile_cnt + + for subtile_idx in cutlass.range(subtile_cnt): + real_subtile_idx = subtile_idx + + tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] + if k_tile_cnt > 0: + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + acc_vec = cute.zeros_like(tiled_copy_r2s.retile(tTR_rAcc), dtype=tTR_rAcc._dtype) + if k_tile_cnt > 0: + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = acc_vec.to(self.c_dtype) + tRS_rC.store(acc_vec) + + c_buffer = (num_prev_subtiles + subtile_idx) % self.num_c_stage + cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) + cute.arch.fence_proxy("async.shared", space="cta") + epilog_sync_barrier.arrive_and_wait() + + if warp_idx == self.epilogue_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, real_subtile_idx)], + tma_desc_ptr=desc_ptr_c, + ) + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + epilog_sync_barrier.arrive_and_wait() + + if k_tile_cnt > 0: + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + num_tiles_executed += cutlass.Int32(1) + + sched_pipeline.consumer_wait(sched_consumer_state) + rmem = cute.make_rmem_tensor((4,), cutlass.Int32) + cute.copy( + sched_copy_atom, + sched_buf_tensor[(None, sched_consumer_state.index)], + rmem, + ) + work_tile_info = MoEWorkTileInfo.from_rmem_tensor(rmem) + cute.arch.fence_acq_rel_cta() + sched_pipeline.consumer_release(sched_consumer_state) + sched_consumer_state.advance() + + c_pipeline.producer_tail() + + tmem.relinquish_alloc_permit() + epilog_sync_barrier.arrive_and_wait() + tmem.free(acc_tmem_ptr) diff --git a/python/cudnn/grouped_gemm/moe_kernel_helpers.py b/python/cudnn/grouped_gemm/moe_kernel_helpers.py index 3cd5862e2..4758cef7d 100644 --- a/python/cudnn/grouped_gemm/moe_kernel_helpers.py +++ b/python/cudnn/grouped_gemm/moe_kernel_helpers.py @@ -754,6 +754,143 @@ def can_implement( return result +def is_valid_bf16_grouped_gemm_dtypes( + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], +) -> bool: + """ + Check if the BF16 grouped GEMM dtypes are valid. + + :return: True if valid, False otherwise + """ + is_valid = True + + valid_output_dtypes = {cutlass.BFloat16, cutlass.Float16, cutlass.Float32} + if ab_dtype != cutlass.BFloat16: + is_valid = False + if c_dtype not in valid_output_dtypes: + is_valid = False + if d_dtype not in valid_output_dtypes: + is_valid = False + if acc_dtype != cutlass.Float32: + is_valid = False + + return is_valid + + +def is_valid_bf16_grouped_gemm_mma_tiler_and_cluster_shape( + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m_aligned: int, + fix_pad_size: int = FIX_PAD_SIZE, + tile_n_align: int = 64, +) -> bool: + """ + Check if the BF16 grouped GEMM MMA tiler and cluster shape are valid. + + :param fix_pad_size: The fixed pad size used by the kernel (default: FIX_PAD_SIZE). + :param tile_n_align: Required alignment of the MMA tile N (mma_tiler_mn[1]). + GLU needs 64 (its epilogue walks 32-column gate/up subtiles in pairs, so the + per-tile subtile count must be even). Plain unfused / DGLU only need 32 (their + epilogue walks independent 32-column subtiles). Default 64. + :return: True if valid, False otherwise + """ + is_valid = True + + if not ((not use_2cta_instrs and mma_tiler_mn[0] in [128]) or (use_2cta_instrs and mma_tiler_mn[0] in [256])): + is_valid = False + if mma_tiler_mn[1] not in range(tile_n_align, 257, tile_n_align): + is_valid = False + if cluster_shape_mn[0] % (2 if use_2cta_instrs else 1) != 0: + is_valid = False + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + cluster_shape_mn[0] * cluster_shape_mn[1] > 16 + or cluster_shape_mn[0] <= 0 + or cluster_shape_mn[1] <= 0 + or not is_power_of_2(cluster_shape_mn[0]) + or not is_power_of_2(cluster_shape_mn[1]) + ): + is_valid = False + cluster_tiler_m = (cluster_shape_mn[0] // (2 if use_2cta_instrs else 1)) * mma_tiler_mn[0] + + if cluster_tiler_m not in [128, 256]: + is_valid = False + + if m_aligned % mma_tiler_mn[0] != 0: + is_valid = False + if m_aligned != fix_pad_size: + is_valid = False + + return is_valid + + +def can_implement_bf16_grouped_gemm( + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + d_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + m: int, + n: int, + k: int, + l: int, + a_major: str, + b_major: str, + cd_major: str, + m_aligned: int, + fix_pad_size: int = FIX_PAD_SIZE, + n_align: int = 64, + tile_n_align: int = 64, +) -> bool: + """ + Check if the BF16 grouped GEMM can be implemented with the given parameters. + + :param fix_pad_size: The fixed pad size used by the kernel (default: FIX_PAD_SIZE). + :param n_align: Required alignment of the problem N dimension. GLU needs 64 + (gate/up are interleaved as even/odd 32-column blocks, so N must contain an + even number of 32-column blocks). Plain unfused / DGLU only need 32 (their + per-tile epilogue walks independent 32-column subtiles). Default 64. + :param tile_n_align: Required alignment of the MMA tile N (mma_tiler_mn[1]). + GLU needs 64 (subtiles consumed in gate/up pairs); unfused / DGLU need 32. + Default 64. + :return: True if implementable, False otherwise + """ + result = True + + if m_aligned != fix_pad_size: + result = False + + if not is_valid_bf16_grouped_gemm_dtypes(ab_dtype, c_dtype, d_dtype, acc_dtype): + result = False + + if not (a_major == "k" and b_major in ("k", "n") and cd_major == "n"): + result = False + + if not is_valid_bf16_grouped_gemm_mma_tiler_and_cluster_shape( + use_2cta_instrs, + mma_tiler_mn, + cluster_shape_mn, + m_aligned, + fix_pad_size, + tile_n_align=tile_n_align, + ): + result = False + + if not is_valid_tensor_alignment(m, n, k, l, ab_dtype, d_dtype, a_major, b_major, cd_major): + result = False + + if n % n_align != 0 or m % 256 != 0: + result = False + + return result + + def compute_stages( tiled_mma: cute.TiledMma, mma_tiler_mnk: Tuple[int, int, int], @@ -916,6 +1053,52 @@ def compute_stages_wgrad( return num_acc_stage, num_ab_stage, num_c_stage +def compute_stages_wgrad_bf16( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + epi_tile: cute.Tile, + c_dtype: Type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + num_smem_capacity: int, + occupancy: int, +) -> Tuple[int, int, int]: + """Compute pipeline stages for BF16 wgrad kernel.""" + num_acc_stage = 2 + num_c_stage = 2 + num_tile_stage = 2 + + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + 1, + ) + b_smem_layout_stage_one = sm100_utils.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + 1, + ) + c_smem_layout_stage_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + + ab_bytes_per_stage = cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + cute.size_in_bytes(b_dtype, b_smem_layout_stage_one) + mbar_helpers_bytes = 1024 + sinfo_bytes = 4 * 4 * num_tile_stage + c_bytes = cute.size_in_bytes(c_dtype, c_smem_layout_stage_one) * num_c_stage + + fixed_overhead = mbar_helpers_bytes + c_bytes + sinfo_bytes + num_ab_stage = (num_smem_capacity // occupancy - fixed_overhead) // ab_bytes_per_stage + + return num_acc_stage, num_ab_stage, num_c_stage + + def get_tma_atom_kind(atom_sm_cnt: cutlass.Int32, mcast: cutlass.Boolean) -> Union[cpasync.CopyBulkTensorTileG2SMulticastOp, cpasync.CopyBulkTensorTileG2SOp]: """ Select the appropriate TMA copy atom based on SM count and multicast flag. diff --git a/python/cudnn/grouped_gemm/moe_sched_extension.py b/python/cudnn/grouped_gemm/moe_sched_extension.py index a82457b58..a8bf915e6 100644 --- a/python/cudnn/grouped_gemm/moe_sched_extension.py +++ b/python/cudnn/grouped_gemm/moe_sched_extension.py @@ -418,3 +418,222 @@ def get_gmem_tensor( return (real, desc) raise ValueError(f"WgradScaledGemmSchedExtension: unknown tensor '{tensor_name}'") + + +class DiscreteWeightGroupedGemmSchedExtension(MoESchedExtension): + """ + MoE scheduler extension for discrete-weight non-scaled grouped GEMM. + + Handles domain conversion for: a, b, c, d, prob, bias. + + B is discrete (per-expert pointer array) and uses expert-wise TMA + descriptors from workspace. A/C/D/prob are contiguous across experts and + indexed by padded token offset. Bias is a dense (N, L) tensor and selected + by expert index. + + Domain conversion: + A: (total_padded_M, K, 1) → domain_offset M by token_offset + B: template (N, K, 1) → rewrite L to dynamic 1, + expert-wise desc + C/D/prob: (total_padded_M, N, 1) → domain_offset M by token_offset + Bias: (N, L) → domain_offset L by expert_idx + + :param tensormap_ctor: Discrete-weight tensormap workspace accessor for B + descriptors. + """ + + def __init__(self, tensormap_ctor: OnlineTensormapDescCreator): + super().__init__(tensormap_ctor) + + def __extract_mlir_values__(self): + return extract_mlir_values(self.tensormap_ctor) + + def __new_from_mlir_values__(self, values): + new_ctor = new_from_mlir_values(self.tensormap_ctor, values) + return DiscreteWeightGroupedGemmSchedExtension(tensormap_ctor=new_ctor) + + def update_expert_info(self, offs, expert_idx): + self.token_offset, self.tokens_i = compute_expert_token_range(offs, expert_idx) + + @cute.jit + def get_gmem_tensor( + self, + tensor_name: str, + gmem_tensor_in_moe_view: cute.Tensor, + offs: cute.Tensor, + work_tile_info: MoEWorkTileInfo, + ): + expert_idx = work_tile_info.expert_idx + if cutlass.const_expr(hasattr(self, "token_offset")): + token_offset, tokens_i = self.token_offset, self.tokens_i + else: + token_offset, tokens_i = compute_expert_token_range(offs, expert_idx) + + shape = gmem_tensor_in_moe_view.shape + c1 = cutlass.Int32(1) + + if cutlass.const_expr(tensor_name == "a"): + real = cute.domain_offset((token_offset, 0, 0), gmem_tensor_in_moe_view) + real = rewrite_tensor_shape(real, (tokens_i, shape[1], c1)) + return (real, None) + + elif cutlass.const_expr(tensor_name == "b"): + real = rewrite_tensor_shape(gmem_tensor_in_moe_view, (shape[0], shape[1], c1)) + desc = tensormap_ptr_for_copy(self.tensormap_ctor.get_desc_ptr("b", expert_idx)) + return (real, desc) + + elif cutlass.const_expr(tensor_name == "bias"): + real = cute.domain_offset((0, expert_idx), gmem_tensor_in_moe_view) + real = rewrite_tensor_shape(real, (shape[0], c1)) + return (real, None) + + else: + real = cute.domain_offset((token_offset, 0, 0), gmem_tensor_in_moe_view) + real = rewrite_tensor_shape(real, (tokens_i, shape[1], c1)) + return (real, None) + + +class ContiguousGroupedGemmSchedExtension(MoESchedExtension): + """ + MoE scheduler extension for contiguous non-scaled grouped GEMM. + + Handles domain conversion for: a, b, c, d, prob, bias. + + All tensors use global TMA descriptors. A/C/D/prob are indexed by padded + token offset. B and bias select the expert through their L dimension. + + Domain conversion: + A: (total_padded_M, K, 1) → domain_offset M by token_offset + B: (N, K, L) → domain_offset L by expert_idx + C/D/prob: (total_padded_M, N, 1) → domain_offset M by token_offset + Bias: (N, L) → domain_offset L by expert_idx + + No constructor parameters are required because contiguous weights do not + use per-expert descriptor workspace. + """ + + def __init__(self): + super().__init__(tensormap_ctor=None) + + def __extract_mlir_values__(self): + return [] + + def __new_from_mlir_values__(self, values): + return ContiguousGroupedGemmSchedExtension() + + def update_expert_info(self, offs, expert_idx): + self.token_offset, self.tokens_i = compute_expert_token_range(offs, expert_idx) + + @cute.jit + def get_gmem_tensor( + self, + tensor_name: str, + gmem_tensor_in_moe_view: cute.Tensor, + offs: cute.Tensor, + work_tile_info: MoEWorkTileInfo, + ): + expert_idx = work_tile_info.expert_idx + if cutlass.const_expr(hasattr(self, "token_offset")): + token_offset, tokens_i = self.token_offset, self.tokens_i + else: + token_offset, tokens_i = compute_expert_token_range(offs, expert_idx) + + shape = gmem_tensor_in_moe_view.shape + c1 = cutlass.Int32(1) + + if cutlass.const_expr(tensor_name == "a"): + real = cute.domain_offset((token_offset, 0, 0), gmem_tensor_in_moe_view) + real = rewrite_tensor_shape(real, (tokens_i, shape[1], c1)) + return (real, None) + + elif cutlass.const_expr(tensor_name == "b"): + real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view) + real = rewrite_tensor_shape(real, (shape[0], shape[1], c1)) + return (real, None) + + elif cutlass.const_expr(tensor_name == "bias"): + real = cute.domain_offset((0, expert_idx), gmem_tensor_in_moe_view) + real = rewrite_tensor_shape(real, (shape[0], c1)) + return (real, None) + + else: + real = cute.domain_offset((token_offset, 0, 0), gmem_tensor_in_moe_view) + real = rewrite_tensor_shape(real, (tokens_i, shape[1], c1)) + return (real, None) + + +class WgradGemmSchedExtension(MoESchedExtension): + """ + BF16 wgrad extension for Dense/Discrete output × Tensor2D/Ragged input. + + Domain conversion (2Dx2D): + A (Tensor2D): (M, total_padded_K, 1) -> domain_offset K, global desc + A (TensorRagged): (M, total_padded_K, 1) -> rewrite shape, expert-wise desc + B (Tensor2D): (N, total_padded_K, 1) -> domain_offset K, global desc + B (TensorRagged): (N, total_padded_K, 1) -> rewrite shape, expert-wise desc + C (Dense): (M, N, expert_cnt) -> domain_offset L by expert_idx, global + C (Discrete): template (M, N, 1) -> rewrite L to dynamic 1, expert-wise + """ + + def __init__( + self, + tensormap_ctor, + weight_mode: MoEWeightMode, + input_order: WGradInputOrder = WGradInputOrder.Tensor2D, + ): + super().__init__(tensormap_ctor) + self.weight_mode = weight_mode + self.input_order = input_order + + def __extract_mlir_values__(self): + return extract_mlir_values(self.tensormap_ctor) + + def __new_from_mlir_values__(self, values): + new_ctor = new_from_mlir_values(self.tensormap_ctor, values) + return WgradGemmSchedExtension( + tensormap_ctor=new_ctor, + weight_mode=self.weight_mode, + input_order=self.input_order, + ) + + def update_expert_info(self, offs, expert_idx): + self.token_offset, self.tokens_i = compute_expert_token_range(offs, expert_idx) + + @cute.jit + def get_gmem_tensor( + self, + tensor_name: str, + gmem_tensor_in_moe_view: cute.Tensor, + offs: cute.Tensor, + work_tile_info: MoEWorkTileInfo, + ): + expert_idx = work_tile_info.expert_idx + if cutlass.const_expr(hasattr(self, "token_offset")): + token_offset, tokens_i = self.token_offset, self.tokens_i + else: + token_offset, tokens_i = compute_expert_token_range(offs, expert_idx) + + shape = gmem_tensor_in_moe_view.shape + c1 = cutlass.Int32(1) + + if cutlass.const_expr(tensor_name in ("a", "b")): + if cutlass.const_expr(self.input_order == WGradInputOrder.TensorRagged): + per_expert_shape = (shape[0], tokens_i, c1) + real = rewrite_tensor_shape(gmem_tensor_in_moe_view, per_expert_shape) + desc = tensormap_ptr_for_copy(self.tensormap_ctor.get_desc_ptr(tensor_name, expert_idx)) + return (real, desc) + real = cute.domain_offset((0, token_offset, 0), gmem_tensor_in_moe_view) + real = rewrite_tensor_shape(real, (shape[0], tokens_i, c1)) + return (real, None) + + elif cutlass.const_expr(tensor_name == "c"): + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DENSE): + real = cute.domain_offset((0, 0, expert_idx), gmem_tensor_in_moe_view) + real = rewrite_tensor_shape(real, (shape[0], shape[1], c1)) + return (real, None) + real = rewrite_tensor_shape(gmem_tensor_in_moe_view, (shape[0], shape[1], c1)) + desc = tensormap_ptr_for_copy(self.tensormap_ctor.get_desc_ptr("c", expert_idx)) + return (real, desc) + + else: + raise ValueError(f"WgradGemmSchedExtension: unknown tensor '{tensor_name}'") diff --git a/python/cudnn/grouped_gemm/moe_utils.py b/python/cudnn/grouped_gemm/moe_utils.py index 0bb540e80..d976b221f 100644 --- a/python/cudnn/grouped_gemm/moe_utils.py +++ b/python/cudnn/grouped_gemm/moe_utils.py @@ -768,3 +768,169 @@ def construct_and_write(self, expert_idx: Int32, dependency=None) -> None: self.epi_tile, ) store_tma_desc(tma_atom_c, self.get_desc_ptr("c", expert_idx)) + + +class WgradTensormapConstructor(OnlineTensormapDescCreator): + """TMA descriptor constructor for BF16 wgrad 2Dx2D A/B/C. + + Workspace slot layouts (per executor): + + +-----------------+--------------+--------------------------+ + | input_order | weight_mode | slot_names | + +-----------------+--------------+--------------------------+ + | Tensor2D | DENSE | [] | + | Tensor2D | DISCRETE | ["c"] | + | TensorRagged | DENSE | ["a", "b"] | + | TensorRagged | DISCRETE | ["a", "b", "c"] | + +-----------------+--------------+--------------------------+ + """ + + def __init__( + self, + weight_mode, + tiled_mma, + mma_tiler, + cluster_layout_vmnk_shape, + offs: cute.Tensor, + workspace_ptr, + c_tma_op=None, + epi_smem_layout=None, + epi_tile=None, + c_ptrs=None, + c_single_expert=None, + expert_cnt: int = 1, + input_order=None, + a_tma_op=None, + b_tma_op=None, + a_smem_layout=None, + b_smem_layout=None, + a_major_mode=None, + b_major_mode=None, + a_tensor: cute.Tensor = None, + b_tensor: cute.Tensor = None, + ) -> None: + super().__init__() + self.weight_mode = weight_mode + self.input_order = input_order if input_order is not None else WGradInputOrder.Tensor2D + self.tiled_mma = tiled_mma + self.mma_tiler = mma_tiler + self.cluster_layout_vmnk_shape = cluster_layout_vmnk_shape + self.offs = offs + self.c_tma_op = c_tma_op + self.epi_smem_layout = epi_smem_layout + self.epi_tile = epi_tile + self.c_ptrs = c_ptrs + self.c_single_expert = c_single_expert + self.expert_cnt = expert_cnt + self.a_tma_op = a_tma_op + self.b_tma_op = b_tma_op + self.a_smem_layout = a_smem_layout + self.b_smem_layout = b_smem_layout + self.a_major_mode = a_major_mode + self.b_major_mode = b_major_mode + self.a_tensor = a_tensor + self.b_tensor = b_tensor + self.workspace = TensormapWorkspace(workspace_ptr, self.slot_names(self.input_order, weight_mode)) + + @staticmethod + def slot_names(input_order, weight_mode) -> list: + names = [] + if input_order == WGradInputOrder.TensorRagged: + names.extend(["a", "b"]) + if weight_mode == MoEWeightMode.DISCRETE: + names.append("c") + return names + + @staticmethod + def get_workspace_size(input_order, weight_mode, expert_cnt: int) -> int: + num_slots = len(WgradTensormapConstructor.slot_names(input_order, weight_mode)) + return TensormapWorkspace.size_bytes(num_slots, expert_cnt) + + @cute.jit + def get_desc_ptr(self, tensor_name: str, executor_idx: Int32) -> Pointer: + return self.workspace.get_ptr(tensor_name, executor_idx) + + @cute.jit + def construct_and_write(self, expert_idx: Int32, dependency=None) -> None: + """Build A/B and/or C per-expert TMA descriptors.""" + from cutlass.cute.nvgpu import cpasync + + token_offset, tokens_i = compute_expert_token_range(self.offs, expert_idx) + c1 = cutlass.Int32(1) + c0 = cutlass.Int32(0) + + if cutlass.const_expr(self.input_order == WGradInputOrder.TensorRagged): + # A + a_dtype = self.a_tensor.element_type + a_m_dim = cute.size(self.a_tensor, mode=[0]) + a_elem_offset = cutlass.Int64(a_m_dim) * cutlass.Int64(token_offset) + a_byte_offset = (a_elem_offset * a_dtype.width) // 8 + a_iter_u8 = cute.recast_ptr(self.a_tensor.iterator, dtype=cutlass.Uint8) + a_iter_e = cute.recast_ptr(a_iter_u8 + a_byte_offset, dtype=a_dtype) + if cutlass.const_expr(self.a_major_mode == OperandMajorMode.K): + a_stride_e = (tokens_i, c1, c0) + else: + a_stride_e = (c1, a_m_dim, c0) + a_tensor_e = cute.make_tensor( + a_iter_e, + cute.make_layout((a_m_dim, tokens_i, c1), stride=a_stride_e), + ) + tma_atom_a, _ = cute.nvgpu.make_tiled_tma_atom_A( + self.a_tma_op, + a_tensor_e, + self.a_smem_layout, + self.mma_tiler, + self.tiled_mma, + self.cluster_layout_vmnk_shape, + ) + store_tma_desc(tma_atom_a, self.get_desc_ptr("a", expert_idx)) + + # B + b_dtype = self.b_tensor.element_type + b_n_dim = cute.size(self.b_tensor, mode=[0]) + b_elem_offset = cutlass.Int64(b_n_dim) * cutlass.Int64(token_offset) + b_byte_offset = (b_elem_offset * b_dtype.width) // 8 + b_iter_u8 = cute.recast_ptr(self.b_tensor.iterator, dtype=cutlass.Uint8) + b_iter_e = cute.recast_ptr(b_iter_u8 + b_byte_offset, dtype=b_dtype) + if cutlass.const_expr(self.b_major_mode == OperandMajorMode.K): + b_stride_e = (tokens_i, c1, c0) + else: + b_stride_e = (c1, b_n_dim, c0) + b_tensor_e = cute.make_tensor( + b_iter_e, + cute.make_layout((b_n_dim, tokens_i, c1), stride=b_stride_e), + ) + tma_atom_b, _ = cute.nvgpu.make_tiled_tma_atom_B( + self.b_tma_op, + b_tensor_e, + self.b_smem_layout, + self.mma_tiler, + self.tiled_mma, + self.cluster_layout_vmnk_shape, + ) + store_tma_desc(tma_atom_b, self.get_desc_ptr("b", expert_idx)) + + if cutlass.const_expr(self.weight_mode == MoEWeightMode.DISCRETE): + c_ptr_tensor = cute.make_tensor( + cute.make_ptr( + cutlass.Int64, + self.c_ptrs.toint(), + cute.typing.AddressSpace.gmem, + assumed_align=8, + ), + cute.make_layout((self.expert_cnt,)), + ) + c_ptr_val = c_ptr_tensor[expert_idx] + c_ptr = cute.make_ptr( + self.c_single_expert.element_type, + c_ptr_val, + cute.typing.AddressSpace.gmem, + ) + c_tensor_i = cute.make_tensor(c_ptr, self.c_single_expert.layout) + tma_atom_c, _ = cpasync.make_tiled_tma_atom( + self.c_tma_op, + c_tensor_i, + self.epi_smem_layout, + self.epi_tile, + ) + store_tma_desc(tma_atom_c, self.get_desc_ptr("c", expert_idx)) diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py index 0f7757d17..1c188cd15 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_dglu.py @@ -7,6 +7,8 @@ import torch import pytest +import cudnn +from unittest.mock import Mock from test_utils import torch_fork_set_rng from fe_api.test_fe_api_utils import DYNAMIC_SHAPES_M_VALUES from fe_api.grouped_gemm.test_grouped_gemm_swiglu_utils import ( @@ -27,6 +29,11 @@ allocate_discrete_dswiglu_output_tensors, check_ref_discrete_dswiglu, ) +from test_grouped_gemm_dglu_bf16_utils import ( + assert_grouped_gemm_dglu_close as assert_grouped_gemm_dglu_bf16_close, + grouped_gemm_dglu_bf16_reference, + make_grouped_gemm_dglu_bf16_problem, +) with_scheduler_modes = pytest.mark.parametrize( "use_dynamic_sched", @@ -64,11 +71,110 @@ def _apply_grouped_gemm_cfg_overrides(cfg, cfg_overrides=None): return cfg +@pytest.mark.L0 +def test_grouped_gemm_dglu_blockscaled_discrete_records_pointer_streams(monkeypatch): + from cudnn.grouped_gemm.grouped_gemm_dglu._blockscaled_api import GroupedGemmDgluBlockScaledAPI + + api = object.__new__(GroupedGemmDgluBlockScaledAPI) + api._logger = Mock() + api._get_default_stream = lambda stream: stream + api._runtime_error_if = lambda condition, message: None + api._has_dbias = False + api.weight_mode = None + api._compiled_kernel = Mock() + + recorded = [] + monkeypatch.setattr( + GroupedGemmDgluBlockScaledAPI, + "_record_pointer_stream", + staticmethod(lambda pointers, stream: recorded.append((pointers, stream))), + raising=False, + ) + + b_ptrs = object() + sfb_ptrs = object() + stream = object() + api.execute( + a_tensor=torch.ones(1), + c_tensor=object(), + d_row_tensor=object(), + d_col_tensor=object(), + sfa_tensor=object(), + padded_offsets=object(), + alpha_tensor=object(), + beta_tensor=object(), + prob_tensor=object(), + dprob_tensor=object(), + b_ptrs=b_ptrs, + sfb_ptrs=sfb_ptrs, + current_stream=stream, + ) + + assert recorded == [(b_ptrs, stream), (sfb_ptrs, stream)] + + +@pytest.mark.L0 +def test_cudnn_all_excludes_module_implementation_helpers(): + unexpected_names = { + "ctypes", + "glob", + "os", + "sys", + "sysconfig", + "importlib", + "is_windows", + "module_name", + "symbols_to_import", + "symbol_name", + "load_cudnn", + "Any", + } + assert unexpected_names.isdisjoint(cudnn.__all__) + assert {"backend_version", "__version__", "Node", "pygraph", "graph", "Graph", "wrapper"}.issubset(cudnn.__all__) + + # --------------------------------------------------------------------------- # Dense mode: Class API # --------------------------------------------------------------------------- +@pytest.mark.L0 +@pytest.mark.parametrize( + ("discrete", "b_major"), + [(False, "k"), (True, "k"), (True, "n")], + ids=["bf16-dense", "bf16-discrete-k-major", "bf16-discrete-n-major"], +) +def test_grouped_gemm_dglu_wrapper_bf16(discrete, b_major): + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("Requires SM100+ for grouped GEMM dGLU BF16 kernel.") + + problem = make_grouped_gemm_dglu_bf16_problem(discrete=discrete, b_major=b_major) + expected_d, expected_dprob, _ = grouped_gemm_dglu_bf16_reference( + problem, + act_func="dswiglu", + linear_offset=0.0, + generate_dbias=False, + ) + kwargs = dict( + a_tensor=problem["a"], + c_tensor=problem["c"], + sfa_tensor=None, + padded_offsets=problem["offsets"], + alpha_tensor=problem["alpha"], + beta_tensor=problem["beta"], + prob_tensor=problem["prob"], + dprob_tensor=problem["dprob"], + d_dtype=torch.bfloat16, + ) + if discrete: + kwargs.update(b_ptrs=problem["b_ptrs"], n=problem["n"], b_dtype=torch.bfloat16, b_major=b_major) + else: + kwargs.update(b_tensor=problem["b"], sfb_tensor=None) + result = cudnn.grouped_gemm_dglu_wrapper_sm100(**kwargs) + assert_grouped_gemm_dglu_bf16_close(result["d_row_tensor"], expected_d) + assert_grouped_gemm_dglu_bf16_close(result["dprob_tensor"], expected_dprob) + + @pytest.mark.L0 @torch_fork_set_rng(seed=0) @with_scheduler_modes diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py index 732b45131..ed86d475d 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_glu.py @@ -7,6 +7,7 @@ import torch import pytest +import cudnn from test_utils import torch_fork_set_rng from fe_api.test_fe_api_utils import DYNAMIC_SHAPES_M_VALUES from fe_api.grouped_gemm.test_grouped_gemm_swiglu_utils import ( @@ -25,6 +26,11 @@ allocate_discrete_output_tensors, check_ref_discrete_grouped_gemm, ) +from test_grouped_gemm_glu_bf16_utils import ( + assert_grouped_gemm_glu_close as assert_grouped_gemm_glu_bf16_close, + grouped_gemm_glu_bf16_reference, + make_grouped_gemm_glu_bf16_problem, +) with_scheduler_modes = pytest.mark.parametrize( "use_dynamic_sched", @@ -50,6 +56,44 @@ def _apply_grouped_gemm_cfg_overrides(cfg, cfg_overrides=None): # --------------------------------------------------------------------------- +@pytest.mark.L0 +@pytest.mark.parametrize( + ("discrete", "b_major"), + [(False, "k"), (True, "k"), (True, "n")], + ids=["bf16-dense", "bf16-discrete-k-major", "bf16-discrete-n-major"], +) +def test_grouped_gemm_glu_wrapper_bf16(discrete, b_major): + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("Requires SM100+ for grouped GEMM GLU BF16 kernel.") + + problem = make_grouped_gemm_glu_bf16_problem(discrete=discrete, b_major=b_major) + expected_c, expected_d = grouped_gemm_glu_bf16_reference( + problem, + act_func="swiglu", + linear_offset=0.0, + geglu_alpha=1.702, + glu_clamp_max=7.0, + glu_clamp_min=-7.0, + ) + kwargs = dict( + a_tensor=problem["a"], + sfa_tensor=None, + padded_offsets=problem["offsets"], + alpha_tensor=problem["alpha"], + bias_tensor=problem["bias"], + prob_tensor=problem["prob"], + c_dtype=torch.bfloat16, + d_dtype=torch.bfloat16, + ) + if discrete: + kwargs.update(b_ptrs=problem["b_ptrs"], n=problem["n"], b_dtype=torch.bfloat16, b_major=b_major) + else: + kwargs.update(b_tensor=problem["b"], sfb_tensor=None) + result = cudnn.grouped_gemm_glu_wrapper_sm100(generate_c=True, **kwargs) + assert_grouped_gemm_glu_bf16_close(result["d_tensor"], expected_d) + assert_grouped_gemm_glu_bf16_close(result["c_tensor"], expected_c) + + @pytest.mark.L0 @torch_fork_set_rng(seed=0) @with_scheduler_modes diff --git a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py index c8952f1fc..76d2bb9a0 100644 --- a/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py +++ b/test/python/fe_api/grouped_gemm/test_grouped_gemm_wgrad.py @@ -14,12 +14,45 @@ check_ref_grouped_gemm_wgrad, wgrad_to_ragged_layout, ) +from test_grouped_gemm_wgrad_bf16_utils import ( + assert_grouped_gemm_wgrad_close as assert_grouped_gemm_wgrad_bf16_close, + grouped_gemm_wgrad_bf16_reference, + make_grouped_gemm_wgrad_bf16_problem, +) # --------------------------------------------------------------------------- # Dense mode: Class API # --------------------------------------------------------------------------- +@pytest.mark.L0 +@pytest.mark.parametrize("discrete", [False, True], ids=["bf16-dense", "bf16-discrete"]) +def test_grouped_gemm_wgrad_wrapper_bf16(discrete): + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("Requires SM100+ for grouped GEMM WGrad BF16 kernel.") + + problem = make_grouped_gemm_wgrad_bf16_problem(discrete=discrete) + expected = grouped_gemm_wgrad_bf16_reference(problem) + kwargs = dict( + a_tensor=problem["a"], + b_tensor=problem["b"], + sfa_tensor=None, + sfb_tensor=None, + offsets_tensor=problem["offsets"], + output_mode="discrete" if discrete else "dense", + wgrad_tensor=problem["output"], + wgrad_ptrs=problem["output_ptrs"], + acc_dtype=torch.float32, + wgrad_dtype=problem["output_dtype"], + mma_tiler_mn=(128, 128), + cluster_shape_mn=(1, 1), + input_order=problem["input_order"], + ) + result = cudnn.grouped_gemm_wgrad_wrapper_sm100(**kwargs) + assert result["wgrad_tensor"] is problem["output"] + assert_grouped_gemm_wgrad_bf16_close(result["wgrad_tensor"], expected) + + def _test_grouped_gemm_wgrad_dense_compile_execute( ab_dtype, wgrad_dtype, @@ -583,6 +616,11 @@ def counted_compile(self): monkeypatch.setattr(grouped_gemm_wgrad_api.GroupedGemmWgradSm100, "check_support", lambda self: True) monkeypatch.setattr(grouped_gemm_wgrad_api.GroupedGemmWgradSm100, "compile", counted_compile) monkeypatch.setattr(grouped_gemm_wgrad_api.GroupedGemmWgradSm100, "execute", lambda self, **kwargs: None) + monkeypatch.setattr( + grouped_gemm_wgrad_api, + "select_grouped_gemm_backend", + lambda **_: grouped_gemm_wgrad_api.GroupedGemmBackend.BLOCK_SCALED, + ) first_inputs = _make_wgrad_wrapper_cache_inputs([8, 12]) second_inputs = _make_wgrad_wrapper_cache_inputs([80, 80]) @@ -628,6 +666,11 @@ def counted_compile(self): monkeypatch.setattr(grouped_gemm_wgrad_api.GroupedGemmWgradSm100, "check_support", lambda self: True) monkeypatch.setattr(grouped_gemm_wgrad_api.GroupedGemmWgradSm100, "compile", counted_compile) monkeypatch.setattr(grouped_gemm_wgrad_api.GroupedGemmWgradSm100, "execute", lambda self, **kwargs: None) + monkeypatch.setattr( + grouped_gemm_wgrad_api, + "select_grouped_gemm_backend", + lambda **_: grouped_gemm_wgrad_api.GroupedGemmBackend.BLOCK_SCALED, + ) inputs = _make_wgrad_wrapper_cache_inputs([8, 12]) diff --git a/test/python/fe_api/test_grouped_gemm_bf16.py b/test/python/fe_api/test_grouped_gemm_bf16.py new file mode 100644 index 000000000..68b85e54c --- /dev/null +++ b/test/python/fe_api/test_grouped_gemm_bf16.py @@ -0,0 +1,51 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""BF16 coverage for the unfused SM100 grouped GEMM wrapper.""" + +from __future__ import annotations + +import pytest +import torch + +from test_grouped_gemm_bf16_utils import ( + assert_grouped_gemm_close, + grouped_gemm_bf16_reference, + make_grouped_gemm_bf16_problem, +) + + +@pytest.fixture(autouse=True) +def require_sm100(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + major, minor = torch.cuda.get_device_capability() + if major * 10 + minor < 100: + pytest.skip("SM100 is required") + + +@pytest.mark.L0 +@pytest.mark.parametrize("discrete", [False, True], ids=["bf16-dense", "bf16-discrete"]) +def test_grouped_gemm_bf16_wrapper(discrete): + from cudnn.grouped_gemm.grouped_gemm_unfused import grouped_gemm_wrapper_sm100 + + problem = make_grouped_gemm_bf16_problem(discrete=discrete, enable_bias=True) + expected_c, expected_d = grouped_gemm_bf16_reference(problem) + kwargs = dict( + a_tensor=problem["a"], + padded_offsets=problem["offsets"], + alpha_tensor=problem["alpha"], + bias_tensor=problem["bias"], + prob_tensor=problem["prob"], + c_dtype=torch.bfloat16, + d_dtype=torch.bfloat16, + generate_c=True, + ) + if discrete: + kwargs.update(b_ptrs=problem["b_ptrs"], n=problem["n"], b_dtype=torch.bfloat16) + else: + kwargs.update(b_tensor=problem["b"]) + + result = grouped_gemm_wrapper_sm100(**kwargs) + assert_grouped_gemm_close(result["c_tensor"], expected_c) + assert_grouped_gemm_close(result["d_tensor"], expected_d) diff --git a/test/python/fe_api/test_grouped_gemm_bf16_utils.py b/test/python/fe_api/test_grouped_gemm_bf16_utils.py new file mode 100644 index 000000000..a73a80330 --- /dev/null +++ b/test/python/fe_api/test_grouped_gemm_bf16_utils.py @@ -0,0 +1,92 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Fixtures and Torch reference for the unfused BF16 grouped GEMM API.""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def make_grouped_gemm_bf16_problem( + *, + m: int = 512, + n: int = 256, + k: int = 128, + experts: int = 2, + discrete: bool = False, + enable_bias: bool = False, + alpha_values: tuple[float, ...] = (0.75, -1.25), +) -> dict[str, Any]: + """Create source-compatible K-major A/B and N-major output inputs.""" + if m % (256 * experts) != 0: + raise ValueError("m must give every expert a 256-aligned token count") + + generator = torch.Generator(device="cuda").manual_seed(20260716 + m + int(discrete)) + a = (torch.randn((1, m, k), generator=generator, device="cuda", dtype=torch.bfloat16) * 0.125).permute(1, 2, 0) + + # Storage is (expert, n, k), so each expert matrix is contiguous K-major. + b_storage = torch.randn((experts, n, k), generator=generator, device="cuda", dtype=torch.bfloat16) * 0.125 + b = b_storage.permute(1, 2, 0) + b_ptrs = torch.tensor( + [b_storage[i].data_ptr() for i in range(experts)], + dtype=torch.int64, + device="cuda", + ) + + group_m = m // experts + offsets = torch.arange(group_m, m + 1, group_m, dtype=torch.int32, device="cuda") + alpha = torch.tensor(alpha_values[:experts], dtype=torch.float32, device="cuda") + prob = torch.linspace(0.25, 0.875, m, dtype=torch.float32, device="cuda").view(m, 1, 1) + bias = None + if enable_bias: + bias = (torch.randn((experts, n), generator=generator, device="cuda", dtype=torch.bfloat16) * 0.125).t() + + return { + "a": a, + "b": None if discrete else b, + "b_storage": b_storage, + "b_ptrs": b_ptrs if discrete else None, + "offsets": offsets, + "alpha": alpha, + "prob": prob, + "bias": bias, + "n": n, + "k": k, + "experts": experts, + } + + +def grouped_gemm_bf16_reference(problem: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: + """Return FP32 C/D references matching the source kernel epilogue exactly.""" + a = problem["a"].float().squeeze(-1) + b_storage = problem["b_storage"].float() + offsets = problem["offsets"].cpu().tolist() + alpha = problem["alpha"].float() + prob = problem["prob"].float().squeeze(-1).squeeze(-1) + bias = problem["bias"] + + c_ref = torch.empty((a.shape[0], problem["n"], 1), dtype=torch.float32, device="cuda") + d_ref = torch.empty_like(c_ref) + begin = 0 + for expert, end in enumerate(offsets): + gemm = torch.matmul(a[begin:end], b_storage[expert].t()) + scaled = alpha[expert] * gemm + if bias is None: + c = scaled + d = prob[begin:end, None] * scaled + else: + c = scaled + prob[begin:end, None] * bias[:, expert].float()[None, :] + d = c + c_ref[begin:end, :, 0] = c + d_ref[begin:end, :, 0] = d + begin = end + return c_ref, d_ref + + +def assert_grouped_gemm_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + """Compare after output conversion with tolerances suitable for BF16 inputs.""" + converted = expected.to(actual.dtype).float() + torch.testing.assert_close(actual.float(), converted, rtol=2e-2, atol=3e-2) diff --git a/test/python/fe_api/test_grouped_gemm_dglu_bf16_utils.py b/test/python/fe_api/test_grouped_gemm_dglu_bf16_utils.py new file mode 100644 index 000000000..191bea197 --- /dev/null +++ b/test/python/fe_api/test_grouped_gemm_dglu_bf16_utils.py @@ -0,0 +1,192 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Fixtures and exact Torch reference for the BF16 dGLU API.""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def make_grouped_gemm_dglu_bf16_problem( + *, + m: int = 512, + n: int = 128, + k: int = 128, + experts: int = 2, + discrete: bool = False, + b_major: str = "k", + alpha_values: tuple[float, ...] = (0.75, -1.25), + beta_values: tuple[float, ...] = (1.5, -0.5), +) -> dict[str, Any]: + """Create source-compatible BF16 dGLU inputs with clamp-sensitive C.""" + if m % (256 * experts) != 0: + raise ValueError("m must give every expert a 256-aligned token count") + if n % 32 != 0: + raise ValueError("n must be divisible by the 32-column dGLU block") + + generator = torch.Generator(device="cuda").manual_seed(20260718 + m + int(discrete)) + a = (torch.randn((1, m, k), generator=generator, device="cuda", dtype=torch.bfloat16) * 0.125).permute(1, 2, 0) + b_reference = ( + torch.randn( + (experts, n, k), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + * 0.125 + ) + if b_major not in ("k", "n"): + raise ValueError(f"b_major must be 'k' or 'n', got {b_major}") + b_storage = b_reference if b_major == "k" else b_reference.transpose(1, 2).contiguous() + b = b_reference.permute(1, 2, 0) + b_ptrs = torch.tensor( + [b_storage[index].data_ptr() for index in range(experts)], + dtype=torch.int64, + device="cuda", + ) + + group_m = m // experts + offsets = torch.arange(group_m, m + 1, group_m, dtype=torch.int32, device="cuda") + alpha = torch.tensor(alpha_values[:experts], dtype=torch.float32, device="cuda") + beta = torch.tensor(beta_values[:experts], dtype=torch.float32, device="cuda") + prob = torch.linspace(-0.75, 0.875, m, dtype=torch.float32, device="cuda").view(m, 1, 1) + + c = ( + torch.randn( + (1, m, 2 * n), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + * 0.5 + ).permute(1, 2, 0) + begin = 0 + for expert, end in enumerate(offsets.cpu().tolist()): + beta_value = float(beta[expert].item()) + # After beta scaling, these values exercise every dGeGLU clamp mask. + c[begin:end, 0, 0] = 8.5 / beta_value + c[begin:end, 1, 0] = -8.5 / beta_value + c[begin:end, 32, 0] = 8.25 / beta_value + c[begin:end, 33, 0] = -8.25 / beta_value + begin = end + + return { + "a": a, + "b": None if discrete else b, + "b_storage": b_storage, + "b_reference": b_reference, + "b_ptrs": b_ptrs if discrete else None, + "c": c, + "offsets": offsets, + "alpha": alpha, + "beta": beta, + "prob": prob, + "dprob": torch.zeros((m, 1, 1), dtype=torch.float32, device="cuda"), + "n": n, + "k": k, + "experts": experts, + } + + +def _interleaved_indices(n: int, device: torch.device) -> tuple[torch.Tensor, ...]: + """Return gate/input destinations and the compact N source order.""" + columns_2n = torch.arange(2 * n, device=device).view((2 * n) // 32, 32) + gate_columns = columns_2n[0::2].reshape(-1) + input_columns = columns_2n[1::2].reshape(-1) + compact_columns = torch.arange(n, device=device).view(n // 32, 32).reshape(-1) + return gate_columns, input_columns, compact_columns + + +def grouped_gemm_dglu_bf16_reference( + problem: dict[str, Any], + *, + act_func: str, + linear_offset: float, + generate_dbias: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Return D, dprob, and optional dbias matching the reference oracle.""" + a = problem["a"].float().squeeze(-1) + b_storage = problem["b_reference"].float() + c = problem["c"].float() + offsets = problem["offsets"].cpu().tolist() + alpha = problem["alpha"].float() + beta = problem["beta"].float() + prob = problem["prob"].float() + n = problem["n"] + + ref = torch.empty((a.shape[0], n, 1), dtype=torch.float32, device="cuda") + c_scaled = torch.empty_like(c, dtype=torch.float32) + begin = 0 + for expert, end in enumerate(offsets): + ref[begin:end, :, 0] = torch.matmul( + a[begin:end] * alpha[expert], + (b_storage[expert] * alpha[expert]).t(), + ) + c_scaled[begin:end, :, 0] = c[begin:end, :, 0] * beta[expert] + begin = end + + gate_columns, input_columns, compact_columns = _interleaved_indices(n, c.device) + gate_unclipped = c_scaled.index_select(1, gate_columns) + input_unclipped = c_scaled.index_select(1, input_columns) + + if act_func == "dswiglu": + gate = gate_unclipped + input_value = input_unclipped + sigmoid = torch.sigmoid(gate) + swish = gate * sigmoid + dprob_terms = swish * input_value * ref + d_gate = ref * prob * input_value * sigmoid * (1.0 + gate * (1.0 - sigmoid)) + d_input = ref * prob * swish + elif act_func == "dgeglu": + gate = torch.clamp(gate_unclipped, max=7.0) + input_value = torch.clamp(input_unclipped, min=-7.0, max=7.0) + sigmoid = torch.sigmoid(1.702 * gate) + swish = gate * sigmoid + dprob_terms = swish * (input_value + linear_offset) * ref + d_gate = ref * sigmoid * (1.0 + 1.702 * gate * (1.0 - sigmoid)) * (input_value + linear_offset) * prob + d_input = ref * gate * sigmoid * prob + + gate_filter = gate_unclipped.clone() + input_filter = input_unclipped.clone() + gate_filter[gate_unclipped > 7.0] = 0.0 + input_filter[(input_unclipped > 7.0) | (input_unclipped < -7.0)] = 0.0 + d_gate = d_gate * gate_filter + d_input = d_input * input_filter + else: + raise ValueError(f"unsupported activation {act_func}") + + chunk_sums = [chunk.sum(dim=1, keepdim=True) for chunk in torch.split(dprob_terms, 32, dim=1)] + dprob = torch.cat(chunk_sums, dim=1).sum(dim=1, keepdim=True) + + d = torch.empty_like(c_scaled) + d.index_copy_(1, gate_columns, d_gate.index_select(1, compact_columns)) + d.index_copy_(1, input_columns, d_input.index_select(1, compact_columns)) + + dbias = None + if generate_dbias: + dbias = torch.zeros( + (problem["experts"], 2 * n, 1), + dtype=torch.bfloat16, + device="cuda", + ) + begin = 0 + for expert, end in enumerate(offsets): + dbias[expert, :, 0] = d[begin:end, :, 0].sum(dim=0).to(torch.bfloat16) + begin = end + return d, dprob, dbias + + +def assert_grouped_gemm_dglu_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + """Compare D/dprob after output conversion with BF16 fast-math tolerance.""" + converted = expected.to(actual.dtype).float() + torch.testing.assert_close(actual.float(), converted, rtol=3e-2, atol=8e-2) + + +def assert_grouped_gemm_dglu_dbias_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + """Use the upstream reduction-order-aware tolerance for BF16 atomics.""" + expected_bf16 = expected.to(torch.bfloat16).float() + atol = max(expected_bf16.abs().max().item() * 0.008 * (4**0.5), 0.1) + torch.testing.assert_close(actual.float(), expected_bf16, rtol=1e-2, atol=atol) diff --git a/test/python/fe_api/test_grouped_gemm_glu_bf16_utils.py b/test/python/fe_api/test_grouped_gemm_glu_bf16_utils.py new file mode 100644 index 000000000..5b419ebd3 --- /dev/null +++ b/test/python/fe_api/test_grouped_gemm_glu_bf16_utils.py @@ -0,0 +1,137 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Fixtures and Torch reference for the BF16 grouped GEMM GLU API.""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def make_grouped_gemm_glu_bf16_problem( + *, + m: int = 512, + n: int = 256, + k: int = 128, + experts: int = 2, + discrete: bool = False, + b_major: str = "k", + enable_bias: bool = False, + alpha_values: tuple[float, ...] = (0.75, -1.25), +) -> dict[str, Any]: + """Create source-compatible BF16 inputs for GLU forward.""" + if m % (256 * experts) != 0: + raise ValueError("m must give every expert a 256-aligned token count") + if n % 64 != 0: + raise ValueError("n must contain paired 32-column gate/up blocks") + + generator = torch.Generator(device="cuda").manual_seed(20260717 + m + int(discrete) + int(enable_bias)) + a = (torch.randn((1, m, k), generator=generator, device="cuda", dtype=torch.bfloat16) * 0.125).permute(1, 2, 0) + b_reference = ( + torch.randn( + (experts, n, k), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + * 0.125 + ) + if b_major not in ("k", "n"): + raise ValueError(f"b_major must be 'k' or 'n', got {b_major}") + b_storage = b_reference if b_major == "k" else b_reference.transpose(1, 2).contiguous() + b = b_reference.permute(1, 2, 0) + b_ptrs = torch.tensor( + [b_storage[index].data_ptr() for index in range(experts)], + dtype=torch.int64, + device="cuda", + ) + + group_m = m // experts + offsets = torch.arange(group_m, m + 1, group_m, dtype=torch.int32, device="cuda") + alpha = torch.tensor(alpha_values[:experts], dtype=torch.float32, device="cuda") + probability = torch.linspace(0.25, 0.875, m, dtype=torch.float32, device="cuda").view(m, 1, 1) + bias = None + if enable_bias: + bias = ( + torch.randn( + (experts, n), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + * 0.125 + ).t() + + return { + "a": a, + "b": None if discrete else b, + "b_storage": b_storage, + "b_reference": b_reference, + "b_ptrs": b_ptrs if discrete else None, + "offsets": offsets, + "alpha": alpha, + "prob": probability, + "bias": bias, + "n": n, + "k": k, + "experts": experts, + } + + +def _interleaved_gate_up(c_tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Split alternating 32-column blocks into gate and up tensors.""" + n = c_tensor.shape[1] + columns = torch.arange(n, device=c_tensor.device).view(n // 32, 32) + gate_columns = columns[0::2].reshape(-1) + up_columns = columns[1::2].reshape(-1) + return ( + c_tensor.index_select(1, gate_columns), + c_tensor.index_select(1, up_columns), + ) + + +def grouped_gemm_glu_bf16_reference( + problem: dict[str, Any], + *, + act_func: str, + linear_offset: float, + geglu_alpha: float = 1.702, + glu_clamp_max: float = 7.0, + glu_clamp_min: float = -7.0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return FP32 C/D matching the upstream 32-column GLU epilogue.""" + a = problem["a"].float().squeeze(-1) + b_storage = problem["b_reference"].float() + offsets = problem["offsets"].cpu().tolist() + alpha = problem["alpha"].float() + probability = problem["prob"].float() + bias = problem["bias"] + + c_reference = torch.empty((a.shape[0], problem["n"], 1), dtype=torch.float32, device="cuda") + begin = 0 + for expert, end in enumerate(offsets): + c_value = alpha[expert] * torch.matmul(a[begin:end], b_storage[expert].t()) + if bias is not None: + c_value = c_value + bias[:, expert].float()[None, :] + c_reference[begin:end, :, 0] = c_value + begin = end + + gate, up = _interleaved_gate_up(c_reference) + if act_func == "swiglu": + activated = up * torch.nn.functional.silu(gate) + elif act_func == "geglu": + gate = torch.clamp(gate, max=glu_clamp_max) + up = torch.clamp(up, min=glu_clamp_min, max=glu_clamp_max) + activated = (up + linear_offset) * gate * torch.sigmoid(geglu_alpha * gate) + else: + raise ValueError(f"unsupported activation {act_func}") + d_reference = activated * probability + return c_reference, d_reference + + +def assert_grouped_gemm_glu_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + """Compare after output conversion with BF16/CuTe fast-math tolerances.""" + converted = expected.to(actual.dtype).float() + torch.testing.assert_close(actual.float(), converted, rtol=2e-2, atol=4e-2) diff --git a/test/python/fe_api/test_grouped_gemm_wgrad_bf16_utils.py b/test/python/fe_api/test_grouped_gemm_wgrad_bf16_utils.py new file mode 100644 index 000000000..571cc90f5 --- /dev/null +++ b/test/python/fe_api/test_grouped_gemm_wgrad_bf16_utils.py @@ -0,0 +1,134 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: MIT + +"""Exact upstream-style fixtures and Torch reference for BF16 MoE wgrad.""" + +from __future__ import annotations + +from typing import Any, Iterable + +import torch + + +def _pack_ragged( + logical: torch.Tensor, + group_k_list: Iterable[int], + *, + operand: str, + major: str, +) -> torch.Tensor: + """Pack expert slices exactly as WgradTensormapConstructor expects.""" + pieces = [] + begin = 0 + for group_k in group_k_list: + if operand == "a": + part = logical[:, begin : begin + group_k] + packed = part.contiguous() if major == "k" else part.t().contiguous() + else: + part = logical[begin : begin + group_k, :] + packed = part.t().contiguous() if major == "k" else part.contiguous() + pieces.append(packed.reshape(-1)) + begin += group_k + storage = torch.cat(pieces) if pieces else logical.new_empty((0,)) + if operand == "a": + m, tokens = logical.shape + stride = (tokens, 1) if major == "k" else (1, m) + return torch.as_strided(storage, (m, tokens), stride) + tokens, n = logical.shape + stride = (1, tokens) if major == "k" else (n, 1) + return torch.as_strided(storage, (tokens, n), stride) + + +def make_grouped_gemm_wgrad_bf16_problem( + *, + group_k_list: tuple[int, ...] = (256, 0, 256), + m: int = 128, + n: int = 128, + a_major: str = "k", + b_major: str = "k", + input_order: str = "tensor2d", + output_dtype: torch.dtype = torch.bfloat16, + discrete: bool = False, + initial_value: float | None = None, +) -> dict[str, Any]: + """Create logical inputs plus source-compatible 2-D/ragged physical views.""" + if a_major not in ("k", "m") or b_major not in ("k", "n"): + raise ValueError("unsupported major mode") + if any(group_k < 0 or group_k % 256 for group_k in group_k_list): + raise ValueError("group K sizes must be non-negative multiples of 256") + tokens = sum(group_k_list) + generator = torch.Generator(device="cuda").manual_seed(20260719 + tokens + 11 * len(group_k_list) + int(discrete)) + logical_a = torch.randn((m, tokens), device="cuda", dtype=torch.bfloat16, generator=generator) * 0.125 + logical_b = torch.randn((tokens, n), device="cuda", dtype=torch.bfloat16, generator=generator) * 0.125 + a = _pack_ragged( + logical_a, + (tokens,) if input_order == "tensor2d" else group_k_list, + operand="a", + major=a_major, + ) + b = _pack_ragged( + logical_b, + (tokens,) if input_order == "tensor2d" else group_k_list, + operand="b", + major=b_major, + ) + offsets = torch.tensor( + [sum(group_k_list[: index + 1]) for index in range(len(group_k_list))], + dtype=torch.int32, + device="cuda", + ) + output = torch.empty((len(group_k_list), m, n), dtype=output_dtype, device="cuda") + if initial_value is not None: + output.fill_(initial_value) + output_ptrs = torch.tensor( + [output[index].data_ptr() for index in range(len(group_k_list))], + dtype=torch.int64, + device="cuda", + ) + return { + "a": a, + "b": b, + "logical_a": logical_a, + "logical_b": logical_b, + "offsets": offsets, + "output": output, + "output_ptrs": output_ptrs if discrete else None, + "group_k_list": group_k_list, + "m": m, + "n": n, + "experts": len(group_k_list), + "a_major": a_major, + "b_major": b_major, + "input_order": input_order, + "output_dtype": output_dtype, + "discrete": discrete, + "initial_value": initial_value, + } + + +def grouped_gemm_wgrad_bf16_reference(problem: dict[str, Any], *, accumulate: bool = False) -> torch.Tensor: + """Compute the exact per-expert grouped-mm oracle in FP32.""" + result = torch.empty( + (problem["experts"], problem["m"], problem["n"]), + dtype=torch.float32, + device="cuda", + ) + begin = 0 + for expert, end in enumerate(problem["offsets"].cpu().tolist()): + if end == begin: + result[expert].zero_() + else: + result[expert] = torch.matmul( + problem["logical_a"][:, begin:end].float(), + problem["logical_b"][begin:end, :].float(), + ) + begin = end + if accumulate and problem["initial_value"] is not None: + result.add_(float(problem["initial_value"])) + return result + + +def assert_grouped_gemm_wgrad_close(actual: torch.Tensor, expected: torch.Tensor) -> None: + """Compare after the requested output conversion with upstream tolerance.""" + converted = expected.to(actual.dtype).float() + torch.testing.assert_close(actual.float(), converted, rtol=3e-2, atol=8e-2)