diff --git a/backends/aoti/common_shims_slim.cpp b/backends/aoti/common_shims_slim.cpp index c8c7408aa62..8c4f72598c8 100644 --- a/backends/aoti/common_shims_slim.cpp +++ b/backends/aoti/common_shims_slim.cpp @@ -132,6 +132,10 @@ int32_t aoti_torch_dtype_float16() { return 5; // ScalarType::Half } +int32_t aoti_torch_dtype_float8_e4m3fn() { + return 24; // ScalarType::Float8_e4m3fn +} + int32_t aoti_torch_dtype_int64() { return 4; // ScalarType::Long } diff --git a/backends/aoti/common_shims_slim.h b/backends/aoti/common_shims_slim.h index c5a5cab9413..5709c86cd8f 100644 --- a/backends/aoti/common_shims_slim.h +++ b/backends/aoti/common_shims_slim.h @@ -76,6 +76,7 @@ aoti_torch_get_device_index(Tensor* tensor, int32_t* ret_device_index); AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float32(); AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_bfloat16(); AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float16(); +AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_float8_e4m3fn(); AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int64(); AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int32(); AOTI_SHIM_EXPORT int32_t aoti_torch_dtype_int16(); diff --git a/backends/aoti/slim/c10/core/ScalarType.h b/backends/aoti/slim/c10/core/ScalarType.h index 9a99aecf992..0c04a0c3e4f 100644 --- a/backends/aoti/slim/c10/core/ScalarType.h +++ b/backends/aoti/slim/c10/core/ScalarType.h @@ -39,6 +39,7 @@ enum class ScalarType : int8_t { // QUInt8 = 13, // QInt32 = 14, BFloat16 = 15, // bfloat16 + Float8_e4m3fn = 24, // float8 E4M3, finite values with NaN Undefined = -1, }; @@ -52,6 +53,7 @@ constexpr ScalarType kHalf = ScalarType::Half; constexpr ScalarType kFloat = ScalarType::Float; constexpr ScalarType kBool = ScalarType::Bool; constexpr ScalarType kBFloat16 = ScalarType::BFloat16; +constexpr ScalarType kFloat8_e4m3fn = ScalarType::Float8_e4m3fn; /// Returns the size in bytes of a single element of the given scalar type. /// @param t The scalar type. @@ -76,6 +78,8 @@ inline size_t elementSize(ScalarType t) { return sizeof(bool); case ScalarType::BFloat16: return sizeof(BFloat16); + case ScalarType::Float8_e4m3fn: + return 1; default: ET_CHECK_MSG(false, "Unknown ScalarType: %d", static_cast(t)); } @@ -104,6 +108,8 @@ inline const char* toString(ScalarType t) { return "Bool"; case ScalarType::BFloat16: return "BFloat16"; + case ScalarType::Float8_e4m3fn: + return "Float8_e4m3fn"; case ScalarType::Undefined: return "Undefined"; default: @@ -116,7 +122,7 @@ inline const char* toString(ScalarType t) { /// @return true if the scalar type is floating point, false otherwise. inline bool isFloatingType(ScalarType t) { return t == ScalarType::Half || t == ScalarType::Float || - t == ScalarType::BFloat16; + t == ScalarType::BFloat16 || t == ScalarType::Float8_e4m3fn; } /// Checks if the scalar type is an integral type (including bool optionally). @@ -159,6 +165,7 @@ inline bool isValidScalarType(ScalarType t) { case ScalarType::Float: case ScalarType::Bool: case ScalarType::BFloat16: + case ScalarType::Float8_e4m3fn: return true; default: return false; diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 654b9225fdb..2e2fd3d364b 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -3,6 +3,21 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") +fbcode_target( + _kind = runtime.python_library, + name = "coalesced_int4_tensor", + srcs = [ + "coalesced_int4_tensor.py", + ], + visibility = [ + "//executorch/...", + ], + deps = [ + "//caffe2:torch", + "//pytorch/ao:torchao", + ], +) + fbcode_target( _kind = runtime.python_library, name = "target_arch", @@ -30,16 +45,15 @@ fbcode_target( fbcode_target( _kind = runtime.python_library, - name = "coalesced_int4_tensor", + name = "q4k_dequant", srcs = [ - "coalesced_int4_tensor.py", + "quantize_op_dispatch/q4k_dequant.py", ], visibility = [ - "//executorch/...", + "//executorch/backends/cuda/...", ], deps = [ "//caffe2:torch", - "//pytorch/ao:torchao", ], ) @@ -91,6 +105,9 @@ fbcode_target( ":coalesced_int4_tensor", ":dp4a_planar_int5_tensor", ":dp4a_planar_int6_tensor", + ":q4k_dequant", + ":target_arch", + ":triton_kernels", "//caffe2:torch", "//pytorch/ao:torchao", ], @@ -156,6 +173,7 @@ fbcode_target( "triton/kernels/__init__.py", "triton/kernels/fused_moe.py", "triton/kernels/int4_matmul.py", + "triton/kernels/q4k_fp8_linear.py", "triton/kernels/sdpa.py", "triton/kernels/topk.py", "triton/kernels/tq4_sdpa.py", @@ -165,6 +183,7 @@ fbcode_target( ], deps = [ ":optimization_config", + ":q4k_dequant", ":target_arch", "//caffe2:torch", ], diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index cbc32fbad72..012ccc90a5e 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -865,10 +865,11 @@ def get_extra_aoti_compile_context_manager( through the unmodified AOTI codepath, which avoids regressions in their cuda CI exports. """ - # Parse compile_specs for low_memory_mode (default OFF). compile_specs + # Parse compile_specs for opt-in features (default OFF). compile_specs # may be None when called without specs (parity with base default). low_memory_mode = "OFF" tma_causal_prefill = False + q4k_fp8_prefill = False for spec in compile_specs or []: if spec.key == "low_memory_mode": mode = spec.value.decode("utf-8").upper() @@ -879,12 +880,19 @@ def get_extra_aoti_compile_context_manager( low_memory_mode = mode elif spec.key == "enable_tma_causal_prefill": tma_causal_prefill = _on_off_compile_spec_value(spec) + elif spec.key == "enable_q4k_fp8_prefill": + q4k_fp8_prefill = _on_off_compile_spec_value(spec) if tma_causal_prefill and not cuda_targets_are_sm90_or_newer(): logging.warning( "enable_tma_causal_prefill requires an SM90+ CUDA target; disabling it" ) tma_causal_prefill = False + if q4k_fp8_prefill and not cuda_targets_are_sm90_or_newer(): + logging.warning( + "enable_q4k_fp8_prefill requires an SM90+ CUDA target; disabling it" + ) + q4k_fp8_prefill = False @contextlib.contextmanager def _combined(): @@ -897,7 +905,10 @@ def _combined(): stack.enter_context(torch.nn.attention.sdpa_kernel([SDPBackend.MATH])) stack.enter_context(target_smem_context()) stack.enter_context( - cuda_optimization_context(tma_causal_prefill=tma_causal_prefill) + cuda_optimization_context( + q4k_fp8_prefill=q4k_fp8_prefill, + tma_causal_prefill=tma_causal_prefill, + ) ) if low_memory_mode == "ON": # Force AOTI's mutated-buffer clones onto CPU during diff --git a/backends/cuda/optimization_config.py b/backends/cuda/optimization_config.py index 8b9ee3a062e..514710da973 100644 --- a/backends/cuda/optimization_config.py +++ b/backends/cuda/optimization_config.py @@ -11,19 +11,30 @@ from typing import Iterator +_Q4K_FP8_PREFILL_ENABLED = contextvars.ContextVar( + "q4k_fp8_prefill_enabled", default=False +) _TMA_CAUSAL_PREFILL_ENABLED = contextvars.ContextVar( "tma_causal_prefill_enabled", default=False ) +def q4k_fp8_prefill_enabled() -> bool: + return _Q4K_FP8_PREFILL_ENABLED.get() + + def tma_causal_prefill_enabled() -> bool: return _TMA_CAUSAL_PREFILL_ENABLED.get() @contextlib.contextmanager -def cuda_optimization_context(*, tma_causal_prefill: bool) -> Iterator[None]: +def cuda_optimization_context( + *, q4k_fp8_prefill: bool, tma_causal_prefill: bool = False +) -> Iterator[None]: + q4k_token = _Q4K_FP8_PREFILL_ENABLED.set(q4k_fp8_prefill) tma_token = _TMA_CAUSAL_PREFILL_ENABLED.set(tma_causal_prefill) try: yield finally: _TMA_CAUSAL_PREFILL_ENABLED.reset(tma_token) + _Q4K_FP8_PREFILL_ENABLED.reset(q4k_token) diff --git a/backends/cuda/quantize_op_dispatch/int4_dispatch.py b/backends/cuda/quantize_op_dispatch/int4_dispatch.py index d6fc4cb2d1a..92fb957867a 100644 --- a/backends/cuda/quantize_op_dispatch/int4_dispatch.py +++ b/backends/cuda/quantize_op_dispatch/int4_dispatch.py @@ -16,12 +16,13 @@ At .pte runtime, the captured graph is executed by the AOTI-generated .so: - The custom op ``executorch_cuda::int4_plain_mm`` maps to a C shim that runs the W4A8 dp4a matvec kernel (backends/cuda/runtime/shims/). - - The inline dequant + F.linear is compiled by inductor into fused Triton - dequant + cuBLAS matmul kernels. + - On SM90+, prefill uses a pure-Triton Q4_K -> FP8 -> BF16-output linear. + - Older GPUs use inline dequant + F.linear, compiled into a Triton dequant + and BF16 matmul. Dispatch strategy (determines what gets captured in the export graph): Decode (M<=4): Custom op ``executorch_cuda::int4_plain_mm`` - Prefill (M>4): Inline dequant + F.linear (standard PyTorch ops) + Prefill (M>4): FP8 Triton linear on SM90+, existing BF16 path otherwise Importing the parent ``quantize_op_dispatch`` package registers this dispatch override (along with the INT8 one) before using nn.Linear with @@ -33,7 +34,11 @@ import torch import torch.nn.functional as F from executorch.backends.cuda.coalesced_int4_tensor import CudaCoalescedInt4Tensor +from executorch.backends.cuda.optimization_config import q4k_fp8_prefill_enabled from executorch.backends.cuda.quantize_op_dispatch._library import lib as _lib +from executorch.backends.cuda.quantize_op_dispatch.q4k_dequant import dequant_matmul +from executorch.backends.cuda.target_arch import cuda_targets_are_sm90_or_newer +from executorch.backends.cuda.triton.kernels.q4k_fp8_linear import q4k_fp8_linear from torch.library import impl # --------------------------------------------------------------------------- @@ -57,84 +62,13 @@ def _cuda(self, qdata, scale, scale_step, zero, zero_point_step, group_size): # Metadata is stored in the coalesced [N, n_groups] layout (transposed at # pack time, see pack_cuda.pack_linear_for_cuda). The scale is a uint8 code # with a per-256 fp16 scale_step; the zero is a uint8 code with a per-256 - # fp16 zero_point_step. _dequant_matmul reconstructs scale = + # fp16 zero_point_step. dequant_matmul reconstructs scale = # code*scale_step[g//8], zero = code*zero_point_step[g//8]. - return _dequant_matmul( + return dequant_matmul( self, qdata, scale, scale_step, zero, zero_point_step, group_size ) -# Chunked dequant for the export GPU budget. The lm_head dequant (N = vocab_size, -# e.g. 262144) runs through the int4_plain_mm custom op (M=1); AOTI executes that -# op's CUDA impl during autotune / cpp_wrapper codegen, where it transiently holds -# ~5 full-size bf16 temporaries (low/high/data/data-z/w_deq) — ~10 GiB for a -# 262144-row weight even though the final w_deq is only ~2.6 GiB. Chunking along N -# caps that at ~chunk rows. It is numerically identical (F.linear output rows are -# independent), and because only the lm_head (custom-op) path crosses the N -# threshold — never the M>4 prefill inline path — it never enters the runtime -# graph: ZERO runtime / accuracy impact. Applied unconditionally to any weight -# whose row count exceeds the threshold. -_DEQUANT_N_THRESHOLD = 65536 -_DEQUANT_N_CHUNK = 32768 - - -def _dequant_matmul(x, qdata, scale, scale_step, zero, zero_point_step, group_size): - """Dequant INT4 weights to input dtype and call F.linear. - - Metadata is in the coalesced [N, n_groups] layout (baked into the weight - constant at pack time), aligned row-for-row with qdata's [N, *]. The scale is - a uint8 code with a per-256-super-block fp16 ``scale_step`` ([N, K/256]); the - real per-group scale is ``scale_code * scale_step[:, g // 8]``. The zero is a - uint8 code with a per-256-super-block fp16 ``zero_point_step`` ([N, K/256]); - the real per-group zero is ``zero_code * zero_point_step[:, g // 8]``. - - Large weights (N > threshold, i.e. the lm_head) are chunked along N to bound - the dequant intermediate (see note above); smaller weights take the original - single-shot dequant. - """ - N, K_half = qdata.shape - K = K_half * 2 - n_groups = K // group_size - gs_half = group_size // 2 - n_super = K // 256 - groups_per_super = n_groups // n_super - dtype = x.dtype - - def _unit_dq_mm(qd, sc, s_step, ze, z_step, rows): - p = qd.to(torch.uint8).reshape(rows, n_groups, gs_half) - low = (p & 0x0F).to(dtype) - high = ((p >> 4) & 0x0F).to(dtype) - data = torch.stack([low, high], dim=-1).reshape(rows, n_groups, group_size) - # Scale: uint8 code * per-256 fp16 step (broadcast over the 8 groups in - # each super-block). - scale_step_g = s_step.to(dtype).repeat_interleave(groups_per_super, dim=1) - s = (sc.to(dtype) * scale_step_g).unsqueeze(-1) - # Zero: uint8 code * per-256 fp16 step (broadcast over the 8 groups in - # each super-block). - zero_point_step_g = z_step.to(dtype).repeat_interleave(groups_per_super, dim=1) - z = (ze.to(dtype) * zero_point_step_g).unsqueeze(-1) - w_deq = ((data - z) * s).reshape(rows, K) - return F.linear(x, w_deq) - - if N <= _DEQUANT_N_THRESHOLD: - return _unit_dq_mm(qdata, scale, scale_step, zero, zero_point_step, N) - - outs = [] - for i in range(0, N, _DEQUANT_N_CHUNK): - j = min(i + _DEQUANT_N_CHUNK, N) - outs.append( - _unit_dq_mm( - qdata[i:j], - scale[i:j], - scale_step[i:j], - zero[i:j], - zero_point_step[i:j], - j - i, - ) - ) - return torch.cat(outs, dim=-1) - - # --------------------------------------------------------------------------- # CudaCoalescedInt4Tensor F.linear dispatch # --------------------------------------------------------------------------- @@ -173,7 +107,33 @@ def _(func, types, args, kwargs): x_2d, qdata, scale, scale_step, zero, zero_point_step, gs ) else: - out = _dequant_matmul(x_2d, qdata, scale, scale_step, zero, zero_point_step, gs) + # CUDA export traces with CPU example tensors, then the CUDA backend + # lowers the captured custom op. Treat tracing as a CUDA-target case; + # requiring ``x_2d.is_cuda`` here silently falls back to BF16 dequant + # during export even when the target GPU is SM90+. + cuda_target = x_2d.is_cuda or torch.compiler.is_compiling() + if ( + q4k_fp8_prefill_enabled() + and cuda_targets_are_sm90_or_newer() + and cuda_target + and x_2d.dtype == torch.bfloat16 + and x_2d.is_contiguous() + and gs == 32 + and x_2d.shape[1] % 256 == 0 + ): + out = q4k_fp8_linear( + x_2d, + qdata, + scale, + scale_step, + zero, + zero_point_step, + gs, + ) + else: + out = dequant_matmul( + x_2d, qdata, scale, scale_step, zero, zero_point_step, gs + ) out = out.reshape(*orig_shape[:-1], -1) if bias is not None: diff --git a/backends/cuda/quantize_op_dispatch/q4k_dequant.py b/backends/cuda/quantize_op_dispatch/q4k_dequant.py new file mode 100644 index 00000000000..46b3ca2013a --- /dev/null +++ b/backends/cuda/quantize_op_dispatch/q4k_dequant.py @@ -0,0 +1,55 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Shared Q4_K dequantization fallback for CUDA linear dispatch.""" + +import torch +import torch.nn.functional as F + + +_DEQUANT_N_THRESHOLD = 65536 +_DEQUANT_N_CHUNK = 32768 + + +def dequant_matmul(x, qdata, scale, scale_step, zero, zero_point_step, group_size): + """Dequantize Q4_K weights to the activation dtype and call F.linear.""" + N, K_half = qdata.shape + K = K_half * 2 + n_groups = K // group_size + gs_half = group_size // 2 + n_super = K // 256 + groups_per_super = n_groups // n_super + dtype = x.dtype + + def _unit_dq_mm(qd, sc, s_step, ze, z_step, rows): + p = qd.to(torch.uint8).reshape(rows, n_groups, gs_half) + low = (p & 0x0F).to(dtype) + high = ((p >> 4) & 0x0F).to(dtype) + data = torch.stack([low, high], dim=-1).reshape(rows, n_groups, group_size) + scale_step_g = s_step.to(dtype).repeat_interleave(groups_per_super, dim=1) + s = (sc.to(dtype) * scale_step_g).unsqueeze(-1) + zero_point_step_g = z_step.to(dtype).repeat_interleave(groups_per_super, dim=1) + z = (ze.to(dtype) * zero_point_step_g).unsqueeze(-1) + w_deq = ((data - z) * s).reshape(rows, K) + return F.linear(x, w_deq) + + if N <= _DEQUANT_N_THRESHOLD: + return _unit_dq_mm(qdata, scale, scale_step, zero, zero_point_step, N) + + outs = [] + for i in range(0, N, _DEQUANT_N_CHUNK): + j = min(i + _DEQUANT_N_CHUNK, N) + outs.append( + _unit_dq_mm( + qdata[i:j], + scale[i:j], + scale_step[i:j], + zero[i:j], + zero_point_step[i:j], + j - i, + ) + ) + return torch.cat(outs, dim=-1) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 04c56c3243d..e27c5fad48b 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -104,6 +104,43 @@ def test_invalid_autotune_at_compile_time_compile_spec(self): [CompileSpec(key="autotune_at_compile_time", value=b"MAYBE")] ) + def test_q4k_fp8_prefill_defaults_off(self): + from executorch.backends.cuda.optimization_config import q4k_fp8_prefill_enabled + + with CudaBackend.get_extra_aoti_compile_context_manager([]): + self.assertFalse(q4k_fp8_prefill_enabled()) + + def test_q4k_fp8_prefill_compile_spec(self): + from executorch.backends.cuda.optimization_config import q4k_fp8_prefill_enabled + + with patch( + "executorch.backends.cuda.cuda_backend.cuda_targets_are_sm90_or_newer", + return_value=True, + ), CudaBackend.get_extra_aoti_compile_context_manager( + [CompileSpec(key="enable_q4k_fp8_prefill", value=b"ON")] + ): + self.assertTrue(q4k_fp8_prefill_enabled()) + + def test_invalid_q4k_fp8_prefill_compile_spec(self): + with self.assertRaisesRegex(ValueError, "Invalid enable_q4k_fp8_prefill"): + CudaBackend.get_extra_aoti_compile_context_manager( + [CompileSpec(key="enable_q4k_fp8_prefill", value=b"MAYBE")] + ) + + def test_q4k_fp8_prefill_unsupported_target_is_disabled(self): + from executorch.backends.cuda.optimization_config import q4k_fp8_prefill_enabled + + with patch( + "executorch.backends.cuda.cuda_backend.cuda_targets_are_sm90_or_newer", + return_value=False, + ), self.assertLogs( + level="WARNING" + ) as logs, CudaBackend.get_extra_aoti_compile_context_manager( + [CompileSpec(key="enable_q4k_fp8_prefill", value=b"ON")] + ): + self.assertFalse(q4k_fp8_prefill_enabled()) + self.assertIn("requires an SM90+ CUDA target", "\n".join(logs.output)) + def test_tma_causal_prefill_defaults_off(self): from executorch.backends.cuda.optimization_config import ( tma_causal_prefill_enabled, diff --git a/backends/cuda/tests/test_int4_dispatch.py b/backends/cuda/tests/test_int4_dispatch.py index b3758683359..5a9921962a2 100644 --- a/backends/cuda/tests/test_int4_dispatch.py +++ b/backends/cuda/tests/test_int4_dispatch.py @@ -25,6 +25,7 @@ """ import contextlib +import os import unittest from unittest import mock @@ -33,7 +34,9 @@ import torch.nn as nn import torch.nn.functional as F from executorch.backends.cuda.coalesced_int4_tensor import CudaCoalescedInt4Tensor -from executorch.backends.cuda.quantize_op_dispatch.int4_dispatch import _dequant_matmul +from executorch.backends.cuda.optimization_config import cuda_optimization_context +from executorch.backends.cuda.quantize_op_dispatch.q4k_dequant import dequant_matmul +from executorch.backends.cuda.target_arch import cuda_targets_are_sm90_or_newer from executorch.examples.models.gemma4_31b.cuda_packers import pack_linear_for_cuda from executorch.extension.llm.export.int4 import ExportableInt4Tensor from executorch.extension.llm.export.quant.quantize import ( @@ -111,6 +114,25 @@ def test_group_size_32(self): x = torch.randn(1, 256, dtype=torch.bfloat16, device="cuda") self._check(module(x), F.linear(x, w_ref)) + def test_prefill_group_size_32(self): + module, _ = _make_int4_linear(256, 512, group_size=32) + x = torch.randn(64, 512, dtype=torch.bfloat16, device="cuda") + out = module(x) + weight = module.weight + ref = dequant_matmul( + x, + weight.qdata, + weight.scale, + weight.scale_step, + weight.zero_point, + weight.zero_point_step, + weight.block_size[-1], + ) + normalized_max_error = ( + out.float() - ref.float() + ).abs().max() / ref.float().abs().max() + self.assertLess(normalized_max_error.item(), 0.05) + def test_symmetric(self): module, w_ref = _make_int4_linear(256, 512, symmetric=True) x = torch.randn(1, 512, dtype=torch.bfloat16, device="cuda") @@ -253,7 +275,7 @@ def _record_int4_plain_mm(): def _fake(self, qdata, scale, scale_step, zero, zero_point_step, group_size): calls.append((tuple(self.shape), group_size)) - return _dequant_matmul( + return dequant_matmul( self, qdata, scale, scale_step, zero, zero_point_step, group_size ) @@ -310,6 +332,93 @@ def test_coalesced_tensor_prefill_uses_dequant(self): ref = F.linear(x, dequantize_weight(t, torch.bfloat16)) self.assertLess(self._rel_err(out, ref), 0.02) + def test_sm90_prefill_export_captures_fp8_linear(self): + """CPU example tensors must still capture the CUDA-target FP8 op.""" + if torch.version.hip is not None: + self.skipTest("Q4_K-to-FP8 prefill is CUDA-only") + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9: + self.skipTest("SM90+ CUDA target required") + + class LinearModule(nn.Module): + def __init__(self): + super().__init__() + t, _ = _make_exportable_int4_tensor(64, 256, group_size=32) + packed = CudaCoalescedInt4Tensor.from_exportable_int4_tensor(t) + self.weight = nn.Parameter(packed, requires_grad=False) + + def forward(self, x): + return F.linear(x, self.weight) + + exported = torch.export.export( + LinearModule(), (torch.randn(8, 256, dtype=torch.bfloat16),), strict=True + ) + self.assertNotIn("triton.q4k_fp8_linear", exported.graph_module.code) + + with cuda_optimization_context(q4k_fp8_prefill=True): + exported = torch.export.export( + LinearModule(), + (torch.randn(8, 256, dtype=torch.bfloat16),), + strict=True, + ) + self.assertIn("triton.q4k_fp8_linear", exported.graph_module.code) + + def test_fp8_gate_falls_back_to_bfloat16_dequant(self): + module, _ = _make_int4_linear(256, 512, group_size=32) + x = torch.randn(8, 512, dtype=torch.bfloat16, device="cuda") + with mock.patch( + "executorch.backends.cuda.quantize_op_dispatch.int4_dispatch.q4k_fp8_linear" + ) as fp8_linear: + out = module(x) + fp8_linear.assert_not_called() + self.assertEqual(out.shape, (8, 256)) + + def test_target_arch_override_controls_fp8_gate(self): + with mock.patch.object(torch.version, "hip", None): + with mock.patch.dict( + os.environ, {"TORCH_CUDA_ARCH_LIST": "8.0"} + ), mock.patch.object(torch.cuda, "get_device_capability") as capability: + self.assertFalse(cuda_targets_are_sm90_or_newer()) + capability.assert_not_called() + with mock.patch.dict(os.environ, {"TORCH_CUDA_ARCH_LIST": "9.0;12.0+PTX"}): + self.assertTrue(cuda_targets_are_sm90_or_newer()) + with mock.patch.object(torch.version, "hip", "6.3"), mock.patch.dict( + os.environ, {"TORCH_CUDA_ARCH_LIST": "9.0"} + ): + self.assertFalse(cuda_targets_are_sm90_or_newer()) + + def test_noncontiguous_prefill_falls_back(self): + t, _ = _make_exportable_int4_tensor(16, 256, group_size=32) + weight = CudaCoalescedInt4Tensor.from_exportable_int4_tensor(t) + x = torch.randn(1, 16, 256, dtype=torch.bfloat16)[:, ::2, :] + self.assertFalse(x.reshape(-1, 256).is_contiguous()) + with mock.patch( + "executorch.backends.cuda.quantize_op_dispatch.int4_dispatch.cuda_targets_are_sm90_or_newer", + return_value=True, + ), mock.patch( + "executorch.backends.cuda.quantize_op_dispatch.int4_dispatch.q4k_fp8_linear" + ) as fp8_linear, mock.patch.object( + torch.compiler, "is_compiling", return_value=True + ): + out = F.linear(x, weight) + fp8_linear.assert_not_called() + self.assertEqual(out.shape, (1, 8, 16)) + + def test_non_bfloat16_prefill_falls_back(self): + t, _ = _make_exportable_int4_tensor(16, 256, group_size=32) + weight = CudaCoalescedInt4Tensor.from_exportable_int4_tensor(t) + x = torch.randn(8, 256, dtype=torch.float32) + with mock.patch( + "executorch.backends.cuda.quantize_op_dispatch.int4_dispatch.cuda_targets_are_sm90_or_newer", + return_value=True, + ), mock.patch( + "executorch.backends.cuda.quantize_op_dispatch.int4_dispatch.q4k_fp8_linear" + ) as fp8_linear, mock.patch.object( + torch.compiler, "is_compiling", return_value=True + ): + out = F.linear(x, weight) + fp8_linear.assert_not_called() + self.assertEqual(out.dtype, torch.float32) + def test_square_shape_not_misrouted(self): """N == n_groups (square scale) stock tensor is still not routed. diff --git a/backends/cuda/triton/kernels/q4k_fp8_linear.py b/backends/cuda/triton/kernels/q4k_fp8_linear.py new file mode 100644 index 00000000000..81463dfce19 --- /dev/null +++ b/backends/cuda/triton/kernels/q4k_fp8_linear.py @@ -0,0 +1,209 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Pure-Triton Q4_K -> FP8 prefill linear for SM90+ GPUs.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +from executorch.backends.cuda.optimization_config import q4k_fp8_prefill_enabled +from executorch.backends.cuda.quantize_op_dispatch.q4k_dequant import dequant_matmul +from torch.library import triton_op, wrap_triton + + +@triton.jit +def _cast_bf16_to_fp8_kernel(x, out, n_elements, XBLOCK: tl.constexpr): + idx = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK) + mask = idx < n_elements + value = tl.load(x + idx, mask=mask, other=0.0) + tl.store(out + idx, value.to(tl.float8e4nv), mask=mask) + + +@triton.jit +def _dequant_q4k_fp8_kernel( + qdata, + scale_code, + scale_step, + zero_code, + zero_step, + out, + n_elements, + K: tl.constexpr, + GROUP_SIZE: tl.constexpr, + GROUPS_PER_SUPER: tl.constexpr, + XBLOCK: tl.constexpr, +): + idx = tl.program_id(0) * XBLOCK + tl.arange(0, XBLOCK) + mask = idx < n_elements + n = idx // K + k = idx - n * K + packed = tl.load(qdata + n * (K // 2) + k // 2, mask=mask, other=0).to(tl.uint8) + q = ((packed >> ((k & 1) * 4)) & 0xF).to(tl.float32) + group = k // GROUP_SIZE + super_group = group // GROUPS_PER_SUPER + groups = K // GROUP_SIZE + supers = groups // GROUPS_PER_SUPER + sc = tl.load(scale_code + n * groups + group, mask=mask, other=0).to(tl.float32) + ss = tl.load(scale_step + n * supers + super_group, mask=mask, other=0.0).to( + tl.float32 + ) + zc = tl.load(zero_code + n * groups + group, mask=mask, other=0).to(tl.float32) + zs = tl.load(zero_step + n * supers + super_group, mask=mask, other=0.0).to( + tl.float32 + ) + value = (q - zc * zs) * (sc * ss) + tl.store(out + idx, value.to(tl.float8e4nv), mask=mask) + + +_FP8_MM_CONFIGS = [ + triton.Config( + {"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, + num_warps=4, + num_stages=3, + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 128, "GROUP_M": 8}, + num_warps=8, + num_stages=3, + ), + triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 128, "GROUP_M": 8}, + num_warps=8, + num_stages=3, + ), +] + + +def _validate_q4k_fp8_inputs( + x: torch.Tensor, qdata: torch.Tensor, group_size: int +) -> None: + if x.dtype != torch.bfloat16: + raise RuntimeError(f"Expected bfloat16 activations, got {x.dtype}") + if not x.is_contiguous(): + raise RuntimeError("Q4_K FP8 requires contiguous activations") + if qdata.dtype != torch.uint8: + raise RuntimeError(f"Expected uint8 packed weights, got {qdata.dtype}") + if group_size != 32: + raise RuntimeError(f"Q4_K FP8 requires group_size=32, got {group_size}") + if x.shape[1] % 256 != 0: + raise RuntimeError(f"Q4_K FP8 requires K divisible by 256, got {x.shape[1]}") + + +@triton.autotune(configs=_FP8_MM_CONFIGS, key=["M", "N", "K"]) +@triton.jit +def _fp8_mm_kernel( + a, + b, + c, + M, + N: tl.constexpr, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_M: tl.constexpr, +): + pid = tl.program_id(0) + n_m = tl.cdiv(M, BLOCK_M) + n_n = tl.cdiv(N, BLOCK_N) + width = GROUP_M * n_n + group = pid // width + first_m = group * GROUP_M + group_m = tl.minimum(n_m - first_m, GROUP_M) + pid_m = first_m + ((pid % width) % group_m) + pid_n = (pid % width) // group_m + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + acc = tl.zeros((BLOCK_M, BLOCK_N), tl.float32) + for kt in range(0, tl.cdiv(K, BLOCK_K)): + k = kt * BLOCK_K + offs_k + av = tl.load( + a + offs_m[:, None] * K + k[None, :], + mask=(offs_m[:, None] < M) & (k[None, :] < K), + other=0.0, + ) + bv = tl.load( + b + offs_n[None, :] * K + k[:, None], + mask=(offs_n[None, :] < N) & (k[:, None] < K), + other=0.0, + ) + acc += tl.dot(av, bv) + tl.store( + c + offs_m[:, None] * N + offs_n[None, :], + acc.to(tl.bfloat16), + mask=(offs_m[:, None] < M) & (offs_n[None, :] < N), + ) + + +@triton_op("triton::q4k_fp8_linear", mutates_args={}) +def q4k_fp8_linear( + x: torch.Tensor, + qdata: torch.Tensor, + scale: torch.Tensor, + scale_step: torch.Tensor, + zero: torch.Tensor, + zero_point_step: torch.Tensor, + group_size: int, +) -> torch.Tensor: + """Run Q4_K prefill linear through FP8 tensor cores on SM90+.""" + if not q4k_fp8_prefill_enabled(): + return dequant_matmul( + x, qdata, scale, scale_step, zero, zero_point_step, group_size + ) + + M, K = x.shape + N = qdata.shape[0] + _validate_q4k_fp8_inputs(x, qdata, group_size) + groups_per_super = 256 // group_size + x_fp8 = torch.empty((M, K), dtype=torch.float8_e4m3fn, device=x.device) + w_fp8 = torch.empty((N, K), dtype=torch.float8_e4m3fn, device=x.device) + output = torch.empty((M, N), dtype=torch.bfloat16, device=x.device) + + wrap_triton(_cast_bf16_to_fp8_kernel)[ + lambda meta: (triton.cdiv(M * K, meta["XBLOCK"]),) + ](x, x_fp8, M * K, XBLOCK=512, num_warps=4, num_stages=1) + wrap_triton(_dequant_q4k_fp8_kernel)[ + lambda meta: (triton.cdiv(N * K, meta["XBLOCK"]),) + ]( + qdata, + scale, + scale_step, + zero, + zero_point_step, + w_fp8, + N * K, + K, + group_size, + groups_per_super, + XBLOCK=512, + num_warps=4, + num_stages=1, + ) + wrap_triton(_fp8_mm_kernel)[ + lambda meta: ( + triton.cdiv(M, meta["BLOCK_M"]) * triton.cdiv(N, meta["BLOCK_N"]), + ) + ](x_fp8, w_fp8, output, M, N, K) + return output + + +@q4k_fp8_linear.register_fake +def _q4k_fp8_linear_fake( + x: torch.Tensor, + qdata: torch.Tensor, + scale: torch.Tensor, + scale_step: torch.Tensor, + zero: torch.Tensor, + zero_point_step: torch.Tensor, + group_size: int, +) -> torch.Tensor: + _validate_q4k_fp8_inputs(x, qdata, group_size) + return torch.empty( + (x.shape[0], qdata.shape[0]), dtype=torch.bfloat16, device=x.device + ) diff --git a/examples/models/muse-glimmer/export/export_dflash.py b/examples/models/muse-glimmer/export/export_dflash.py index 1e4ee8b4fec..07b39dc7aa5 100644 --- a/examples/models/muse-glimmer/export/export_dflash.py +++ b/examples/models/muse-glimmer/export/export_dflash.py @@ -58,6 +58,7 @@ def export_dflash( backend: str = "mlx", mmproj: str | None = None, max_vision_patches: int = 16384, + enable_q4k_fp8_prefill: bool = False, enable_tma_causal_prefill: bool = False, ) -> None: """Export DFlash target + draft to one CUDA or MLX .pte. @@ -157,7 +158,10 @@ def export_dflash( cuda_optimization_context, ) - with cuda_optimization_context(tma_causal_prefill=enable_tma_causal_prefill): + with cuda_optimization_context( + q4k_fp8_prefill=enable_q4k_fp8_prefill, + tma_causal_prefill=enable_tma_causal_prefill, + ): _export_dflash_cuda( target_model, target_config, @@ -169,6 +173,7 @@ def export_dflash( max_seq_len, activation_dtype, max_vision_patches, + enable_q4k_fp8_prefill=enable_q4k_fp8_prefill, enable_tma_causal_prefill=enable_tma_causal_prefill, ) @@ -442,6 +447,7 @@ def _export_dflash_cuda( max_seq_len: int, activation_dtype: torch.dtype, max_vision_patches: int, + enable_q4k_fp8_prefill: bool, enable_tma_causal_prefill: bool, ) -> None: """Export the CUDA DFlash contract. @@ -648,6 +654,10 @@ def cuda_partitioner(method: str) -> CudaPartitioner: [ CudaBackend.generate_method_name_compile_spec(method), CompileSpec("low_memory_mode", b"ON"), + CompileSpec( + "enable_q4k_fp8_prefill", + b"ON" if enable_q4k_fp8_prefill else b"OFF", + ), CompileSpec( "enable_tma_causal_prefill", b"ON" if enable_tma_causal_prefill else b"OFF", @@ -796,6 +806,11 @@ def main() -> None: help="Activation / KV-cache / unquantized-weight dtype. Defaults to " "float16 for MLX and bfloat16 for CUDA.", ) + parser.add_argument( + "--enable-q4k-fp8-prefill", + action="store_true", + help="Enable the experimental SM90+ Q4_K-to-FP8 prefill linear path.", + ) parser.add_argument( "--enable-tma-causal-prefill", action="store_true", @@ -829,6 +844,7 @@ def main() -> None: backend=args.backend, mmproj=args.mmproj, max_vision_patches=args.max_vision_patches, + enable_q4k_fp8_prefill=args.enable_q4k_fp8_prefill, enable_tma_causal_prefill=args.enable_tma_causal_prefill, ) diff --git a/examples/models/muse-glimmer/export/export_solo.py b/examples/models/muse-glimmer/export/export_solo.py index 9576264e70d..6962a87ff79 100644 --- a/examples/models/muse-glimmer/export/export_solo.py +++ b/examples/models/muse-glimmer/export/export_solo.py @@ -84,6 +84,7 @@ def export_and_lower( pos_embed_table: torch.Tensor | None = None, max_vision_patches: int = 16384, vision_fp32_mm: str = "none", + enable_q4k_fp8_prefill: bool = False, enable_tma_causal_prefill: bool = False, ) -> None: if backend == "cuda": @@ -91,7 +92,10 @@ def export_and_lower( cuda_optimization_context, ) - with cuda_optimization_context(tma_causal_prefill=enable_tma_causal_prefill): + with cuda_optimization_context( + q4k_fp8_prefill=enable_q4k_fp8_prefill, + tma_causal_prefill=enable_tma_causal_prefill, + ): _export_cuda( model, config, @@ -102,6 +106,7 @@ def export_and_lower( pos_embed_table=pos_embed_table, max_vision_patches=max_vision_patches, vision_fp32_mm=vision_fp32_mm, + enable_q4k_fp8_prefill=enable_q4k_fp8_prefill, enable_tma_causal_prefill=enable_tma_causal_prefill, ) elif backend == "mlx": @@ -165,6 +170,7 @@ def _export_cuda( pos_embed_table: torch.Tensor | None = None, max_vision_patches: int = 16384, vision_fp32_mm: str = "none", + enable_q4k_fp8_prefill: bool = False, enable_tma_causal_prefill: bool = False, ) -> None: import torch._inductor.config as inductor_config @@ -287,6 +293,10 @@ def _partitioner_for(name: str) -> "CudaPartitioner": [ CudaBackend.generate_method_name_compile_spec(name), CompileSpec("low_memory_mode", b"ON"), + CompileSpec( + "enable_q4k_fp8_prefill", + b"ON" if enable_q4k_fp8_prefill else b"OFF", + ), CompileSpec( "enable_tma_causal_prefill", b"ON" if enable_tma_causal_prefill else b"OFF", @@ -623,6 +633,11 @@ def main() -> None: help="Optional FP32-output linear implementation for vision blocks " "0-34. The default preserves the original all-BF16 encoder.", ) + parser.add_argument( + "--enable-q4k-fp8-prefill", + action="store_true", + help="Enable the experimental SM90+ Q4_K-to-FP8 prefill linear path.", + ) parser.add_argument( "--enable-tma-causal-prefill", action="store_true", @@ -702,6 +717,7 @@ def main() -> None: pos_embed_table=pos_embed_table, max_vision_patches=args.max_vision_patches, vision_fp32_mm=args.vision_fp32_mm, + enable_q4k_fp8_prefill=args.enable_q4k_fp8_prefill, enable_tma_causal_prefill=args.enable_tma_causal_prefill, )