Implement full CPU forward for MultiScaleDeformableAttention

Replace NotImplementedError stub with the actual pure-Python CPU
implementation using F.grid_sample bilinear interpolation, matching
the upstream mmcv multi_scale_deformable_attn_pytorch logic.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

setup_mmcv_shim.py

@@ -170,35 +170,129 @@ class ConcatCell:
     def __init__(self, *args, **kwargs): raise NotImplementedError
 ''')
 
-# Create multi_scale_deform_attn stub with registry registration
+# Create multi_scale_deform_attn with full CPU implementation
 with open(os.path.join(ops_dir, 'multi_scale_deform_attn.py'), 'w') as f:
-    f.write('''import warnings
+    f.write('''import math
+import warnings
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
+from mmcv.cnn import constant_init, xavier_init
 from mmcv.cnn.bricks.registry import ATTENTION
 from mmcv.runner import BaseModule
 
+
+def multi_scale_deformable_attn_pytorch(value, value_spatial_shapes,
+                                        sampling_locations, attention_weights):
+    """CPU version of multi-scale deformable attention."""
+    bs, _, num_heads, embed_dims = value.shape
+    _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape
+    value_list = value.split([int(H_ * W_) for H_, W_ in value_spatial_shapes], dim=1)
+    sampling_grids = 2 * sampling_locations - 1
+    sampling_value_list = []
+    for level, (H_, W_) in enumerate(value_spatial_shapes):
+        value_l_ = value_list[level].flatten(2).transpose(1, 2).reshape(
+            bs * num_heads, embed_dims, int(H_), int(W_))
+        sampling_grid_l_ = sampling_grids[:, :, :, level].transpose(1, 2).flatten(0, 1)
+        sampling_value_l_ = F.grid_sample(
+            value_l_, sampling_grid_l_, mode=\'bilinear\', padding_mode=\'zeros\',
+            align_corners=False)
+        sampling_value_list.append(sampling_value_l_)
+    attention_weights = attention_weights.transpose(1, 2).reshape(
+        bs * num_heads, 1, num_queries, num_levels * num_points)
+    output = (torch.stack(sampling_value_list, dim=-2).flatten(-2) *
+              attention_weights).sum(-1).view(bs, num_heads * embed_dims, num_queries)
+    return output.transpose(1, 2).contiguous()
+
+
 @ATTENTION.register_module()
 class MultiScaleDeformableAttention(BaseModule):
     def __init__(self, embed_dims=256, num_heads=8, num_levels=4, num_points=4,
                  im2col_step=64, dropout=0.1, batch_first=False, norm_cfg=None,
                  init_cfg=None, **kwargs):
         super().__init__(init_cfg)
+        if embed_dims % num_heads != 0:
+            raise ValueError(f\'embed_dims must be divisible by num_heads, \'
+                             f\'but got {embed_dims} and {num_heads}\')
+        self.norm_cfg = norm_cfg
+        self.dropout = nn.Dropout(dropout)
+        self.batch_first = batch_first
+        self.im2col_step = im2col_step
         self.embed_dims = embed_dims
-        self.num_heads = num_heads
         self.num_levels = num_levels
+        self.num_heads = num_heads
         self.num_points = num_points
-        self.batch_first = batch_first
         self.sampling_offsets = nn.Linear(embed_dims, num_heads * num_levels * num_points * 2)
         self.attention_weights = nn.Linear(embed_dims, num_heads * num_levels * num_points)
         self.value_proj = nn.Linear(embed_dims, embed_dims)
         self.output_proj = nn.Linear(embed_dims, embed_dims)
-        self.dropout = nn.Dropout(dropout)
+        self.init_weights()
+
+    def init_weights(self):
+        constant_init(self.sampling_offsets, 0.)
+        device = next(self.parameters()).device
+        thetas = torch.arange(self.num_heads, dtype=torch.float32, device=device) * (2.0 * math.pi / self.num_heads)
+        grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
+        grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view(
+            self.num_heads, 1, 1, 2).repeat(1, self.num_levels, self.num_points, 1)
+        for i in range(self.num_points):
+            grid_init[:, :, i, :] *= i + 1
+        self.sampling_offsets.bias.data = grid_init.view(-1)
+        constant_init(self.attention_weights, val=0., bias=0.)
+        xavier_init(self.value_proj, distribution=\'uniform\', bias=0.)
+        xavier_init(self.output_proj, distribution=\'uniform\', bias=0.)
+        self._is_init = True
 
     def forward(self, query, key=None, value=None, identity=None,
                 query_pos=None, key_padding_mask=None, reference_points=None,
                 spatial_shapes=None, level_start_index=None, **kwargs):
-        raise NotImplementedError("MultiScaleDeformableAttention CPU shim - forward not implemented")
+        if value is None:
+            value = query
+        if identity is None:
+            identity = query
+        if query_pos is not None:
+            query = query + query_pos
+        if not self.batch_first:
+            query = query.permute(1, 0, 2)
+            value = value.permute(1, 0, 2)
+
+        bs, num_query, _ = query.shape
+        bs, num_value, _ = value.shape
+
+        value = self.value_proj(value)
+        if key_padding_mask is not None:
+            value = value.masked_fill(key_padding_mask[..., None], 0.0)
+        value = value.view(bs, num_value, self.num_heads, -1)
+
+        sampling_offsets = self.sampling_offsets(query).view(
+            bs, num_query, self.num_heads, self.num_levels, self.num_points, 2)
+        attention_weights = self.attention_weights(query).view(
+            bs, num_query, self.num_heads, self.num_levels * self.num_points)
+        attention_weights = attention_weights.softmax(-1)
+        attention_weights = attention_weights.view(
+            bs, num_query, self.num_heads, self.num_levels, self.num_points)
+
+        if reference_points.shape[-1] == 2:
+            offset_normalizer = torch.stack(
+                [spatial_shapes[..., 1], spatial_shapes[..., 0]], -1)
+            sampling_locations = reference_points[:, :, None, :, None, :] \\
+                + sampling_offsets / offset_normalizer[None, None, None, :, None, :]
+        elif reference_points.shape[-1] == 4:
+            sampling_locations = reference_points[:, :, None, :, None, :2] \\
+                + sampling_offsets / self.num_points \\
+                * reference_points[:, :, None, :, None, 2:] * 0.5
+        else:
+            raise ValueError(f\'Last dim of reference_points must be 2 or 4, \'
+                             f\'but get {reference_points.shape[-1]} instead.\')
+
+        output = multi_scale_deformable_attn_pytorch(
+            value, spatial_shapes, sampling_locations, attention_weights)
+        output = self.output_proj(output)
+
+        if not self.batch_first:
+            output = output.permute(1, 0, 2)
+
+        return self.dropout(output) + identity
 ''')
 
 print('mmcv.ops shim created successfully')