diff --git a/testing/taskcluster/mach_commands.py b/testing/taskcluster/mach_commands.py
index 877ad2f083305..6db5cba809e32 100644
--- a/testing/taskcluster/mach_commands.py
+++ b/testing/taskcluster/mach_commands.py
@@ -336,15 +336,20 @@ def create_graph(self, **params):
                 test_parameters = copy.copy(build_parameters)
                 test_parameters['build_url'] = build_url
                 test_parameters['tests_url'] = tests_url
-                test_parameters['total_chunks'] = 1
 
+                test_definition = templates.load(test['task'], {})['task']
+                chunk_config = test_definition['extra']['chunks']
+
+                # Allow branch configs to override task level chunking...
                 if 'chunks' in test:
-                    test_parameters['total_chunks'] = test['chunks']
+                    chunk_config['total'] = test['chunks']
+
+                test_parameters['total_chunks'] = chunk_config['total']
 
-                for chunk in range(1, test_parameters['total_chunks'] + 1):
+                for chunk in range(1, chunk_config['total'] + 1):
                     if 'only_chunks' in test and \
                         chunk not in test['only_chunks']:
-                        continue;
+                        continue
 
                     test_parameters['chunk'] = chunk
                     test_task = templates.load(test['task'], test_parameters)