minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[jit_kernel] Migrate cast (downcast_fp8) from sgl-kernel AOT to JIT (#19103)

Browse Source
This commit is contained in:
Johnsonms
2026-03-27 13:21:44 +08:00
committed by GitHub
parent c531be455e
commit 8a56a7b04d
6 changed files with 638 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/include/sgl_kernel/type.cuh
+20
View File
@@ -77,6 +77,10 @@ SGL_REGISTER_DTYPE_TRAIT(
SGL_REGISTER_DTYPE_TRAIT(
bf16x2_t, void, SGL_REGISTER_TYPE_END; SGL_REGISTER_FROM_FUNCTION(fp32x2_t, __float22bfloat162_rn););
#ifndef USE_ROCM
SGL_REGISTER_DTYPE_TRAIT(fp8_e4m3_t, fp8x2_e4m3_t);
#endif
#undef SGL_REGISTER_DTYPE_TRAIT
#undef SGL_REGISTER_FROM_FUNCTION
@@ -98,3 +102,19 @@ SGL_DEVICE To cast(const From& value) {
}
} // namespace device
// ---------------------------------------------------------------------------
// FP8 max clamp value — platform-dependent
// CUDA (e4m3fn): 448.0f
// AMD FNUZ (e4m3fnuz): 224.0f
// AMD E4M3 (e4m3fn): 448.0f
// ---------------------------------------------------------------------------
#ifndef USE_ROCM
constexpr float kFP8E4M3Max = 448.0f;
#else // USE_ROCM
#if HIP_FP8_TYPE_FNUZ
constexpr float kFP8E4M3Max = 224.0f;
#else // HIP_FP8_TYPE_E4M3
constexpr float kFP8E4M3Max = 448.0f;
#endif // HIP_FP8_TYPE_FNUZ
#endif // USE_ROCM