[CUDA] Add Triton Q4_K-to-FP8 prefill linear

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
 }

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();

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<int>(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;

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",
     ],

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

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)

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: