Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
SIYIXNI
vllm
Commits
f8ecb84c
"...composable_kernel_rocm.git" did not exist on "07f818cb17d1c264f69999077df49475949cee97"
Unverified
Commit
f8ecb84c
authored
Jan 27, 2024
by
Woosuk Kwon
Committed by
GitHub
Jan 27, 2024
Browse files
Speed up Punica compilation (#2632)
parent
5f036d2b
Changes
21
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
27 additions
and
0 deletions
+27
-0
csrc/punica/bgmv/generator.py
csrc/punica/bgmv/generator.py
+27
-0
No files found.
csrc/punica/bgmv/generator.py
0 → 100644
View file @
f8ecb84c
DTYPES
=
[
"fp16"
,
"bf16"
,
"fp32"
]
DTYPE_MAP
=
{
"fp16"
:
"nv_half"
,
"bf16"
:
"nv_bfloat16"
,
"fp32"
:
"float"
,
}
TEMPLATE
=
"""
#include "bgmv_config.h"
#include "bgmv_impl.cuh"
FOR_BGMV_WIDE_NARROW(INST_BGMV_TWOSIDE, {input_dtype}, {output_dtype}, {weight_dtype})
"""
.
lstrip
()
for
input_dtype
in
DTYPES
:
for
output_dtype
in
DTYPES
:
for
weight_dtype
in
DTYPES
:
if
weight_dtype
==
"fp32"
:
# FP32 weights are not supported.
continue
kernel_definition
=
TEMPLATE
.
format
(
input_dtype
=
DTYPE_MAP
[
input_dtype
],
output_dtype
=
DTYPE_MAP
[
output_dtype
],
weight_dtype
=
DTYPE_MAP
[
weight_dtype
])
filename
=
f
"bgmv_
{
input_dtype
}
_
{
output_dtype
}
_
{
weight_dtype
}
.cu"
with
open
(
filename
,
"w"
)
as
f
:
f
.
write
(
kernel_definition
)
Prev
1
2
Next
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment