mirror of
https://github.com/vllm-project/vllm.git
synced 2026-10-28 09:45:26 +08:00
[Bugfix] Support CuTe DSL 4.8.0 block-scale API (#59480)
Signed-off-by: hjjq <50634613+hjjq@users.noreply.github.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Shang Wang <samshang.wang@mail.utoronto.ca>
This commit is contained in:
1 file changed
+6
-1
@@ -148,6 +148,11 @@ def mma_mxfp8(
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> None:
|
||||
# CuTe DSL 4.8 renamed the block-scale enum and keyword.
|
||||
if hasattr(nvvm, "Tcgen05MMABlockScale"):
|
||||
scale_kwargs = {"block_scale": nvvm.Tcgen05MMABlockScale.DEFAULT}
|
||||
else:
|
||||
scale_kwargs = {"scale_vec_size": nvvm.Tcgen05MMAScaleVecSize.X1}
|
||||
with cute.arch.elect_one():
|
||||
nvvm.tcgen05_mma_block_scale(
|
||||
nvvm.Tcgen05MMAKind.MXF8F6F4,
|
||||
@@ -159,9 +164,9 @@ def mma_mxfp8(
|
||||
Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
|
||||
_make_tmem_llvm_ptr(sfa_tmem, loc=loc, ip=ip),
|
||||
_make_tmem_llvm_ptr(sfb_tmem, loc=loc, ip=ip),
|
||||
scale_vec_size=nvvm.Tcgen05MMAScaleVecSize.X1,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
**scale_kwargs,
|
||||
)
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user