DCNv4 INT8 patch (Level 1): int8 storage + fp32 arithmetic for SOFIA/MERIDIAN E9
- Add dcnv4_int8_cuda.cu/.h: CUDA kernel (int8 values, fp16 offsets, fp32 interp, requantize with value_scale/output_scale) - Add dcnv4_int8_forward(): inference-only functional wrapper (@no_grad) - Add DCNv4Strip.forward_int8(): module-level INT8 forward (without_pointwise=True) - Add scripts/test_dcnv4_int8.py: correctness gate (<=1 LSB, >=99% exact) and informational fp16 vs int8 benchmark - Update README: INT8 API section, updated structure tree, SOFIA/MERIDIAN context Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -8,4 +8,4 @@
|
||||
|
||||
# from .ms_flash_deform_attn_func import FlashMSDeformAttnFunction
|
||||
from .flash_deform_attn_func import FlashDeformAttnFunction
|
||||
from .dcnv4_func import DCNv4Function
|
||||
from .dcnv4_func import DCNv4Function, dcnv4_int8_forward
|
||||
@@ -59,6 +59,35 @@ def find_spec_bwd(B, H, W, G, C):
|
||||
n_thread = multiplier * G * C // d_stride
|
||||
return d_stride, n_thread
|
||||
|
||||
@torch.no_grad()
|
||||
def dcnv4_int8_forward(
|
||||
value_int8, offset_mask_fp16,
|
||||
kernel_h, kernel_w, pad_h, pad_w,
|
||||
group, group_channels,
|
||||
offset_scale, value_scale, output_scale):
|
||||
"""Inference-only INT8 DCNv4 (Level 1: int8 storage + fp32 math).
|
||||
|
||||
Args:
|
||||
value_int8: int8 tensor [B, H, W, C], quantized as
|
||||
real_value = value_int8 * value_scale (per-tensor, symmetric).
|
||||
offset_mask_fp16: fp16 tensor [B, H, W, padded_offset_dim] — the raw
|
||||
(real-valued) offsets and masks, NOT quantized.
|
||||
kernel_h/kernel_w/pad_h/pad_w/group/group_channels/offset_scale:
|
||||
same semantics as DCNv4Function (stride=1, dilation=1 implied).
|
||||
value_scale / output_scale: per-tensor quantization scales; the kernel
|
||||
folds them into a single requantization multiplier
|
||||
value_scale / output_scale.
|
||||
|
||||
Returns:
|
||||
int8 tensor [B, H, W, C]; real output = result * output_scale.
|
||||
"""
|
||||
requant = float(value_scale) / float(output_scale)
|
||||
return ext.dcnv4_int8_forward(
|
||||
value_int8, offset_mask_fp16,
|
||||
kernel_h, kernel_w, 1, 1, pad_h, pad_w, 1, 1,
|
||||
group, group_channels, float(offset_scale), requant)
|
||||
|
||||
|
||||
class DCNv4Function(Function):
|
||||
@staticmethod
|
||||
@custom_fwd
|
||||
|
||||
@@ -13,7 +13,7 @@ import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.init import xavier_uniform_, constant_
|
||||
from ..functions import DCNv4Function
|
||||
from ..functions import DCNv4Function, dcnv4_int8_forward
|
||||
|
||||
class CenterFeatureScaleModule(nn.Module):
|
||||
def forward(self,
|
||||
@@ -284,3 +284,41 @@ class DCNv4Strip(nn.Module):
|
||||
x = self.output_proj(x)
|
||||
return x
|
||||
|
||||
@torch.no_grad()
|
||||
def forward_int8(self, input_int8, shape, value_scale, output_scale):
|
||||
"""INT8 inference forward (Level 1: int8 storage + fp32 math).
|
||||
|
||||
The deformable sampling consumes/produces int8; the (small) offset/mask
|
||||
Linear runs in the module's own float dtype on the dequantized input.
|
||||
Requires ``without_pointwise=True`` (the int8 path quantizes the
|
||||
sampling op only; pointwise projections belong to the surrounding
|
||||
block's quantized 1x1 convs).
|
||||
|
||||
Args:
|
||||
input_int8: int8 [N, L, C]; real value = input_int8 * value_scale.
|
||||
shape: (H, W) of the token grid.
|
||||
value_scale: per-tensor input scale.
|
||||
output_scale: per-tensor output scale.
|
||||
|
||||
Returns:
|
||||
int8 [N, L, C]; real output = result * output_scale.
|
||||
"""
|
||||
assert self.without_pointwise, \
|
||||
"forward_int8 supports without_pointwise=True only"
|
||||
assert input_int8.dtype == torch.int8, "input must be int8"
|
||||
N, L, C = input_int8.shape
|
||||
H, W = shape
|
||||
|
||||
w_dtype = self.offset_mask.weight.dtype
|
||||
x_deq = input_int8.to(w_dtype) * float(value_scale)
|
||||
offset_mask = self.offset_mask(x_deq).to(torch.float16)
|
||||
offset_mask = offset_mask.reshape(N, H, W, -1).contiguous()
|
||||
|
||||
out = dcnv4_int8_forward(
|
||||
input_int8.reshape(N, H, W, C).contiguous(), offset_mask,
|
||||
self.kernel_h, self.kernel_w, self.pad_h, self.pad_w,
|
||||
self.group, self.group_channels,
|
||||
self.offset_scale, value_scale, output_scale,
|
||||
)
|
||||
return out.reshape(N, L, C)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user