A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 13:33 UTC
[AMD][Kernel][BugFix] Use correct scale in concat_and_cache_ds_mla_kernel when on gfx942 (#32976)
Signed-off-by: Randall Smith <ransmith@amd.com> Signed-off-by: Randall Smith <Randall.Smith@amd.com> Co-authored-by: Randall Smith <ransmith@amd.com>
rasmith committed
Jan 27, 2026 at 01:16 UTC
58996f3589434d99c320e6ee2460a231135f9641
1 file changed
+9
-7
csrc/cache_kernels.cu
+9
-7
@@ -24,6 +24,12 @@
24
typedef __hip_bfloat16 __nv_bfloat16;
25
#endif
26
27
+#if defined(__gfx942__)
28
+constexpr float kFp8ScaleDivisor = 224.f;
29
+#else
30
+constexpr float kFp8ScaleDivisor = 448.f;
31
+#endif
32
+
33
void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
34
int64_t block_size_in_bytes,
35
const torch::Tensor& block_mapping) {
@@ -401,8 +407,7 @@ __global__ void concat_and_cache_ds_mla_kernel(
407
}
408
409
// Compute the scale for the tile
404
- float tile_scale = max_abs / 448.f;
405
- tile_scale = fmaxf(tile_scale, FLT_MIN);
410
+ float tile_scale = fmaxf(max_abs / kFp8ScaleDivisor, FLT_MIN);
411
412
// The first lane of each half-warp writes the scale to kv_cache
413
if ((lane_idx == 0) || (lane_idx == 16)) {
@@ -471,11 +476,8 @@ __global__ void indexer_k_quant_and_cache_kernel(
476
#endif
477
}
478
474
-#if defined(__gfx942__)
475
- float scale = fmaxf(amax, 1e-4) / 224.0f;
476
-#else
477
- float scale = fmaxf(amax, 1e-4) / 448.0f;
478
-#endif
479
+ float scale = fmaxf(amax, 1e-4) / kFp8ScaleDivisor;
480
+
481
if (use_ue8m0) {
482
scale = exp2f(ceilf(log2f(scale)));
483
}