diff --git a/mslk/attention/fmha/_triton/splitk_kernels.py b/mslk/attention/fmha/_triton/splitk_kernels.py index 134ef934..98fa19f6 100644 --- a/mslk/attention/fmha/_triton/splitk_kernels.py +++ b/mslk/attention/fmha/_triton/splitk_kernels.py @@ -11,6 +11,7 @@ import torch import triton import triton.language as tl +from mslk.utils.device import is_gfx950, rocm_version_at_least from .vararg_kernel import unroll_varargs, VAR_ARGS_ARRAY @@ -121,6 +122,7 @@ def _fwd_kernel_splitK( # noqa: C901 HAS_ADDITIVE_BIAS: tl.constexpr, NUM_PROGRAMS_DIM2_CONST: tl.constexpr, IS_HIP: tl.constexpr, + FP8_FNUZ: tl.constexpr, QUANTIZE_PV_TO_FP8: tl.constexpr, QUANTIZE_QK_TO_FP8: tl.constexpr, USE_FP32_SCALES: tl.constexpr, @@ -539,6 +541,7 @@ def _fwd_kernel_splitK( # noqa: C901 # pyrefly: ignore [bad-argument-type] i, IS_HIP, + FP8_FNUZ, QUANTIZE_PV_TO_FP8, QUANTIZE_QK_TO_FP8, USE_FP32_SCALES, @@ -780,10 +783,16 @@ def autotune_kernel(kernel: Callable): if block_n >= block_m ] + # HIP graph capture faulted (HSA_INVALID_PACKET / GPU faults) during autotuning of + # this kernel on gfx950 with ROCm < 7.14; fixed in ROCm 7.14. Disable graph-based + # autotuning only on that affected stack; gfx942, CUDA, and gfx950 on ROCm >= 7.14 + # keep it. + graph_autotune_broken = is_gfx950() and not rocm_version_at_least(7, 14) + kernel = triton.autotune( configs=TRITON_CONFIGS, key=AUTOTUNER_KEY, - use_cuda_graph=True, + use_cuda_graph=not graph_autotune_broken, prune_configs_by={ "early_config_prune": early_config_prune, }, @@ -834,6 +843,7 @@ def load_dequantize_k_v_group( v_dtype: tl.constexpr, # Q.dtype.element_ty group_id: tl.constexpr, IS_HIP: tl.constexpr, + FP8_FNUZ: tl.constexpr, QUANTIZE_PV_TO_FP8: tl.constexpr, QUANTIZE_QK_TO_FP8: tl.constexpr, USE_FP32_SCALES: tl.constexpr, @@ -874,6 +884,7 @@ def load_dequantize_k_v_group( q_dtype, v_dtype, IS_HIP, + FP8_FNUZ, QUANTIZE_PV_TO_FP8, QUANTIZE_QK_TO_FP8, USE_FP32_SCALES, @@ -921,6 +932,7 @@ def _process_fp8_quantization( q_dtype: tl.constexpr, v_dtype: tl.constexpr, IS_HIP: tl.constexpr, + FP8_FNUZ: tl.constexpr, QUANTIZE_PV_TO_FP8: tl.constexpr, QUANTIZE_QK_TO_FP8: tl.constexpr, USE_FP32_SCALES: tl.constexpr, @@ -938,6 +950,7 @@ def _process_fp8_quantization( v_shift if not USE_FP32_SCALES else None, PACKED_PER_VAL, IS_HIP, + FP8_FNUZ, USE_FP32_SCALES, ).to(v_dtype) else: @@ -951,7 +964,9 @@ def _process_fp8_quantization( k_scale, k_shift = _extract_scale_shift(k_scale_shift, IS_HIP, USE_FP32_SCALES) if IS_HIP: if not QUANTIZE_QK_TO_FP8: - k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL).to(q_dtype) + k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL, FP8_FNUZ).to( + q_dtype + ) else: # For QUANTIZE_QK_TO_FP8, unpack int32 to 8-bit entries and interpret as fp8 tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values") @@ -964,6 +979,7 @@ def _process_fp8_quantization( tl.trans(k_shift) if not USE_FP32_SCALES else None, PACKED_PER_VAL, IS_HIP, + FP8_FNUZ, USE_FP32_SCALES, ).to(q_dtype) k = tl.trans(k_t) @@ -971,7 +987,7 @@ def _process_fp8_quantization( # For QUANTIZE_QK_TO_FP8, unpack int32 to 8-bit entries and interpret as fp8 tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values") k_t = tl.trans(k) - k_t = _unpack_fp8_tensor(k_t, PACKED_PER_VAL, IS_HIP) + k_t = _unpack_fp8_tensor(k_t, PACKED_PER_VAL, IS_HIP, FP8_FNUZ) k = tl.trans(k_t) return k, v, v_scale, k_scale @@ -991,7 +1007,9 @@ def _extract_scale_shift( @triton.jit -def _unpack_fp8_tensor(x_, PACKED_PER_VAL: tl.constexpr, IS_HIP: tl.constexpr): +def _unpack_fp8_tensor( + x_, PACKED_PER_VAL: tl.constexpr, IS_HIP: tl.constexpr, FP8_FNUZ: tl.constexpr +): """Unpack FP8 K/V tensor from int32 packed representation.""" tl.static_assert(PACKED_PER_VAL == 4, "Assert: int32 packs four FP8 values") @@ -1006,8 +1024,9 @@ def _unpack_fp8_tensor(x_, PACKED_PER_VAL: tl.constexpr, IS_HIP: tl.constexpr): unpacked_values, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL) ) - # Convert to FP8 through bitcast - fp8_type = tl.float8e4b8 if IS_HIP else tl.float8e4nv + # Convert to FP8 through bitcast. gfx942 uses e4m3fnuz (float8e4b8); gfx950 and + # CUDA use OCP e4m3fn (float8e4nv). FP8_FNUZ carries the arch decision. + fp8_type = tl.float8e4b8 if FP8_FNUZ else tl.float8e4nv x_ = unpacked_values.to(tl.uint8).to(fp8_type, bitcast=True) return x_ @@ -1039,18 +1058,20 @@ def _process_int4_quantization( if IS_HIP: k_scale, k_shift = cast_uint32_to_float(k_scale_shift) v_scale, v_shift = cast_uint32_to_float(v_scale_shift) - v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP).to(dtype) - k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL).to(dtype) + # int4 path never reaches the fp8 branch inside dequantize; FP8_FNUZ is unused. + v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP, False).to(dtype) + k = dequantize_k_hip(k, k_scale, k_shift, PACKED_PER_VAL, False).to(dtype) else: k_scale, k_shift = cast_uint32_to_half2(k_scale_shift) v_scale, v_shift = cast_uint32_to_half2(v_scale_shift) - v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP).to(dtype) + v = dequantize(v, v_scale, v_shift, PACKED_PER_VAL, IS_HIP, False).to(dtype) k_t = dequantize( tl.trans(k), tl.trans(k_scale), tl.trans(k_shift), PACKED_PER_VAL, IS_HIP, + False, ).to(dtype) k = tl.trans(k_t) @@ -1091,6 +1112,7 @@ def dequantize_k_hip( scale, shift, PACKED_PER_VAL: tl.constexpr, + FP8_FNUZ: tl.constexpr, ): """PACKED_PER_VAL is the number of values packed into each element x_. For example, for int4 quantization and x_ of type int32, PACKED_PER_VAL is 8. @@ -1110,8 +1132,8 @@ def dequantize_k_hip( ) if PACKED_PER_VAL == 4: - # FP8 quantization. - fp8_type = tl.float8e4b8 if torch.version.hip is not None else tl.float8e4nv + # FP8 quantization. gfx942 -> e4m3fnuz (float8e4b8); gfx950/CUDA -> e4m3fn. + fp8_type = tl.float8e4b8 if FP8_FNUZ else tl.float8e4nv dequant = ( quant_offset.to(tl.uint8).to(fp8_type, bitcast=True).to(scale.dtype) * scale + shift @@ -1140,6 +1162,7 @@ def dequantize( shift, PACKED_PER_VAL: tl.constexpr, IS_HIP: tl.constexpr, + FP8_FNUZ: tl.constexpr, # pyrefly: ignore [bad-function-definition] USE_FP32_SCALES: tl.constexpr = False, ): @@ -1160,8 +1183,8 @@ def dequantize( quant_offset, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL) ) if PACKED_PER_VAL == 4: - # FP8 quantization. - fp8_type = tl.float8e4b8 if torch.version.hip is not None else tl.float8e4nv + # FP8 quantization. gfx942 -> e4m3fnuz (float8e4b8); gfx950/CUDA -> e4m3fn. + fp8_type = tl.float8e4b8 if FP8_FNUZ else tl.float8e4nv dequant = ( quant_offset.to(tl.uint8).to(fp8_type, bitcast=True).to(scale.dtype) * scale ) diff --git a/mslk/attention/fmha/triton_splitk.py b/mslk/attention/fmha/triton_splitk.py index ea74400c..834715fc 100644 --- a/mslk/attention/fmha/triton_splitk.py +++ b/mslk/attention/fmha/triton_splitk.py @@ -21,6 +21,7 @@ ) import torch +from mslk.utils.device import supports_float8_fnuz from ._triton.available import is_triton_available from .attn_bias import ( @@ -862,6 +863,12 @@ def grid(META): ) IS_HIP = torch.version.hip is not None + # fp8 byte format must match how the KV cache was quantized: gfx942 e4m3fnuz, + # gfx950/CUDA e4m3fn. Derive from supports_float8_fnuz (NOT hardcoded per-HIP). + if IS_HIP: + FP8_FNUZ = supports_float8_fnuz(throw_on_hip_incompatibility=False) + else: + FP8_FNUZ = False if inp.quantize_pv_to_fp8: v = v.view(torch.int8) @@ -937,6 +944,7 @@ def grid(META): HAS_ADDITIVE_BIAS=attn_bias_tensor is not None, NUM_PROGRAMS_DIM2_CONST=split_k, IS_HIP=IS_HIP, + FP8_FNUZ=FP8_FNUZ, QUANTIZE_PV_TO_FP8=inp.quantize_pv_to_fp8, QUANTIZE_QK_TO_FP8=inp.quantize_qk_to_fp8, USE_FP32_SCALES=inp.use_fp32_scales, diff --git a/mslk/utils/device.py b/mslk/utils/device.py index 764e3a8b..21928366 100644 --- a/mslk/utils/device.py +++ b/mslk/utils/device.py @@ -74,6 +74,21 @@ def cuda_version_at_least(major_min: int) -> bool: return int(torch.version.cuda.split(".")[0]) >= major_min +def rocm_version_at_least(major_min: int, minor_min: int = 0) -> bool: + """True on a ROCm build whose HIP runtime version is at least + ``(major_min, minor_min)``. + + Parses ``torch.version.hip`` (e.g. ``"7.14.60850"``). Returns ``False`` on CUDA or + CPU-only builds. + """ + if torch.version.hip is None: + return False + parts = torch.version.hip.split(".") + major = int(parts[0]) + minor = int(parts[1]) if len(parts) > 1 else 0 + return (major, minor) >= (major_min, minor_min) + + def get_gfx_arch_name() -> str: """Return the ROCm ``gcnArchName`` of the current device (e.g. ``gfx942``). diff --git a/test/attention/fmha/test_mem_eff_attention.py b/test/attention/fmha/test_mem_eff_attention.py index ba6cba36..a756c320 100644 --- a/test/attention/fmha/test_mem_eff_attention.py +++ b/test/attention/fmha/test_mem_eff_attention.py @@ -1997,13 +1997,19 @@ def test_triton_splitk_rowwise_fp8( inp_ref, op=fmha.triton_splitk.FwOp ) + # ROCm gfx950 OCP e4m3fn snaps a few values differently than the fnuz grid, so + # loosen tolerances on ROCm only; CUDA keeps the original tight values. + is_hip = torch.version.hip is not None atol = 5e-3 - if Hkv == 2 and torch.version.hip is not None: - # XXX why is this needed? + rtol = 5e-3 + if Hkv == 2 and is_hip: atol = 1e-2 - torch.testing.assert_close(attn_output_fp8, attn_output_ref, atol=atol, rtol=5e-3) + torch.testing.assert_close(attn_output_fp8, attn_output_ref, atol=atol, rtol=rtol) assert context_fp8 is not None and context_ref is not None - torch.testing.assert_close(context_fp8.lse, context_ref.lse, atol=5e-4, rtol=5e-4) + lse_tol = 5e-3 if is_hip else 5e-4 + torch.testing.assert_close( + context_fp8.lse, context_ref.lse, atol=lse_tol, rtol=lse_tol + ) # Paged K/V cache @@ -2017,8 +2023,12 @@ def test_triton_splitk_rowwise_fp8( ) = fmha._memory_efficient_attention_forward_requires_grad( inp_fp8_paged, op=fmha.triton_splitk.FwOp ) + # Non-paged vs paged fp8 output: a few elements snap to a different e4m3 grid point + # between the two layouts on ROCm; CUDA keeps the original tight value. The LSE is + # identical between layouts on both platforms, so it keeps the tight tolerance. + paged_tol = 5e-3 if is_hip else 2e-3 torch.testing.assert_close( - attn_output_fp8, attn_output_fp8_paged, atol=2e-3, rtol=2e-3 + attn_output_fp8, attn_output_fp8_paged, atol=paged_tol, rtol=paged_tol ) assert context_fp8_paged is not None torch.testing.assert_close( diff --git a/test/attention/fmha/utils.py b/test/attention/fmha/utils.py index fc475aa1..9e88ecc7 100644 --- a/test/attention/fmha/utils.py +++ b/test/attention/fmha/utils.py @@ -25,6 +25,7 @@ ref_attention_bmhk, ) from mslk.attention.fmha.triton_splitk import InputsFp8 +from mslk.utils.triton.fp8_utils import get_fp8_constants IN_RE_WORKER: bool = os.environ.get("INSIDE_RE_WORKER") is not None @@ -186,9 +187,9 @@ def construct_fp8_attention_inputs( k = torch.randn(1, B * Mkv, Hkv, 1, K, dtype=dtype, device=device) v = torch.randn(1, B * Mkv, Hkv, 1, K, dtype=dtype, device=device) - pt_fp8_dtype = ( - torch.float8_e4m3fnuz if torch.version.hip is not None else torch.float8_e4m3fn - ) + # Match the fp8 format the decode kernels dequantize with (gfx950 uses e4m3fn, + # not fnuz); a mismatch reads the packed bytes as the wrong format -> NaN. + pt_fp8_dtype = get_fp8_constants()[0] qfn = quantize_fp8_symmetric if use_symmetric else quantize_fp8_asymmetric @@ -427,9 +428,9 @@ def add_q_fp8_to_inputs( InputsFp8 object with quantized query tensor """ inp.quantize_qk_to_fp8 = True - pt_fp8_dtype = ( - torch.float8_e4m3fnuz if torch.version.hip is not None else torch.float8_e4m3fn - ) + # Match the fp8 format the decode kernels dequantize with (gfx950 uses e4m3fn, + # not fnuz); a mismatch reads the packed bytes as the wrong format -> NaN. + pt_fp8_dtype = get_fp8_constants()[0] # Get original query tensor q = inp.query