@@ -142,7 +142,6 @@ def _export_decoder_and_embedding(
142142 qembedding ,
143143 qembedding_group_size ,
144144 use_aoti_packed_int4 = False ,
145- use_aoti_matvec = False ,
146145 device = "cpu" ,
147146):
148147 """Export text_decoder and token_embedding into programs dict."""
@@ -155,6 +154,7 @@ def _export_decoder_and_embedding(
155154 text_decoder .eval ()
156155
157156 packed_linear_count = 0
157+ use_packed_matvec = use_aoti_packed_int4 and qlinear == "4w"
158158 if qlinear :
159159 print (f" Quantizing decoder ({ qlinear } )..." )
160160 quantize_model_ (
@@ -163,16 +163,14 @@ def _export_decoder_and_embedding(
163163 qlinear_group_size = qlinear_group_size ,
164164 qlinear_packing_format = qlinear_packing_format ,
165165 )
166- if use_aoti_packed_int4 and qlinear == "4w" :
166+ if use_packed_matvec :
167167 packed_linear_count = _pack_aoti_int4_weights (
168168 text_decoder ,
169- use_matvec = use_aoti_matvec ,
169+ use_matvec = True ,
170170 )
171171
172- if use_aoti_matvec :
173- # TODO: Resolve fixed-shape greedy-output drift before enabling this by
174- # default; the same drift reproduces with int4_matmul.
175- # Both native runner paths invoke the decoder one token at a time.
172+ # Native runners decode one token per call; static M=1 enables matvec dispatch.
173+ if use_packed_matvec :
176174 sample_embeds = torch .randn (
177175 1 , 1 , model .config .dim , dtype = param_dtype , device = device
178176 )
@@ -232,7 +230,6 @@ def export_all(
232230 qembedding = None ,
233231 qembedding_group_size = None ,
234232 use_aoti_packed_int4 = False ,
235- use_aoti_matvec = False ,
236233 backend = "xnnpack" ,
237234):
238235 """Export all three model components with per-component quantization."""
@@ -297,7 +294,6 @@ def export_all(
297294 qembedding ,
298295 qembedding_group_size ,
299296 use_aoti_packed_int4 = use_aoti_packed_int4 ,
300- use_aoti_matvec = use_aoti_matvec ,
301297 device = device ,
302298 )
303299 if packed_linear_count :
@@ -331,7 +327,6 @@ def export_streaming(
331327 qembedding = None ,
332328 qembedding_group_size = None ,
333329 use_aoti_packed_int4 = False ,
334- use_aoti_matvec = False ,
335330 backend = "xnnpack" ,
336331):
337332 """Export streaming model components with per-component quantization."""
@@ -392,7 +387,6 @@ def export_streaming(
392387 qembedding ,
393388 qembedding_group_size ,
394389 use_aoti_packed_int4 = use_aoti_packed_int4 ,
395- use_aoti_matvec = use_aoti_matvec ,
396390 device = device ,
397391 )
398392 if packed_linear_count :
@@ -583,15 +577,11 @@ def _validate_rocm_args(parser, args):
583577 "tile_packed_to_4d requires a CUDA-only int4 fallback; "
584578 "omit the packing format for ROCm"
585579 )
586- if args .rocm_packed_matvec and args .qlinear != "4w" :
587- parser .error ("--rocm-packed-matvec requires --qlinear=4w" )
588580
589581
590582def _validate_export_args (parser , args , backend_for_export ):
591583 if args .backend == "rocm" :
592584 _validate_rocm_args (parser , args )
593- elif args .rocm_packed_matvec :
594- parser .error ("--rocm-packed-matvec requires --backend=rocm" )
595585
596586 if args .qlinear == "fpa4w" and backend_for_export != "metal" :
597587 parser .error ("--qlinear=fpa4w can only be used with --backend=metal" )
@@ -708,11 +698,6 @@ def main():
708698 "typically 8192). Smaller values reduce memory and improve decode speed "
709699 "but limit how far back the decoder can attend. Only used with --streaming." ,
710700 )
711- parser .add_argument (
712- "--rocm-packed-matvec" ,
713- action = "store_true" ,
714- help = "Use the experimental fixed-shape packed INT4 decoder matvec on ROCm." ,
715- )
716701 parser .add_argument (
717702 "--dtype" ,
718703 default = "fp32" ,
@@ -771,7 +756,6 @@ def main():
771756 "qembedding" : args .qembedding ,
772757 "qembedding_group_size" : args .qembedding_group_size ,
773758 "use_aoti_packed_int4" : args .backend == "rocm" ,
774- "use_aoti_matvec" : args .rocm_packed_matvec ,
775759 "backend" : backend_for_export ,
776760 }
777761 if args .streaming :
0 commit comments