v4.5 tag update (#3202)
* Python DSL examples reorganization. * v4.5 tag update.
This commit is contained in:
Junkai-Wu
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
+397
@@ -0,0 +1,397 @@
|
||||
# Copyright (c) 2025 - 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.
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def ssd_reference_fp32_all(x, a, delta, B, C, Y_out, Fstate_out, D, has_d, d_has_hdim):
|
||||
"""
|
||||
Rearrange tensor dimensions from cuda layout to reference layout, then directly call TriDao's ssd implementation
|
||||
Arguments:
|
||||
X/x: (D, L, C, H, B):(C*L, 1, L, D*C*L, H*D*C*L)
|
||||
A/delta: (L, C, H, B):(1, L, C*L, H*C*L)
|
||||
a: (H):(1)
|
||||
B/C: (L, N, C, G, B):(1, C*L, L, N*C*L, G*N*C*L)
|
||||
D: (1, H):(0, 1) or (D, H):(1, D)
|
||||
has_d: bool
|
||||
d_has_hdim: bool
|
||||
Return:
|
||||
Y_out: (L, D, C, H, B):(1, C*L, L, D*C*L, H*D*C*L)
|
||||
Fstate_out: (D, N, H, B):(N, 1, D*N, H*D*N)
|
||||
"""
|
||||
assert x.dtype == a.dtype == delta.dtype == B.dtype == C.dtype
|
||||
|
||||
A = delta * a.view(1, 1, -1, 1)
|
||||
X = x * delta.unsqueeze(0)
|
||||
|
||||
# Rearrange to match cutlass layout to tridao's layout
|
||||
block_len = A.shape[0]
|
||||
initial_states = None
|
||||
# A: l c h b-> b c l h
|
||||
A = A.permute(3, 1, 0, 2)
|
||||
# X: p l c h b -> b c l h p
|
||||
X = X.permute(4, 2, 1, 3, 0)
|
||||
# B: l n c g b -> b c l g n
|
||||
B = B.permute(4, 2, 0, 3, 1)
|
||||
# C: l n c g b -> b c l g n
|
||||
C = C.permute(4, 2, 0, 3, 1)
|
||||
# X/A/B/C: b c l ... -> b (c l) ...
|
||||
X, A, B, C = [x.reshape(x.shape[0], -1, *x.shape[3:]) for x in (X, A, B, C)]
|
||||
|
||||
# Ngroup (g to h) mapping
|
||||
B_val, CL_val, G_val, N_val = B.shape
|
||||
H_val = X.shape[2]
|
||||
ngroup_ratio = H_val // G_val
|
||||
# B/C: (B, CL, H, N)
|
||||
h_to_g_mapping = torch.arange(H_val, device=B.device) // ngroup_ratio
|
||||
B = B.gather(2, h_to_g_mapping.view(1, 1, -1, 1).expand(B_val, CL_val, -1, N_val))
|
||||
C = C.gather(2, h_to_g_mapping.view(1, 1, -1, 1).expand(B_val, CL_val, -1, N_val))
|
||||
|
||||
###################################################################
|
||||
# Call reference implementation from Tri Dao ssd_minimal_discrete
|
||||
Y, final_state = ssd_minimal_discrete_fp32_all(
|
||||
X, A, B, C, block_len, initial_states
|
||||
)
|
||||
###################################################################
|
||||
|
||||
if has_d:
|
||||
D_val = Y.shape[3]
|
||||
if not d_has_hdim:
|
||||
D = D.expand(D_val, -1)
|
||||
Y = Y + torch.einsum("bchp,ph->bchp", X, D)
|
||||
|
||||
# Rearrange to match tridao's layout to cutlass layout
|
||||
# Y: b (c l) h p -> b c l h p
|
||||
Y = Y.reshape(Y.shape[0], -1, block_len, Y.shape[2], Y.shape[3])
|
||||
# Y: b c l h p -> l p c h b
|
||||
Y = Y.permute(2, 4, 1, 3, 0)
|
||||
# Fstate_out: b h p n -> p n h b
|
||||
Fstate_out.copy_(final_state.permute(2, 3, 1, 0))
|
||||
Y_out.copy_(Y)
|
||||
return
|
||||
|
||||
|
||||
def ssd_reference_lowprecision_intermediates(
|
||||
x, a, delta, B, C, Y_out, Fstate_out, intermediate_dtype, D, has_d, d_has_hdim
|
||||
):
|
||||
"""
|
||||
Rearrange tensor dimensions from cuda layout to reference layout, then call a reduced intermediate dtype version of ssd implementation
|
||||
Arguments:
|
||||
X/x: (D, L, C, H, B):(C*L, 1, L, D*C*L, H*D*C*L)
|
||||
A/delta: (L, C, H, B):(1, L, C*L, H*C*L)
|
||||
a: (H):(1)
|
||||
B/C: (L, N, C, G, B):(1, C*L, L, N*C*L, G*N*C*L)
|
||||
intermediate_dtype: input and intermediate data type
|
||||
D: (1, H):(0, 1) or (D, H):(1, D)
|
||||
has_d: bool
|
||||
d_has_hdim: bool
|
||||
Return:
|
||||
Y_out: (L, D, C, H, B):(1, C*L, L, D*C*L, H*D*C*L)
|
||||
Fstate_out: (D, N, H, B):(N, 1, D*N, H*D*N)
|
||||
"""
|
||||
assert x.dtype == a.dtype == delta.dtype == B.dtype == C.dtype
|
||||
|
||||
A = delta * a.view(1, 1, -1, 1)
|
||||
|
||||
# Rearrange to match cutlass layout to tridao's layout
|
||||
block_len = A.shape[0]
|
||||
initial_states = None
|
||||
# A: l c h b-> b c l h
|
||||
A = A.permute(3, 1, 0, 2)
|
||||
# delta: l c h b-> b c l h
|
||||
delta = delta.permute(3, 1, 0, 2)
|
||||
# x: p l c h b -> b c l h p
|
||||
x = x.permute(4, 2, 1, 3, 0)
|
||||
# B: l n c g b -> b c l g n
|
||||
B = B.permute(4, 2, 0, 3, 1)
|
||||
# C: l n c g b -> b c l g n
|
||||
C = C.permute(4, 2, 0, 3, 1)
|
||||
# x/A/delta/B/C: b c l ... -> b (c l) ...
|
||||
x, A, delta, B, C = [
|
||||
tensor.reshape(tensor.shape[0], -1, *tensor.shape[3:])
|
||||
for tensor in (x, A, delta, B, C)
|
||||
]
|
||||
|
||||
# Ngroup (g to h) mapping
|
||||
B_val, CL_val, G_val, N_val = B.shape
|
||||
H_val = x.shape[2]
|
||||
ngroup_ratio = H_val // G_val
|
||||
# B/C: (B, CL, H, N)
|
||||
h_to_g_mapping = torch.arange(H_val, device=B.device) // ngroup_ratio
|
||||
B = B.gather(2, h_to_g_mapping.view(1, 1, -1, 1).expand(B_val, CL_val, -1, N_val))
|
||||
C = C.gather(2, h_to_g_mapping.view(1, 1, -1, 1).expand(B_val, CL_val, -1, N_val))
|
||||
|
||||
# Type convert input tensors to input dtype (same as intermediate dtype)
|
||||
x = x.to(intermediate_dtype).to(torch.float32)
|
||||
A = A.to(intermediate_dtype).to(torch.float32)
|
||||
delta = delta.to(intermediate_dtype).to(torch.float32)
|
||||
B = B.to(intermediate_dtype).to(torch.float32)
|
||||
C = C.to(intermediate_dtype).to(torch.float32)
|
||||
|
||||
#########################################################################
|
||||
# Call reference implementation ssd_minimal_discrete_bf16_intermediates
|
||||
Y, final_state = ssd_minimal_discrete_lowprecision_intermediates(
|
||||
x, A, delta, B, C, block_len, intermediate_dtype, initial_states
|
||||
)
|
||||
#########################################################################
|
||||
|
||||
if has_d:
|
||||
D = D.to(intermediate_dtype).to(torch.float32)
|
||||
D_val = Y.shape[3]
|
||||
if not d_has_hdim:
|
||||
D = D.expand(D_val, -1)
|
||||
Y = Y + torch.einsum("bchp,ph->bchp", x, D)
|
||||
|
||||
# Type convert output tensors to output dtype (same as intermediate dtype)
|
||||
Y = Y.to(intermediate_dtype).to(torch.float32)
|
||||
final_state = final_state.to(intermediate_dtype).to(torch.float32)
|
||||
|
||||
# Rearrange to match tridao's layout to cutlass layout
|
||||
# Y: b (c l) h p -> b c l h p
|
||||
Y = Y.reshape(Y.shape[0], -1, block_len, Y.shape[2], Y.shape[3])
|
||||
# Y: b c l h p -> l p c h b
|
||||
Y = Y.permute(2, 4, 1, 3, 0)
|
||||
# Fstate_out: b h p n -> p n h b
|
||||
Fstate_out.copy_(final_state.permute(2, 3, 1, 0))
|
||||
Y_out.copy_(Y)
|
||||
return
|
||||
|
||||
|
||||
def analyze_relative_diffs(actual, expected):
|
||||
"""
|
||||
Print statistics of relative differences between actual and expected tensors
|
||||
"""
|
||||
# Calculate relative differences
|
||||
abs_diff = (actual - expected).abs()
|
||||
rel_diff = abs_diff / (torch.maximum(expected.abs(), actual.abs()) + 0.00001)
|
||||
|
||||
total_elements = rel_diff.numel()
|
||||
|
||||
# Handle special cases first
|
||||
nan_mask = torch.isnan(rel_diff)
|
||||
inf_mask = torch.isinf(rel_diff)
|
||||
nan_count = nan_mask.sum().item()
|
||||
inf_count = inf_mask.sum().item()
|
||||
|
||||
# Find position and value of maximum relative difference
|
||||
max_rel_diff = (
|
||||
rel_diff[~nan_mask & ~inf_mask].max()
|
||||
if (~nan_mask & ~inf_mask).any()
|
||||
else float("nan")
|
||||
)
|
||||
max_rel_diff_pos = (
|
||||
rel_diff[~nan_mask & ~inf_mask].argmax()
|
||||
if (~nan_mask & ~inf_mask).any()
|
||||
else -1
|
||||
)
|
||||
|
||||
# Print max relative difference info
|
||||
print("Maximum relative difference:")
|
||||
print(f"Position: {max_rel_diff_pos}")
|
||||
print(f"Value: {max_rel_diff:.6e}")
|
||||
print(f"Actual value: {actual.flatten()[max_rel_diff_pos]}")
|
||||
print(f"Expected value: {expected.flatten()[max_rel_diff_pos]}")
|
||||
print(f"NaN values: {nan_count} ({100.0 * nan_count / total_elements:.2f}%)")
|
||||
print(f"Inf values: {inf_count} ({100.0 * inf_count / total_elements:.2f}%)\n")
|
||||
|
||||
# Check different rtol thresholds
|
||||
rtol_levels = [1e-5, 1e-4, 1e-3, 1e-2, 5e-02, 1e-01]
|
||||
|
||||
for i, rtol in enumerate(rtol_levels):
|
||||
if i == 0:
|
||||
mask = rel_diff <= rtol
|
||||
else:
|
||||
mask = (rel_diff <= rtol) & (rel_diff > rtol_levels[i - 1])
|
||||
|
||||
count = mask.sum().item()
|
||||
percentage = (count / total_elements) * 100
|
||||
|
||||
if i == 0:
|
||||
print(f"Elements with rtol <= {rtol:.0e}: {count} ({percentage:.2f}%)")
|
||||
else:
|
||||
print(
|
||||
f"Elements with {rtol_levels[i - 1]:.0e} < rtol <= {rtol:.0e}: {count} ({percentage:.2f}%)"
|
||||
)
|
||||
|
||||
# Print elements exceeding the largest rtol
|
||||
mask = rel_diff > rtol_levels[-1]
|
||||
count = mask.sum().item()
|
||||
percentage = (count / total_elements) * 100
|
||||
print(f"Elements with rtol > {rtol_levels[-1]:.0e}: {count} ({percentage:.2f}%)\n")
|
||||
|
||||
|
||||
def segsum(x):
|
||||
"""
|
||||
More stable segment sum calculation.
|
||||
x: b h c l
|
||||
"""
|
||||
T = x.size(-1)
|
||||
# x: b h c l -> b h c l l
|
||||
x = x.unsqueeze(-1).expand(*x.shape, T)
|
||||
mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=-1)
|
||||
x = x.masked_fill(~mask, 0)
|
||||
x_segsum = torch.cumsum(x, dim=-2)
|
||||
mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool), diagonal=0)
|
||||
x_segsum = x_segsum.masked_fill(~mask, -torch.inf)
|
||||
return x_segsum
|
||||
|
||||
|
||||
def ssd_minimal_discrete_fp32_all(X, A, B, C, block_len, initial_states=None):
|
||||
"""
|
||||
This is same with https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/ssd_minimal.py
|
||||
(all accumulation and intermediate results in fp32)
|
||||
|
||||
Arguments:
|
||||
X: (batch(B), length(C*L), n_heads(H), d_head(D))
|
||||
A: (batch(B), length(C*L), n_heads(H))
|
||||
B: (batch(B), length(C*L), n_heads(H), d_state(N))
|
||||
C: (batch(B), length(C*L), n_heads(H), d_state(N))
|
||||
Return:
|
||||
Y: (batch(B), length(C*L), n_heads(H), d_head(D))
|
||||
final_state: (B, H, D, N)
|
||||
"""
|
||||
assert X.dtype == A.dtype == B.dtype == C.dtype
|
||||
assert X.shape[1] % block_len == 0
|
||||
|
||||
# Rearrange into blocks/chunks
|
||||
# X/A/B/C:b (c l) ... -> b c l ...
|
||||
X, A, B, C = [
|
||||
x.reshape(x.shape[0], -1, block_len, *x.shape[2:]) for x in (X, A, B, C)
|
||||
]
|
||||
|
||||
# A: b c l h -> b h c l
|
||||
A = A.permute(0, 3, 1, 2)
|
||||
# A_cumsum: (B, H, C, L)
|
||||
A_cumsum = torch.cumsum(A, dim=-1)
|
||||
|
||||
# 1. Compute the output for each intra-chunk (diagonal blocks)
|
||||
segsum_A = segsum(A)
|
||||
L = torch.exp(segsum_A)
|
||||
Y_diag = torch.einsum("bclhn,bcshn,bhcls,bcshp->bclhp", C, B, L, X)
|
||||
|
||||
# 2. Compute the state for each intra-chunk
|
||||
# (right term of low-rank factorization of off-diagonal blocks; B terms)
|
||||
decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum))
|
||||
states = torch.einsum("bclhn,bhcl,bclhp->bchpn", B, decay_states, X)
|
||||
|
||||
# 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
|
||||
# (middle term of factorization of off-diag blocks; A terms)
|
||||
if initial_states is None:
|
||||
initial_states = torch.zeros_like(states[:, :1])
|
||||
states = torch.cat([initial_states, states], dim=1)
|
||||
decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
|
||||
new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
|
||||
states, final_state = new_states[:, :-1], new_states[:, -1]
|
||||
|
||||
# 4. Compute state -> output conversion per chunk
|
||||
# (left term of low-rank factorization of off-diagonal blocks; C terms)
|
||||
state_decay_out = torch.exp(A_cumsum)
|
||||
Y_off = torch.einsum("bclhn,bchpn,bhcl->bclhp", C, states, state_decay_out)
|
||||
|
||||
# Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
|
||||
# Y: b c l h p -> b (c l) h p
|
||||
Y = (Y_diag + Y_off).reshape(Y_diag.shape[0], -1, Y_diag.shape[3], Y_diag.shape[4])
|
||||
return Y, final_state
|
||||
|
||||
|
||||
def ssd_minimal_discrete_lowprecision_intermediates(
|
||||
X, A, delta, B, C, block_len, intermediate_dtype, initial_states=None
|
||||
):
|
||||
"""
|
||||
This is adjusted from ssd_minimal_discrete_fp32_all, with exceptions:
|
||||
1. accumulation in fp32 but intermediates Q/b_tmem/P are in intermediate_dtype
|
||||
2. delta is not pre-multiplied with X, delta was applied to generate Q/b_tmem to match GPU implementation
|
||||
|
||||
Arguments:
|
||||
X: (batch(B), length(C*L), n_heads(H), d_head(D))
|
||||
A: (batch(B), length(C*L), n_heads(H))
|
||||
delta: (batch(B), length(C*L), n_heads(H))
|
||||
B: (batch(B), length(C*L), n_heads(H), d_state(N))
|
||||
C: (batch(B), length(C*L), n_heads(H), d_state(N))
|
||||
Return:
|
||||
Y: (batch(B), length(C*L), n_heads(H), d_head(D))
|
||||
final_state: (B, H, D, N)
|
||||
"""
|
||||
assert X.dtype == A.dtype == B.dtype == C.dtype
|
||||
assert X.shape[1] % block_len == 0
|
||||
|
||||
# Rearrange into blocks/chunks
|
||||
# X/A/delta/B/C: b (c l) ... -> b c l ...
|
||||
X, A, delta, B, C = [
|
||||
x.reshape(x.shape[0], -1, block_len, *x.shape[2:]) for x in (X, A, delta, B, C)
|
||||
]
|
||||
|
||||
# A: b c l h -> b h c l
|
||||
A = A.permute(0, 3, 1, 2)
|
||||
# delta: b c l h -> b h c l
|
||||
delta = delta.permute(0, 3, 1, 2)
|
||||
# A_cumsum: (B, H, C, L)
|
||||
A_cumsum = torch.cumsum(A, dim=-1)
|
||||
|
||||
# 1. Compute the output for each intra-chunk (diagonal blocks)
|
||||
segsum_A = segsum(A)
|
||||
L = torch.exp(segsum_A)
|
||||
intra_acc_0 = torch.einsum("bclhn,bcshn->bclhs", C, B)
|
||||
Q = torch.einsum("bclhs,bhcls,bhcs->bclhs", intra_acc_0, L, delta)
|
||||
Y_diag = torch.einsum(
|
||||
"bclhs,bcshp->bclhp", Q.to(intermediate_dtype).to(torch.float32), X
|
||||
)
|
||||
|
||||
# 2. Compute the state for each intra-chunk
|
||||
# (right term of low-rank factorization of off-diagonal blocks; B terms)
|
||||
decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum))
|
||||
b_tmem = torch.einsum("bclhn,bhcl,bhcl->bclhn", B, decay_states, delta)
|
||||
states = torch.einsum(
|
||||
"bclhn,bclhp->bchpn", b_tmem.to(intermediate_dtype).to(torch.float32), X
|
||||
)
|
||||
|
||||
# 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
|
||||
# (middle term of factorization of off-diag blocks; A terms)
|
||||
if initial_states is None:
|
||||
initial_states = torch.zeros_like(states[:, :1])
|
||||
states = torch.cat([initial_states, states], dim=1)
|
||||
decay_chunk = torch.exp(segsum(F.pad(A_cumsum[:, :, :, -1], (1, 0))))
|
||||
new_states = torch.einsum("bhzc,bchpn->bzhpn", decay_chunk, states)
|
||||
states, final_state = new_states[:, :-1], new_states[:, -1]
|
||||
final_state = final_state
|
||||
|
||||
# 4. Compute state -> output conversion per chunk
|
||||
# (left term of low-rank factorization of off-diagonal blocks; C terms)
|
||||
state_decay_out = torch.exp(A_cumsum)
|
||||
Y_off_tmp = torch.einsum(
|
||||
"bclhn,bchpn->bclhp", C, states.to(intermediate_dtype).to(torch.float32)
|
||||
)
|
||||
Y_off = torch.einsum("bclhp,bhcl->bclhp", Y_off_tmp, state_decay_out)
|
||||
|
||||
# Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
|
||||
# Y: b c l h p -> b (c l) h p
|
||||
Y = (Y_diag + Y_off).reshape(
|
||||
Y_diag.shape[0], -1, Y_diag.shape[3], Y_diag.shape[4]
|
||||
) # b (c l) h p
|
||||
return Y, final_state
|
||||
+199
@@ -0,0 +1,199 @@
|
||||
# Copyright (c) 2025 - 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.
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
from cutlass.cutlass_dsl import (
|
||||
Integer,
|
||||
Int32,
|
||||
min,
|
||||
extract_mlir_values,
|
||||
new_from_mlir_values,
|
||||
dsl_user_op,
|
||||
)
|
||||
from cutlass._mlir import ir
|
||||
import cutlass.cute as cute
|
||||
from cutlass.utils import WorkTileInfo
|
||||
|
||||
|
||||
class Mamba2SSDTileSchedulerParams:
|
||||
def __init__(
|
||||
self,
|
||||
problem_shape_ntiles: int,
|
||||
eh: int,
|
||||
ngroup_ratio: int,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
):
|
||||
self.problem_shape_ntiles = problem_shape_ntiles
|
||||
self.eh = eh
|
||||
self.ngroup_ratio = ngroup_ratio
|
||||
self._loc = loc
|
||||
|
||||
def __extract_mlir_values__(self):
|
||||
values, self._values_pos = [], []
|
||||
for obj in [self.problem_shape_ntiles, self.eh, self.ngroup_ratio]:
|
||||
obj_values = extract_mlir_values(obj)
|
||||
values += obj_values
|
||||
self._values_pos.append(len(obj_values))
|
||||
return values
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
obj_list = []
|
||||
for obj, n_items in zip(
|
||||
[self.problem_shape_ntiles, self.eh, self.ngroup_ratio], self._values_pos
|
||||
):
|
||||
obj_list.append(new_from_mlir_values(obj, values[:n_items]))
|
||||
values = values[n_items:]
|
||||
return Mamba2SSDTileSchedulerParams(*(tuple(obj_list)), loc=self._loc)
|
||||
|
||||
@dsl_user_op
|
||||
def get_grid_shape(
|
||||
self, max_active_clusters: Int32, *, loc=None, ip=None
|
||||
) -> Tuple[Integer, Integer, Integer]:
|
||||
return (min(self.problem_shape_ntiles, max_active_clusters), 1, 1)
|
||||
|
||||
|
||||
class Mamba2SSDTileScheduler:
|
||||
def __init__(
|
||||
self,
|
||||
params: Mamba2SSDTileSchedulerParams,
|
||||
num_persistent_ctas: Int32,
|
||||
current_work_linear_idx: Int32,
|
||||
num_tiles_executed: Int32,
|
||||
):
|
||||
self.params = params
|
||||
self.num_persistent_ctas = num_persistent_ctas
|
||||
self._current_work_linear_idx = current_work_linear_idx
|
||||
self._num_tiles_executed = num_tiles_executed
|
||||
|
||||
def __extract_mlir_values__(self) -> list[ir.Value]:
|
||||
values = extract_mlir_values(self.num_persistent_ctas)
|
||||
values.extend(extract_mlir_values(self._current_work_linear_idx))
|
||||
values.extend(extract_mlir_values(self._num_tiles_executed))
|
||||
return values
|
||||
|
||||
def __new_from_mlir_values__(
|
||||
self, values: list[ir.Value]
|
||||
) -> "Mamba2SSDTileScheduler":
|
||||
assert len(values) == 3
|
||||
new_num_persistent_ctas = new_from_mlir_values(
|
||||
self.num_persistent_ctas, [values[0]]
|
||||
)
|
||||
new_current_work_linear_idx = new_from_mlir_values(
|
||||
self._current_work_linear_idx, [values[1]]
|
||||
)
|
||||
new_num_tiles_executed = new_from_mlir_values(
|
||||
self._num_tiles_executed, [values[2]]
|
||||
)
|
||||
return Mamba2SSDTileScheduler(
|
||||
self.params,
|
||||
new_num_persistent_ctas,
|
||||
new_current_work_linear_idx,
|
||||
new_num_tiles_executed,
|
||||
)
|
||||
|
||||
# called by host
|
||||
@staticmethod
|
||||
@dsl_user_op
|
||||
def create(
|
||||
params: Mamba2SSDTileSchedulerParams,
|
||||
block_idx: Tuple[Integer, Integer, Integer],
|
||||
grid_dim: Tuple[Integer, Integer, Integer],
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
):
|
||||
params = params
|
||||
|
||||
# Calculate the number of persistent clusters by dividing the total grid size
|
||||
# by the number of CTAs per cluster
|
||||
num_persistent_ctas = Int32(cute.size(grid_dim, loc=loc, ip=ip))
|
||||
|
||||
bidx, bidy, bidz = block_idx
|
||||
|
||||
# Initialize workload index equals to the cluster index in the grid
|
||||
current_work_linear_idx = Int32(bidx)
|
||||
|
||||
# Initialize number of tiles executed to zero
|
||||
num_tiles_executed = Int32(0)
|
||||
return Mamba2SSDTileScheduler(
|
||||
params,
|
||||
num_persistent_ctas,
|
||||
current_work_linear_idx,
|
||||
num_tiles_executed,
|
||||
)
|
||||
|
||||
# called by host
|
||||
@staticmethod
|
||||
def get_grid_shape(
|
||||
params: Mamba2SSDTileSchedulerParams,
|
||||
max_active_clusters: Int32,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Tuple[Integer, Integer, Integer]:
|
||||
return params.get_grid_shape(max_active_clusters, loc=loc, ip=ip)
|
||||
|
||||
# private method
|
||||
def _get_current_work_for_linear_idx(
|
||||
self, current_work_linear_idx: Int32, *, loc=None, ip=None
|
||||
) -> WorkTileInfo:
|
||||
is_valid = current_work_linear_idx < cute.size(
|
||||
self.params.problem_shape_ntiles, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
eh_idx = current_work_linear_idx % self.params.eh
|
||||
b_idx = current_work_linear_idx // self.params.eh
|
||||
g_idx = eh_idx // self.params.ngroup_ratio
|
||||
# cur_tile_coord is (b_idx, eh_idx, g_idx)
|
||||
cur_tile_coord = tuple(Int32(x) for x in (b_idx, eh_idx, g_idx))
|
||||
|
||||
return WorkTileInfo(cur_tile_coord, is_valid)
|
||||
|
||||
@dsl_user_op
|
||||
def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
|
||||
return self._get_current_work_for_linear_idx(
|
||||
self._current_work_linear_idx, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
@dsl_user_op
|
||||
def initial_work_tile_info(self, *, loc=None, ip=None) -> WorkTileInfo:
|
||||
return self.get_current_work(loc=loc, ip=ip)
|
||||
|
||||
@dsl_user_op
|
||||
def advance_to_next_work(self, *, advance_count: int = 1, loc=None, ip=None):
|
||||
self._current_work_linear_idx += Int32(advance_count) * Int32(
|
||||
self.num_persistent_ctas
|
||||
)
|
||||
self._num_tiles_executed += Int32(1)
|
||||
|
||||
@property
|
||||
def num_tiles_executed(self) -> Int32:
|
||||
return self._num_tiles_executed
|
||||
+2096
File diff suppressed because it is too large
Load Diff
+2030
File diff suppressed because it is too large
Load Diff
+2172
File diff suppressed because it is too large
Load Diff
+400
@@ -0,0 +1,400 @@
|
||||
# Copyright (c) 2025 - 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.
|
||||
|
||||
from typing import Tuple, Optional
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
import cutlass.cute.nvgpu.tcgen05 as tcgen05
|
||||
import cutlass.pipeline as pipeline
|
||||
|
||||
|
||||
@cute.jit
|
||||
def load_qk(
|
||||
iterations: int,
|
||||
kv_step: cutlass.Int32,
|
||||
k_args: Tuple,
|
||||
scale_k_args: Optional[Tuple] = None,
|
||||
q_args: Optional[Tuple] = None,
|
||||
) -> Tuple[pipeline.PipelineProducer, pipeline.PipelineProducer]:
|
||||
if cutlass.const_expr(q_args is not None):
|
||||
tQgQ, tQsQ, tma_atom_q, load_q_producer = q_args
|
||||
else:
|
||||
tQgQ, tQsQ, tma_atom_q, load_q_producer = None, None, None, None
|
||||
tKgK, tKsK, tma_atom_k, load_k_producer = k_args
|
||||
tKgScaleK, tKsScaleK, tma_atom_scale_k, load_scale_k_producer = scale_k_args
|
||||
|
||||
scale_k_handle = load_scale_k_producer.acquire_and_advance()
|
||||
cute.copy(
|
||||
tma_atom_scale_k,
|
||||
tKgScaleK[None, kv_step],
|
||||
tKsScaleK[None, scale_k_handle.index],
|
||||
tma_bar_ptr=scale_k_handle.barrier,
|
||||
)
|
||||
for iter in cutlass.range(iterations, unroll=1):
|
||||
if cutlass.const_expr(q_args is not None):
|
||||
q_handle = load_q_producer.acquire_and_advance()
|
||||
cute.copy(
|
||||
tma_atom_q,
|
||||
tQgQ[None, iter],
|
||||
tQsQ[None, q_handle.index],
|
||||
tma_bar_ptr=q_handle.barrier,
|
||||
)
|
||||
k_handle = load_k_producer.acquire_and_advance()
|
||||
cute.copy(
|
||||
tma_atom_k,
|
||||
tKgK[None, kv_step, iter],
|
||||
tKsK[None, k_handle.index],
|
||||
tma_bar_ptr=k_handle.barrier,
|
||||
)
|
||||
if cutlass.const_expr(q_args is not None):
|
||||
return load_k_producer, load_scale_k_producer, load_q_producer
|
||||
else:
|
||||
return load_k_producer, load_scale_k_producer
|
||||
|
||||
|
||||
@cute.jit
|
||||
def load_v(
|
||||
iterations: int,
|
||||
kv_step: cutlass.Int32,
|
||||
v_args: Tuple,
|
||||
scale_v_args: Tuple,
|
||||
) -> pipeline.PipelineProducer:
|
||||
tVgV, tVsV, tma_atom_v, load_v_producer = v_args
|
||||
tScaleVgV, tScaleVsV, tma_atom_scale_v, load_scale_v_producer = scale_v_args
|
||||
scale_v_handle = load_scale_v_producer.acquire_and_advance()
|
||||
cute.copy(
|
||||
tma_atom_scale_v,
|
||||
tScaleVgV[None, kv_step],
|
||||
tScaleVsV[None, scale_v_handle.index],
|
||||
tma_bar_ptr=scale_v_handle.barrier,
|
||||
)
|
||||
for iter in cutlass.range(iterations, unroll=1):
|
||||
v_handle = load_v_producer.acquire_and_advance()
|
||||
cute.copy(
|
||||
tma_atom_v,
|
||||
tVgV[None, iter, kv_step],
|
||||
tVsV[None, v_handle.index],
|
||||
tma_bar_ptr=v_handle.barrier,
|
||||
)
|
||||
return load_v_producer, load_scale_v_producer
|
||||
|
||||
|
||||
@cute.jit
|
||||
def get_scale_smem_layout(
|
||||
scale_granularity: int,
|
||||
d_r: int,
|
||||
mma_tiler: cute.Tile,
|
||||
major_mode: tcgen05.OperandMajorMode,
|
||||
) -> Tuple[cute.Layout, cute.Tile]:
|
||||
size_mn = mma_tiler[1] // 2 # 2cta by default
|
||||
if cutlass.const_expr(major_mode == tcgen05.OperandMajorMode.MN): # v
|
||||
scale_tiler = (mma_tiler[2] * d_r,)
|
||||
tma_view_layout = cute.make_layout(
|
||||
(mma_tiler[2] * d_r),
|
||||
)
|
||||
assert scale_granularity % mma_tiler[1] == 0, (
|
||||
"scale_granularity must be divisible by mma_tiler[1]"
|
||||
)
|
||||
rest_l = scale_granularity // mma_tiler[1]
|
||||
s2r_view_layout = cute.make_layout(
|
||||
(size_mn, mma_tiler[2], (rest_l, d_r)),
|
||||
stride=(0, d_r, (0, 1)),
|
||||
)
|
||||
else: # k
|
||||
scale_tiler = (mma_tiler[1] * d_r,)
|
||||
tma_view_layout = cute.make_layout((size_mn * d_r))
|
||||
assert scale_granularity % mma_tiler[2] == 0, (
|
||||
"scale_granularity must be divisible by mma_tiler[2]"
|
||||
)
|
||||
rest_l = scale_granularity // mma_tiler[2]
|
||||
s2r_view_layout = cute.make_layout(
|
||||
(size_mn, mma_tiler[2], (rest_l, d_r)),
|
||||
stride=(d_r, 0, (0, 1)),
|
||||
)
|
||||
# Apply a trivial swizzle to make it a composed layout, which could be used to construct TMA atom
|
||||
tma_view_smem_layout = cute.make_composed_layout(
|
||||
cute.make_swizzle(0, 4, 3), 0, tma_view_layout
|
||||
)
|
||||
return tma_view_smem_layout, scale_tiler, s2r_view_layout
|
||||
|
||||
|
||||
@cute.jit
|
||||
def mma_qk(
|
||||
iterations: int,
|
||||
qk_tiled_mma: cute.TiledMma,
|
||||
tensor_args: Tuple,
|
||||
pipeline_args: Tuple,
|
||||
):
|
||||
tStS, tSrQ, tSrK_trans = tensor_args
|
||||
mma_s_producer, load_q_consumer, load_q_releaser, dequant_kv_consumer = (
|
||||
pipeline_args
|
||||
)
|
||||
cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster())
|
||||
is_leader_cta = cta_rank_in_cluster % 2 == 0
|
||||
if is_leader_cta:
|
||||
s_handle = mma_s_producer.acquire_and_advance()
|
||||
tStS_slice = tStS[None, None, None, s_handle.index]
|
||||
qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False)
|
||||
for iter in cutlass.range(iterations, unroll=1):
|
||||
if cutlass.const_expr(load_q_consumer is not None):
|
||||
load_q_consumer.wait_and_advance()
|
||||
tSrQ_slice = tSrQ[None, None, None, iter]
|
||||
k_trans_handle = dequant_kv_consumer.wait_and_advance()
|
||||
tSrK_trans_slice = tSrK_trans[None, None, None, k_trans_handle.index]
|
||||
num_kphases = cute.size(tSrQ_slice, mode=[2])
|
||||
for kphase_idx in cutlass.range(num_kphases, unroll_full=True):
|
||||
kphase_coord = (None, None, kphase_idx)
|
||||
cute.gemm(
|
||||
qk_tiled_mma,
|
||||
tStS_slice,
|
||||
tSrQ_slice[kphase_coord],
|
||||
tSrK_trans_slice[kphase_coord],
|
||||
tStS_slice,
|
||||
)
|
||||
qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True)
|
||||
k_trans_handle.release()
|
||||
if cutlass.const_expr(load_q_releaser is not None):
|
||||
load_q_releaser.release()
|
||||
load_q_releaser.advance()
|
||||
s_handle.commit()
|
||||
return mma_s_producer, load_q_consumer, dequant_kv_consumer
|
||||
|
||||
|
||||
@cute.jit
|
||||
def dequant_k(
|
||||
iterations: int,
|
||||
transform_warp_ids: Tuple,
|
||||
dtype_args: Tuple,
|
||||
tensor_args: Tuple,
|
||||
pipeline_args: Tuple,
|
||||
):
|
||||
(k_dtype, q_dtype) = dtype_args
|
||||
(sOrig, sScale, sTrans) = tensor_args
|
||||
(load_kv_consumer, load_scale_consumer, dequant_kv_producer) = pipeline_args
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
THREADS_PER_WARP = 32
|
||||
thread_idx = tidx % (THREADS_PER_WARP * len(transform_warp_ids))
|
||||
r2s_copy_atom = cute.make_copy_atom(
|
||||
cute.nvgpu.CopyUniversalOp(), k_dtype, num_bits_per_copy=32
|
||||
)
|
||||
# Construct tiled_copy satisfying 16 contiguous elts per copy atom
|
||||
r2s_tiled_copy = cute.make_cotiled_copy(
|
||||
r2s_copy_atom,
|
||||
cute.make_layout((256, 16), stride=(16, 1)),
|
||||
sTrans[(None, None, None, 0)].layout,
|
||||
)
|
||||
thr_r2s_tiled_copy = r2s_tiled_copy.get_slice(thread_idx)
|
||||
tOsOrig = thr_r2s_tiled_copy.partition_S(sOrig)
|
||||
tTsTrans = thr_r2s_tiled_copy.partition_D(sTrans)
|
||||
tOrOrig = cute.make_rmem_tensor_like(
|
||||
cute.append(
|
||||
tOsOrig[None, None, None, None, 0].layout,
|
||||
cute.make_layout(
|
||||
2, stride=cute.cosize(tOsOrig[None, None, None, None, 0].layout)
|
||||
),
|
||||
),
|
||||
k_dtype,
|
||||
)
|
||||
tTrTrans = cute.make_rmem_tensor_like(
|
||||
cute.append(
|
||||
tTsTrans[None, None, None, None, 0].layout,
|
||||
cute.make_layout(
|
||||
2, stride=cute.cosize(tTsTrans[None, None, None, None, 0].layout)
|
||||
),
|
||||
),
|
||||
q_dtype,
|
||||
)
|
||||
tSsScale = thr_r2s_tiled_copy.partition_S(sScale)
|
||||
tSrScale = cute.make_rmem_tensor_like(tSsScale[None, None, None, None, None, 0])
|
||||
scale_handle = load_scale_consumer.wait_and_advance()
|
||||
cute.autovec_copy(
|
||||
tSsScale[None, None, None, None, None, scale_handle.index], tSrScale
|
||||
)
|
||||
cute.arch.fence_view_async_shared()
|
||||
scale_handle.release()
|
||||
# prefetch iter = 0
|
||||
kv_handle = load_kv_consumer.wait_and_advance()
|
||||
cute.autovec_copy(
|
||||
tOsOrig[None, None, None, None, kv_handle.index],
|
||||
tOrOrig[None, None, None, None, 0],
|
||||
)
|
||||
transformed_tensor = tOrOrig[None, None, None, None, 0].load().to(q_dtype)
|
||||
scale = cute.TensorSSA(
|
||||
tSrScale[None, None, None, None, 0].load(),
|
||||
transformed_tensor.shape,
|
||||
q_dtype,
|
||||
)
|
||||
transformed_tensor = transformed_tensor * scale
|
||||
tTrTrans[None, None, None, None, 0].store(transformed_tensor)
|
||||
cute.arch.fence_view_async_shared()
|
||||
kv_handle.release()
|
||||
for iter in cutlass.range(1, iterations, unroll_full=True):
|
||||
kv_trans_handle = dequant_kv_producer.acquire_and_advance()
|
||||
cute.autovec_copy(
|
||||
tTrTrans[None, None, None, None, (iter - 1) % 2],
|
||||
tTsTrans[None, None, None, None, kv_trans_handle.index],
|
||||
)
|
||||
cute.arch.fence_view_async_shared()
|
||||
kv_trans_handle.commit()
|
||||
kv_handle = load_kv_consumer.wait_and_advance()
|
||||
cute.autovec_copy(
|
||||
tOsOrig[None, None, None, None, kv_handle.index],
|
||||
tOrOrig[None, None, None, None, iter % 2],
|
||||
)
|
||||
transformed_tensor = (
|
||||
tOrOrig[None, None, None, None, iter % 2].load().to(q_dtype)
|
||||
)
|
||||
scale = cute.TensorSSA(
|
||||
tSrScale[None, None, None, None, iter].load(),
|
||||
transformed_tensor.shape,
|
||||
q_dtype,
|
||||
)
|
||||
transformed_tensor = transformed_tensor * scale
|
||||
tTrTrans[None, None, None, None, iter % 2].store(transformed_tensor)
|
||||
cute.arch.fence_view_async_shared()
|
||||
kv_handle.release()
|
||||
kv_trans_handle = dequant_kv_producer.acquire_and_advance()
|
||||
cute.autovec_copy(
|
||||
tTrTrans[None, None, None, None, (iterations - 1) % 2],
|
||||
tTsTrans[None, None, None, None, kv_trans_handle.index],
|
||||
)
|
||||
cute.arch.fence_view_async_shared()
|
||||
kv_trans_handle.commit()
|
||||
return load_kv_consumer, load_scale_consumer, dequant_kv_producer
|
||||
|
||||
|
||||
@cute.jit
|
||||
def dequant_v(
|
||||
iterations: int,
|
||||
transform_warp_ids: Tuple,
|
||||
dtype_args: Tuple,
|
||||
tensor_args: Tuple,
|
||||
pipeline_args: Tuple,
|
||||
):
|
||||
(v_dtype, q_dtype) = dtype_args
|
||||
(sOrig, sScale, sTrans) = tensor_args
|
||||
(load_kv_consumer, load_scale_consumer, dequant_kv_producer) = pipeline_args
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
THREADS_PER_WARP = 32
|
||||
thread_idx = tidx % (THREADS_PER_WARP * len(transform_warp_ids))
|
||||
r2s_copy_atom = cute.make_copy_atom(
|
||||
cute.nvgpu.CopyUniversalOp(), v_dtype, num_bits_per_copy=32
|
||||
)
|
||||
# Construct tiled_copy satisfying 16 contiguous elts per copy atom
|
||||
r2s_tiled_copy = cute.make_cotiled_copy(
|
||||
r2s_copy_atom,
|
||||
cute.make_layout((256, 16), stride=(16, 1)),
|
||||
sTrans[(None, None, None, 0)].layout,
|
||||
)
|
||||
thr_r2s_tiled_copy = r2s_tiled_copy.get_slice(thread_idx)
|
||||
tOsOrig = thr_r2s_tiled_copy.partition_S(sOrig)
|
||||
tTsTrans = thr_r2s_tiled_copy.partition_D(sTrans)
|
||||
# double buffer for better perf
|
||||
tOrOrig = cute.make_rmem_tensor_like(
|
||||
cute.append(
|
||||
tOsOrig[None, None, None, None, 0].layout,
|
||||
cute.make_layout(
|
||||
2, stride=cute.cosize(tOsOrig[None, None, None, None, 0].layout)
|
||||
),
|
||||
),
|
||||
v_dtype,
|
||||
)
|
||||
tTrTrans = cute.make_rmem_tensor_like(
|
||||
cute.append(
|
||||
tTsTrans[None, None, None, None, 0].layout,
|
||||
cute.make_layout(
|
||||
2, stride=cute.cosize(tTsTrans[None, None, None, None, 0].layout)
|
||||
),
|
||||
),
|
||||
q_dtype,
|
||||
)
|
||||
tSsScale = thr_r2s_tiled_copy.partition_S(sScale)
|
||||
tSrScale = cute.make_rmem_tensor_like(tSsScale[None, None, None, None, None, 0])
|
||||
scale_v_handle = load_scale_consumer.wait_and_advance()
|
||||
cute.autovec_copy(
|
||||
tSsScale[None, None, None, None, None, scale_v_handle.index],
|
||||
tSrScale,
|
||||
)
|
||||
cute.arch.fence_view_async_shared()
|
||||
scale_v_handle.release()
|
||||
# prefetch iter = 0
|
||||
kv_handle = load_kv_consumer.wait_and_advance()
|
||||
cute.autovec_copy(
|
||||
tOsOrig[None, None, None, None, kv_handle.index],
|
||||
tOrOrig[None, None, None, None, 0],
|
||||
)
|
||||
transformed_tensor = tOrOrig[None, None, None, None, 0].load().to(q_dtype)
|
||||
scale = cute.TensorSSA(
|
||||
tSrScale[None, None, None, None, 0].load(),
|
||||
transformed_tensor.shape,
|
||||
q_dtype,
|
||||
)
|
||||
transformed_tensor = transformed_tensor * scale
|
||||
tTrTrans[None, None, None, None, 0].store(transformed_tensor)
|
||||
cute.arch.fence_view_async_shared()
|
||||
kv_handle.release()
|
||||
for iter in cutlass.range(1, iterations, unroll_full=True):
|
||||
kv_trans_handle = dequant_kv_producer.acquire_and_advance()
|
||||
cute.autovec_copy(
|
||||
tTrTrans[None, None, None, None, (iter - 1) % 2],
|
||||
tTsTrans[None, None, None, None, kv_trans_handle.index],
|
||||
)
|
||||
cute.arch.fence_view_async_shared()
|
||||
kv_trans_handle.commit()
|
||||
kv_handle = load_kv_consumer.wait_and_advance()
|
||||
cute.autovec_copy(
|
||||
tOsOrig[None, None, None, None, kv_handle.index],
|
||||
tOrOrig[None, None, None, None, iter % 2],
|
||||
)
|
||||
transformed_tensor = (
|
||||
tOrOrig[None, None, None, None, iter % 2].load().to(q_dtype)
|
||||
)
|
||||
scale = cute.TensorSSA(
|
||||
tSrScale[
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
iter,
|
||||
].load(),
|
||||
transformed_tensor.shape,
|
||||
q_dtype,
|
||||
)
|
||||
transformed_tensor = transformed_tensor * scale
|
||||
tTrTrans[None, None, None, None, iter % 2].store(transformed_tensor)
|
||||
cute.arch.fence_view_async_shared()
|
||||
kv_handle.release()
|
||||
kv_trans_handle = dequant_kv_producer.acquire_and_advance()
|
||||
cute.autovec_copy(
|
||||
tTrTrans[None, None, None, None, (iterations - 1) % 2],
|
||||
tTsTrans[None, None, None, None, kv_trans_handle.index],
|
||||
)
|
||||
cute.arch.fence_view_async_shared()
|
||||
kv_trans_handle.commit()
|
||||
return load_kv_consumer, load_scale_consumer, dequant_kv_producer
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,304 @@
|
||||
# Copyright (c) 2025 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.
|
||||
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
|
||||
|
||||
class MLAStaticTileSchedulerParams:
|
||||
def __init__(
|
||||
self,
|
||||
is_persistent: bool,
|
||||
problem_shape_b: cute.Int32,
|
||||
problem_shape_s: cute.Int32,
|
||||
cluster_shape_mnk: cute.Shape,
|
||||
split_kv: cutlass.Int32,
|
||||
*,
|
||||
problem_shape_b_fdd: cute.FastDivmodDivisor = None,
|
||||
problem_shape_s_fdd: cute.FastDivmodDivisor = None,
|
||||
split_kv_fdd: cute.FastDivmodDivisor = None,
|
||||
loc=None,
|
||||
ip=None,
|
||||
):
|
||||
"""The static tile scheduler parameters prepared for MLA static tile scheduler.
|
||||
|
||||
:param is_persistent: Whether to use persistent kernel mode
|
||||
:type is_persistent: bool
|
||||
:param problem_shape_b: The shape of the problem
|
||||
:type problem_shape_b: cute.Int32
|
||||
:param problem_shape_s: The shape of the problem in sequence length Q dimension
|
||||
:type problem_shape_s: cute.Int32
|
||||
:param cluster_shape_mnk: The shape of the cluster
|
||||
:type cluster_shape_mnk: cute.Shape
|
||||
:param split_kv: The scalar factor for split KV
|
||||
"""
|
||||
self.is_persistent = is_persistent
|
||||
self.problem_shape_b = problem_shape_b
|
||||
self.problem_shape_s = problem_shape_s
|
||||
self.problem_shape_b_fdd = problem_shape_b_fdd
|
||||
self.problem_shape_s_fdd = problem_shape_s_fdd
|
||||
self.cluster_shape_mnk = cluster_shape_mnk
|
||||
self.split_kv = split_kv
|
||||
self.split_kv_fdd = split_kv_fdd
|
||||
if cutlass.const_expr(problem_shape_b_fdd is None):
|
||||
self.problem_shape_b_fdd = cute.fast_divmod_create_divisor(
|
||||
problem_shape_b, loc=loc, ip=ip
|
||||
)
|
||||
if cutlass.const_expr(problem_shape_s_fdd is None):
|
||||
self.problem_shape_s_fdd = cute.fast_divmod_create_divisor(
|
||||
problem_shape_s, loc=loc, ip=ip
|
||||
)
|
||||
if cutlass.const_expr(split_kv_fdd is None):
|
||||
self.split_kv_fdd = cute.fast_divmod_create_divisor(
|
||||
split_kv, loc=loc, ip=ip
|
||||
)
|
||||
self.loc = loc
|
||||
self.ip = ip
|
||||
|
||||
def __extract_mlir_values__(self):
|
||||
values = cutlass.extract_mlir_values(self.problem_shape_b)
|
||||
values += cutlass.extract_mlir_values(self.problem_shape_s)
|
||||
values += cutlass.extract_mlir_values(self.split_kv)
|
||||
values += cutlass.extract_mlir_values(self.problem_shape_b_fdd)
|
||||
values += cutlass.extract_mlir_values(self.problem_shape_s_fdd)
|
||||
values += cutlass.extract_mlir_values(self.split_kv_fdd)
|
||||
return values
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
problem_shape_b = cutlass.new_from_mlir_values(
|
||||
self.problem_shape_b, (values[0],)
|
||||
)
|
||||
problem_shape_s = cutlass.new_from_mlir_values(
|
||||
self.problem_shape_s, (values[1],)
|
||||
)
|
||||
split_kv = cutlass.new_from_mlir_values(self.split_kv, (values[2],))
|
||||
problem_shape_b_fdd = cutlass.new_from_mlir_values(
|
||||
self.problem_shape_b_fdd, (values[3],)
|
||||
)
|
||||
problem_shape_s_fdd = cutlass.new_from_mlir_values(
|
||||
self.problem_shape_s_fdd, (values[4],)
|
||||
)
|
||||
split_kv_fdd = cutlass.new_from_mlir_values(self.split_kv_fdd, (values[5],))
|
||||
return MLAStaticTileSchedulerParams(
|
||||
self.is_persistent,
|
||||
problem_shape_b,
|
||||
problem_shape_s,
|
||||
self.cluster_shape_mnk,
|
||||
split_kv,
|
||||
problem_shape_b_fdd=problem_shape_b_fdd,
|
||||
problem_shape_s_fdd=problem_shape_s_fdd,
|
||||
split_kv_fdd=split_kv_fdd,
|
||||
loc=self.loc,
|
||||
)
|
||||
|
||||
|
||||
def create_mla_static_tile_scheduler_params(
|
||||
is_persistent: bool,
|
||||
problem_shape_b: cute.Int32,
|
||||
problem_shape_s: cute.Int32,
|
||||
cluster_shape_mnk: cute.Shape,
|
||||
split_kv: cutlass.Int32,
|
||||
) -> MLAStaticTileSchedulerParams:
|
||||
return MLAStaticTileSchedulerParams(
|
||||
is_persistent, problem_shape_b, problem_shape_s, cluster_shape_mnk, split_kv
|
||||
)
|
||||
|
||||
|
||||
class WorkTileInfo:
|
||||
def __init__(self, blk_coord: cute.Coord, is_valid: bool):
|
||||
self.blk_coord = blk_coord
|
||||
self.is_valid = cutlass.Boolean(is_valid)
|
||||
|
||||
def __extract_mlir_values__(self):
|
||||
values = cutlass.extract_mlir_values(self.blk_coord)
|
||||
values += cutlass.extract_mlir_values(self.is_valid)
|
||||
return values
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
new_tile_idx = cutlass.new_from_mlir_values(self.blk_coord, values[:-1])
|
||||
new_is_valid_tile = cutlass.new_from_mlir_values(self.is_valid, [values[-1]])
|
||||
return WorkTileInfo(new_tile_idx, new_is_valid_tile)
|
||||
|
||||
@property
|
||||
def is_valid_tile(self) -> cutlass.Boolean:
|
||||
return self.is_valid
|
||||
|
||||
@property
|
||||
def tile_idx(self) -> cute.Coord:
|
||||
return self.blk_coord
|
||||
|
||||
|
||||
class MLAStaticTileScheduler:
|
||||
def __init__(
|
||||
self,
|
||||
params: MLAStaticTileSchedulerParams,
|
||||
current_work_linear_idx: cutlass.Int32,
|
||||
blk_coord: cute.Coord,
|
||||
grid_shape: cute.Shape,
|
||||
*,
|
||||
is_valid: bool = True,
|
||||
loc=None,
|
||||
ip=None,
|
||||
):
|
||||
"""The static tile scheduler for MLA split kv kernel.
|
||||
Based on `is_persistent`, it provides 2 modes for use:
|
||||
- Persistent mode: Launch fixed blocks and reschedule the data blocks.
|
||||
- Non-persistent mode: Launch dynamic blocks and exit when the current work is done.
|
||||
|
||||
:param params: The static tile scheduler parameters
|
||||
:type params: MLAStaticTileSchedulerParams
|
||||
:param current_work_linear_idx: The linear index of the current work
|
||||
:type current_work_linear_idx: cutlass.Int32
|
||||
:param blk_coord: The coordinate of the current work
|
||||
:type blk_coord: cute.Coord
|
||||
:param grid_shape: The shape of the grid
|
||||
:type grid_shape: cute.Shape
|
||||
:param is_valid: Whether the current work is valid
|
||||
:type is_valid: bool
|
||||
"""
|
||||
self.params = params
|
||||
self.blk_coord = blk_coord
|
||||
self.grid_shape = grid_shape
|
||||
self.current_work_linear_idx = current_work_linear_idx
|
||||
if params.is_persistent:
|
||||
self.persistent_blk_layout = cute.make_layout(
|
||||
(
|
||||
params.cluster_shape_mnk[0],
|
||||
params.problem_shape_s,
|
||||
params.problem_shape_b,
|
||||
params.split_kv,
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
self.num_blocks = cute.size(self.persistent_blk_layout, loc=loc, ip=ip)
|
||||
# Used for persistent scheduling
|
||||
self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip)
|
||||
else:
|
||||
self.is_valid = is_valid
|
||||
self.loc = loc
|
||||
self.ip = ip
|
||||
|
||||
@staticmethod
|
||||
def get_grid_shape(
|
||||
params: MLAStaticTileSchedulerParams,
|
||||
max_active_clusters: int,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> cute.Shape:
|
||||
# called by host
|
||||
grid_shape = (
|
||||
params.cluster_shape_mnk[0],
|
||||
params.problem_shape_b * params.problem_shape_s,
|
||||
params.split_kv,
|
||||
)
|
||||
if params.is_persistent:
|
||||
return (
|
||||
cutlass.min(
|
||||
max_active_clusters * cute.size(params.cluster_shape_mnk),
|
||||
cute.size(grid_shape, loc=loc, ip=ip),
|
||||
),
|
||||
1,
|
||||
1,
|
||||
)
|
||||
else:
|
||||
return grid_shape
|
||||
|
||||
def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
|
||||
is_valid = (
|
||||
self.current_work_linear_idx < self.num_blocks
|
||||
if self.params.is_persistent
|
||||
else self.is_valid
|
||||
)
|
||||
|
||||
if self.params.is_persistent:
|
||||
current_work_cluster_batch, cluster_idx = (
|
||||
self.current_work_linear_idx // self.params.cluster_shape_mnk[0],
|
||||
self.current_work_linear_idx % self.params.cluster_shape_mnk[0],
|
||||
)
|
||||
current_work_s_batch, s_idx = divmod(
|
||||
current_work_cluster_batch, self.params.problem_shape_s_fdd
|
||||
)
|
||||
current_work_b_batch, b_idx = divmod(
|
||||
current_work_s_batch, self.params.problem_shape_b_fdd
|
||||
)
|
||||
_, split_kv_idx = divmod(current_work_b_batch, self.params.split_kv_fdd)
|
||||
|
||||
blk_coord = (cluster_idx, s_idx, b_idx, split_kv_idx)
|
||||
else:
|
||||
s_idx, b_idx = divmod(self.blk_coord[1], self.params.problem_shape_b_fdd)
|
||||
blk_coord = (self.blk_coord[0], s_idx, b_idx, self.blk_coord[2])
|
||||
|
||||
return WorkTileInfo(blk_coord, is_valid)
|
||||
|
||||
def initial_work_tile_info(self, *, loc=None, ip=None):
|
||||
return self.get_current_work(loc=loc, ip=ip)
|
||||
|
||||
def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None):
|
||||
if self.params.is_persistent:
|
||||
self.current_work_linear_idx += advance_count * self.num_persistent_sm
|
||||
else:
|
||||
self.is_valid = False
|
||||
|
||||
def __extract_mlir_values__(self):
|
||||
values = cutlass.extract_mlir_values(self.params)
|
||||
values.extend(cutlass.extract_mlir_values(self.current_work_linear_idx))
|
||||
values.extend(cutlass.extract_mlir_values(self.blk_coord))
|
||||
values.extend(cutlass.extract_mlir_values(self.grid_shape))
|
||||
return values
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
assert len(values) == 13
|
||||
new_params = cutlass.new_from_mlir_values(self.params, values[0:6])
|
||||
new_current_work_linear_idx = cutlass.new_from_mlir_values(
|
||||
self.current_work_linear_idx, [values[6]]
|
||||
)
|
||||
new_blk_coord = cutlass.new_from_mlir_values(self.blk_coord, values[7:10])
|
||||
new_grid_shape = cutlass.new_from_mlir_values(self.grid_shape, values[10:])
|
||||
return MLAStaticTileScheduler(
|
||||
new_params, new_current_work_linear_idx, new_blk_coord, new_grid_shape
|
||||
)
|
||||
|
||||
|
||||
def create_mla_static_tile_scheduler(
|
||||
params: MLAStaticTileSchedulerParams,
|
||||
blk_coord: cute.Coord,
|
||||
grid_shape: cute.Shape,
|
||||
) -> MLAStaticTileScheduler:
|
||||
return MLAStaticTileScheduler(params, blk_coord[0], blk_coord, grid_shape)
|
||||
|
||||
|
||||
LOG2_E = 1.4426950408889634074
|
||||
# avoid register indexing on array.
|
||||
MAX_SPLITS = 256
|
||||
|
||||
|
||||
def ceil_div(a: int, b: int) -> int:
|
||||
return (a + b - 1) // b
|
||||
Reference in New Issue
Block a user