Commit 7d6258fa authored by Po Yen Chen's avatar Po Yen Chen
Browse files

Generate prefill combine instances only for group mode

parent 94802528
...@@ -728,10 +728,11 @@ def get_fwd_splitkv_combine_blobs(kernel_filter : Optional[str], receipt) -> Lis ...@@ -728,10 +728,11 @@ def get_fwd_splitkv_combine_blobs(kernel_filter : Optional[str], receipt) -> Lis
continue continue
#for hdim_str, mode, mask, bias, lse in itertools.product(d.keys(), MODE_MAP.keys(), MASK_MAP.keys(), ["t", "f"], ["t", "f"]): #for hdim_str, mode, mask, bias, lse in itertools.product(d.keys(), MODE_MAP.keys(), MASK_MAP.keys(), ["t", "f"], ["t", "f"]):
for hdim_str, mode in itertools.product(d.keys(), MODE_MAP.keys()): for hdim_str, mode in itertools.product(d.keys(), MODE_MAP.keys()):
tiles = d[hdim_str] # include prefill tile size if in group mode
tiles = d[hdim_str][0 : 2 if mode == 'group' else 1]
hdim = int(hdim_str) hdim = int(hdim_str)
for tile, pipeline in itertools.product(tiles, get_pipelines(dtype, hdim)): for tile, pipeline in itertools.product(tiles, get_pipelines(dtype, hdim)):
if mode == "group": if mode == 'group':
if pipeline.F_spad != 't': if pipeline.F_spad != 't':
# in group mode, spad/skpad must be true, since we can't predict if seqlen of current batch need pad or not # in group mode, spad/skpad must be true, since we can't predict if seqlen of current batch need pad or not
continue continue
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment