Skip to content

Commit 609d1d0

Browse files
syuonisyuoni
authored
[None][fix] Fix Illegal Memory Access for CuteDSL Grouped GEMM (#10008)
Signed-off-by: Enwei Zhu <[email protected]>
1 parent 6a238ca commit 609d1d0

File tree

6 files changed

+165
-174
lines changed

6 files changed

+165
-174
lines changed

‎tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_gather_grouped_gemm_swiglu_fusion.py‎

Lines changed: 28 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -273,7 +273,7 @@ def __init__(
273273
mma_tiler_mn: Tuple[int, int],
274274
cluster_shape_mn: Tuple[int, int],
275275
vectorized_f32: bool,
276-
topk: int,
276+
topk: cutlass.Int64,
277277
):
278278
"""Initializes the configuration for a Blackwell blockscaled dense GEMM kernel with
279279
gather operation and SwiGLU fusion.
@@ -310,7 +310,7 @@ def __init__(
310310
:param vectorized_f32: Enable vectorized f32x2 operations for better performance.
311311
:type vectorized_f32: bool
312312
:param topk: Number of experts selected per token (used for token ID mapping).
313-
:type topk: int
313+
:type topk: cutlass.Int64
314314
"""
315315

316316
self.sf_vec_size = sf_vec_size
@@ -1083,7 +1083,7 @@ def kernel(
10831083
gC_mnl = cute.local_tile(
10841084
mC_mnl, cute.slice_(self.mma_tiler_c, (None, None, 0)), (None, None, None)
10851085
)
1086-
k_tile_cnt = cute.size(gA_mkl, mode=[3])
1086+
k_tile_cnt = cutlass.Int32(cute.size(gA_mkl, mode=[3]))
10871087

10881088
#
10891089
# Partition global tensor for TiledMMA_A/B/C
@@ -2738,7 +2738,7 @@ def is_valid_mma_tiler_and_cluster_shape(
27382738
use_2cta_instrs: bool,
27392739
mma_tiler_mn: Tuple[int, int],
27402740
cluster_shape_mn: Tuple[int, int],
2741-
m_aligned: int,
2741+
m_aligned: cutlass.Int64,
27422742
) -> bool:
27432743
"""
27442744
Check if the mma tiler and cluster shape are valid
@@ -2750,7 +2750,7 @@ def is_valid_mma_tiler_and_cluster_shape(
27502750
:param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
27512751
:type cluster_shape_mn: Tuple[int, int]
27522752
:param m_aligned: The alignment requirement for group M dimension (default: 128)
2753-
:type m_aligned: int
2753+
:type m_aligned: cutlass.Int64
27542754
27552755
:return: True if the mma tiler and cluster shape are valid, False otherwise
27562756
:rtype: bool
@@ -2803,10 +2803,10 @@ def is_valid_mma_tiler_and_cluster_shape(
28032803

28042804
@staticmethod
28052805
def is_valid_tensor_alignment(
2806-
m: int,
2807-
n: int,
2808-
k: int,
2809-
l: int, # noqa: E741
2806+
m: cutlass.Int64,
2807+
n: cutlass.Int64,
2808+
k: cutlass.Int64,
2809+
l: cutlass.Int64, # noqa: E741
28102810
ab_dtype: Type[cutlass.Numeric],
28112811
c_dtype: Type[cutlass.Numeric],
28122812
a_major: str,
@@ -2817,13 +2817,13 @@ def is_valid_tensor_alignment(
28172817
Check if the tensor alignment is valid
28182818
28192819
:param m: The number of rows in the A tensor
2820-
:type m: int
2820+
:type m: cutlass.Int64
28212821
:param n: The number of columns in the B tensor
2822-
:type n: int
2822+
:type n: cutlass.Int64
28232823
:param k: The number of columns in the A tensor
2824-
:type k: int
2824+
:type k: cutlass.Int64
28252825
:param l: The number of columns in the C tensor
2826-
:type l: int
2826+
:type l: cutlass.Int64
28272827
:param ab_dtype: The data type of the A and B operands
28282828
:type ab_dtype: Type[cutlass.Numeric]
28292829
:param c_dtype: The data type of the output tensor
@@ -2863,14 +2863,14 @@ def can_implement(
28632863
c_dtype: Type[cutlass.Numeric],
28642864
mma_tiler_mn: Tuple[int, int],
28652865
cluster_shape_mn: Tuple[int, int],
2866-
m: int,
2867-
n: int,
2868-
k: int,
2869-
l: int, # noqa: E741
2866+
m: cutlass.Int64,
2867+
n: cutlass.Int64,
2868+
k: cutlass.Int64,
2869+
l: cutlass.Int64, # noqa: E741
28702870
a_major: str,
28712871
b_major: str,
28722872
c_major: str,
2873-
m_aligned: int,
2873+
m_aligned: cutlass.Int64,
28742874
) -> bool:
28752875
"""
28762876
Check if the gemm can be implemented
@@ -2892,21 +2892,21 @@ def can_implement(
28922892
:param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
28932893
:type cluster_shape_mn: Tuple[int, int]
28942894
:param m: The number of rows in the A tensor
2895-
:type m: int
2895+
:type m: cutlass.Int64
28962896
:param n: The number of columns in the B tensor
2897-
:type n: int
2897+
:type n: cutlass.Int64
28982898
:param k: The number of columns in the A tensor
2899-
:type k: int
2899+
:type k: cutlass.Int64
29002900
:param l: The number of columns in the C tensor
2901-
:type l: int
2901+
:type l: cutlass.Int64
29022902
:param a_major: The major axis of the A tensor
29032903
:type a_major: str
29042904
:param b_major: The major axis of the B tensor
29052905
:type b_major: str
29062906
:param c_major: The major axis of the C tensor
29072907
:type c_major: str
29082908
:param m_aligned: The alignment requirement for group M dimension (default: 128)
2909-
:type m_aligned: int
2909+
:type m_aligned: cutlass.Int64
29102910
29112911
:return: True if the gemm can be implemented, False otherwise
29122912
:rtype: bool
@@ -2955,11 +2955,11 @@ def wrapper(
29552955
token_id_mapping_ptr: cute.Pointer,
29562956
num_non_exiting_tiles_ptr: cute.Pointer,
29572957
global_sf_ptr: cute.Pointer,
2958-
orig_m: int,
2959-
m: int,
2960-
n: int,
2961-
k: int,
2962-
l: int, # noqa: E741
2958+
orig_m: cutlass.Int64,
2959+
m: cutlass.Int64,
2960+
n: cutlass.Int64,
2961+
k: cutlass.Int64,
2962+
l: cutlass.Int64, # noqa: E741
29632963
tile_size: cutlass.Constexpr,
29642964
scaling_vector_size: cutlass.Constexpr,
29652965
max_active_clusters: cutlass.Constexpr,

‎tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm.py‎

Lines changed: 25 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -837,7 +837,7 @@ def kernel(
837837
gC_mnl = cute.local_tile(
838838
mC_mnl, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None)
839839
)
840-
k_tile_cnt = cute.size(gA_mkl, mode=[3])
840+
k_tile_cnt = cutlass.Int32(cute.size(gA_mkl, mode=[3]))
841841

842842
#
843843
# Partition global tensor for TiledMMA_A/B/C
@@ -2021,7 +2021,7 @@ def is_valid_mma_tiler_and_cluster_shape(
20212021
use_2cta_instrs: bool,
20222022
mma_tiler_mn: Tuple[int, int],
20232023
cluster_shape_mn: Tuple[int, int],
2024-
m_aligned: int,
2024+
m_aligned: cutlass.Int64,
20252025
) -> bool:
20262026
"""
20272027
Check if the mma tiler and cluster shape are valid
@@ -2033,7 +2033,7 @@ def is_valid_mma_tiler_and_cluster_shape(
20332033
:param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
20342034
:type cluster_shape_mn: Tuple[int, int]
20352035
:param m_aligned: The alignment requirement for group M dimension (default: 128)
2036-
:type m_aligned: int
2036+
:type m_aligned: cutlass.Int64
20372037
20382038
:return: True if the mma tiler and cluster shape are valid, False otherwise
20392039
:rtype: bool
@@ -2086,10 +2086,10 @@ def is_valid_mma_tiler_and_cluster_shape(
20862086

20872087
@staticmethod
20882088
def is_valid_tensor_alignment(
2089-
m: int,
2090-
n: int,
2091-
k: int,
2092-
l: int, # noqa: E741
2089+
m: cutlass.Int64,
2090+
n: cutlass.Int64,
2091+
k: cutlass.Int64,
2092+
l: cutlass.Int64, # noqa: E741
20932093
ab_dtype: Type[cutlass.Numeric],
20942094
c_dtype: Type[cutlass.Numeric],
20952095
a_major: str,
@@ -2100,13 +2100,13 @@ def is_valid_tensor_alignment(
21002100
Check if the tensor alignment is valid
21012101
21022102
:param m: The number of rows in the A tensor
2103-
:type m: int
2103+
:type m: cutlass.Int64
21042104
:param n: The number of columns in the B tensor
2105-
:type n: int
2105+
:type n: cutlass.Int64
21062106
:param k: The number of columns in the A tensor
2107-
:type k: int
2107+
:type k: cutlass.Int64
21082108
:param l: The number of columns in the C tensor
2109-
:type l: int
2109+
:type l: cutlass.Int64
21102110
:param ab_dtype: The data type of the A and B operands
21112111
:type ab_dtype: Type[cutlass.Numeric]
21122112
:param c_dtype: The data type of the output tensor
@@ -2148,14 +2148,14 @@ def can_implement(
21482148
use_2cta_instrs: bool,
21492149
mma_tiler_mn: Tuple[int, int],
21502150
cluster_shape_mn: Tuple[int, int],
2151-
m: int,
2152-
n: int,
2153-
k: int,
2154-
l: int, # noqa: E741
2151+
m: cutlass.Int64,
2152+
n: cutlass.Int64,
2153+
k: cutlass.Int64,
2154+
l: cutlass.Int64, # noqa: E741
21552155
a_major: str,
21562156
b_major: str,
21572157
c_major: str,
2158-
m_aligned: int,
2158+
m_aligned: cutlass.Int64,
21592159
) -> bool:
21602160
"""
21612161
Check if the gemm can be implemented
@@ -2177,21 +2177,21 @@ def can_implement(
21772177
:param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
21782178
:type cluster_shape_mn: Tuple[int, int]
21792179
:param m: The number of rows in the A tensor
2180-
:type m: int
2180+
:type m: cutlass.Int64
21812181
:param n: The number of columns in the B tensor
2182-
:type n: int
2182+
:type n: cutlass.Int64
21832183
:param k: The number of columns in the A tensor
2184-
:type k: int
2184+
:type k: cutlass.Int64
21852185
:param l: The number of columns in the C tensor
2186-
:type l: int
2186+
:type l: cutlass.Int64
21872187
:param a_major: The major axis of the A tensor
21882188
:type a_major: str
21892189
:param b_major: The major axis of the B tensor
21902190
:type b_major: str
21912191
:param c_major: The major axis of the C tensor
21922192
:type c_major: str
21932193
:param m_aligned: The alignment requirement for group M dimension (default: 128)
2194-
:type m_aligned: int
2194+
:type m_aligned: cutlass.Int64
21952195
21962196
:return: True if the gemm can be implemented, False otherwise
21972197
:rtype: bool
@@ -2233,10 +2233,10 @@ def wrapper(
22332233
alpha_ptr: cute.Pointer,
22342234
tile_idx_to_group_idx_ptr: cute.Pointer,
22352235
num_non_exiting_tiles_ptr: cute.Pointer,
2236-
m: int,
2237-
n: int,
2238-
k: int,
2239-
l: int, # noqa: E741
2236+
m: cutlass.Int64,
2237+
n: cutlass.Int64,
2238+
k: cutlass.Int64,
2239+
l: cutlass.Int64, # noqa: E741
22402240
tile_size: cutlass.Constexpr,
22412241
scaling_vector_size: cutlass.Constexpr,
22422242
max_active_clusters: cutlass.Constexpr,

‎tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py‎

Lines changed: 27 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -998,7 +998,7 @@ def kernel(
998998
(None, None, None),
999999
)
10001000

1001-
k_tile_cnt = cute.size(gA_mkl, mode=[3])
1001+
k_tile_cnt = cutlass.Int32(cute.size(gA_mkl, mode=[3]))
10021002

10031003
#
10041004
# Partition global tensor for TiledMMA_A/B
@@ -2030,7 +2030,7 @@ def is_valid_mma_tiler_and_cluster_shape(
20302030
use_2cta_instrs: bool,
20312031
mma_tiler_mn: Tuple[int, int],
20322032
cluster_shape_mn: Tuple[int, int],
2033-
m_aligned: int,
2033+
m_aligned: cutlass.Int64,
20342034
) -> bool:
20352035
"""
20362036
Check if the mma tiler and cluster shape are valid
@@ -2042,7 +2042,7 @@ def is_valid_mma_tiler_and_cluster_shape(
20422042
:param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
20432043
:type cluster_shape_mn: Tuple[int, int]
20442044
:param m_aligned: The alignment requirement for group M dimension (default: 128)
2045-
:type m_aligned: int
2045+
:type m_aligned: cutlass.Int64
20462046
20472047
:return: True if the mma tiler and cluster shape are valid, False otherwise
20482048
:rtype: bool
@@ -2095,10 +2095,10 @@ def is_valid_mma_tiler_and_cluster_shape(
20952095

20962096
@staticmethod
20972097
def is_valid_tensor_alignment(
2098-
m: int,
2099-
n: int,
2100-
k: int,
2101-
l: int, # noqa: E741
2098+
m: cutlass.Int64,
2099+
n: cutlass.Int64,
2100+
k: cutlass.Int64,
2101+
l: cutlass.Int64, # noqa: E741
21022102
ab_dtype: Type[cutlass.Numeric],
21032103
out_dtype: Type[cutlass.Numeric],
21042104
a_major: str,
@@ -2109,13 +2109,13 @@ def is_valid_tensor_alignment(
21092109
Check if the tensor alignment is valid
21102110
21112111
:param m: The number of rows in the A tensor
2112-
:type m: int
2112+
:type m: cutlass.Int64
21132113
:param n: The number of columns in the B tensor
2114-
:type n: int
2114+
:type n: cutlass.Int64
21152115
:param k: The number of columns in the A tensor
2116-
:type k: int
2116+
:type k: cutlass.Int64
21172117
:param l: The number of columns in the C tensor
2118-
:type l: int
2118+
:type l: cutlass.Int64
21192119
:param ab_dtype: The data type of the A and B operands
21202120
:type ab_dtype: Type[cutlass.Numeric]
21212121
:param out_dtype: The data type of the output tensor
@@ -2157,14 +2157,14 @@ def can_implement(
21572157
use_2cta_instrs: bool,
21582158
mma_tiler_mn: Tuple[int, int],
21592159
cluster_shape_mn: Tuple[int, int],
2160-
m: int,
2161-
n: int,
2162-
k: int,
2163-
l: int, # noqa: E741
2160+
m: cutlass.Int64,
2161+
n: cutlass.Int64,
2162+
k: cutlass.Int64,
2163+
l: cutlass.Int64, # noqa: E741
21642164
a_major: str,
21652165
b_major: str,
21662166
c_major: str,
2167-
m_aligned: int,
2167+
m_aligned: cutlass.Int64,
21682168
) -> bool:
21692169
"""
21702170
Check if the gemm can be implemented
@@ -2186,21 +2186,21 @@ def can_implement(
21862186
:param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster
21872187
:type cluster_shape_mn: Tuple[int, int]
21882188
:param m: The number of rows in the A tensor
2189-
:type m: int
2189+
:type m: cutlass.Int64
21902190
:param n: The number of columns in the B tensor
2191-
:type n: int
2191+
:type n: cutlass.Int64
21922192
:param k: The number of columns in the A tensor
2193-
:type k: int
2193+
:type k: cutlass.Int64
21942194
:param l: The number of columns in the C tensor
2195-
:type l: int
2195+
:type l: cutlass.Int64
21962196
:param a_major: The major axis of the A tensor
21972197
:type a_major: str
21982198
:param b_major: The major axis of the B tensor
21992199
:type b_major: str
22002200
:param c_major: The major axis of the C tensor
22012201
:type c_major: str
22022202
:param m_aligned: The alignment requirement for group M dimension (default: 128)
2203-
:type m_aligned: int
2203+
:type m_aligned: cutlass.Int64
22042204
22052205
:return: True if the gemm can be implemented, False otherwise
22062206
:rtype: bool
@@ -2245,12 +2245,12 @@ def wrapper(
22452245
permuted_idx_to_expanded_idx_ptr: cute.Pointer,
22462246
num_non_exiting_tiles_ptr: cute.Pointer,
22472247
token_final_scales_ptr: cute.Pointer,
2248-
m: int,
2249-
n: int,
2250-
k: int,
2251-
l: int, # noqa: E741
2252-
num_tokens: int,
2253-
top_k: int,
2248+
m: cutlass.Int64,
2249+
n: cutlass.Int64,
2250+
k: cutlass.Int64,
2251+
l: cutlass.Int64, # noqa: E741
2252+
num_tokens: cutlass.Int64,
2253+
top_k: cutlass.Int64,
22542254
tile_size: cutlass.Constexpr,
22552255
scaling_vector_size: cutlass.Constexpr,
22562256
max_active_clusters: cutlass.Constexpr,

0 commit comments

Comments
 (0)